312 lines
11 KiB
Python
312 lines
11 KiB
Python
"""Conservative prompt budget for personal journal generation.
|
||
|
||
Token counts are estimates, not tokenizer-accurate. Kanshō overestimates input
|
||
so a request is refused before it can overflow a context window. There is no
|
||
silent middle-out deletion of the day.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import math
|
||
import os
|
||
from dataclasses import dataclass, field
|
||
from typing import Any
|
||
|
||
# Journal generation requires a wide window so a full day can be reconstructed
|
||
# without dropping middle events. 32_768 is the product default, not a claim
|
||
# that every model tokenizer uses this unit identically.
|
||
JOURNAL_MIN_CONTEXT_TOKENS = 32_768
|
||
JOURNAL_TARGET_COMPLETION_TOKENS = 4_096
|
||
# Overestimate: German mixed with prompt English is typically closer to 3–4
|
||
# characters per token. Two characters per token is deliberately conservative.
|
||
DEFAULT_CHARS_PER_TOKEN = 2.0
|
||
DEFAULT_SAFETY_MARGIN = 0.15
|
||
CHAT_FORMAT_OVERHEAD_TOKENS = 64
|
||
JOURNAL_PURPOSES = {"journal_generate", "journal_reconstruct"}
|
||
|
||
ERROR_MODEL_CONTEXT_TOO_SMALL = "model_context_too_small"
|
||
ERROR_PROMPT_BUDGET_EXCEEDED = "prompt_budget_exceeded"
|
||
ERROR_MODEL_METADATA_UNKNOWN = "model_metadata_unknown"
|
||
ERROR_RECONSTRUCTION_INVALID = "reconstruction_invalid"
|
||
ERROR_NO_USER_SOURCES = "no_user_sources"
|
||
ERROR_OUTPUT_LIMIT_UNSUPPORTED = "output_limit_unsupported"
|
||
ERROR_PROVIDER_CONTEXT_LENGTH = "provider_context_length_rejected"
|
||
|
||
USER_MESSAGES = {
|
||
ERROR_MODEL_CONTEXT_TOO_SMALL: (
|
||
"Das gewählte Modell hat ein zu kleines Kontextfenster für die "
|
||
"Journalgenerierung. Es sind mindestens 32.000 Tokens erforderlich. "
|
||
"Es wurde kein anderes Modell gewählt."
|
||
),
|
||
ERROR_PROMPT_BUDGET_EXCEEDED: (
|
||
"Der Tagesdialog ist für eine sichere Verarbeitung zu umfangreich. "
|
||
"Kanshō hat nichts stillschweigend aus der Mitte entfernt. "
|
||
"Bitte wähle weniger Gespräche oder ein Modell mit größerem Kontext."
|
||
),
|
||
ERROR_MODEL_METADATA_UNKNOWN: (
|
||
"Die Kontextgrenzen des Modells konnten nicht sicher bestimmt werden. "
|
||
"Der Aufruf wurde nicht gesendet."
|
||
),
|
||
ERROR_RECONSTRUCTION_INVALID: (
|
||
"Die inhaltliche Rekonstruktion war unvollständig oder ungültig. "
|
||
"Es wurde kein Journalentwurf erzeugt."
|
||
),
|
||
ERROR_NO_USER_SOURCES: (
|
||
"In den ausgewählten Gesprächen gibt es noch keinen Nutzertext. "
|
||
"Es wurde kein Journalentwurf erzeugt."
|
||
),
|
||
ERROR_OUTPUT_LIMIT_UNSUPPORTED: (
|
||
"Das Modell unterstützt die für die Journalgenerierung reservierte "
|
||
"Ausgabelänge nicht. Der Aufruf wurde nicht gesendet."
|
||
),
|
||
ERROR_PROVIDER_CONTEXT_LENGTH: (
|
||
"Der Anbieter hat die Anfrage wegen der Kontextlänge abgelehnt. "
|
||
"Kanshō hat den Tagesdialog nicht automatisch gekürzt."
|
||
),
|
||
}
|
||
|
||
|
||
class JournalBudgetError(Exception):
|
||
def __init__(
|
||
self,
|
||
code: str,
|
||
message: str | None = None,
|
||
status_code: int = 422,
|
||
diagnostics: dict[str, Any] | None = None,
|
||
):
|
||
super().__init__(message or USER_MESSAGES.get(code) or code)
|
||
self.code = code
|
||
self.message = message or USER_MESSAGES.get(code) or code
|
||
self.status_code = status_code
|
||
self.diagnostics = diagnostics or {}
|
||
|
||
|
||
@dataclass
|
||
class ModelWindow:
|
||
model: str
|
||
context_length: int
|
||
max_completion_tokens: int | None
|
||
source: str
|
||
provider: str | None = None
|
||
cached: bool = False
|
||
|
||
|
||
@dataclass
|
||
class JournalBudget:
|
||
model: str
|
||
purpose: str
|
||
effective_context_window: int
|
||
reserved_output_tokens: int
|
||
safety_margin: float
|
||
estimated_input_tokens: int = 0
|
||
available_input_tokens: int = 0
|
||
chars_per_token: float = DEFAULT_CHARS_PER_TOKEN
|
||
budget_ok: bool = True
|
||
abort_reason: str | None = None
|
||
context_compression: str = "disabled"
|
||
extras: dict[str, Any] = field(default_factory=dict)
|
||
|
||
def as_diagnostics(self) -> dict[str, Any]:
|
||
payload = {
|
||
"model": self.model,
|
||
"purpose": self.purpose,
|
||
"estimated_input_tokens": self.estimated_input_tokens,
|
||
"effective_context_window": self.effective_context_window,
|
||
"reserved_output_tokens": self.reserved_output_tokens,
|
||
"safety_margin": self.safety_margin,
|
||
"available_input_tokens": self.available_input_tokens,
|
||
"chars_per_token": self.chars_per_token,
|
||
"budget_ok": self.budget_ok,
|
||
"context_compression": self.context_compression,
|
||
"abort_reason": self.abort_reason,
|
||
"estimation": "conservative_char_ratio",
|
||
}
|
||
payload.update(self.extras)
|
||
return payload
|
||
|
||
|
||
def _float_env(name: str, default: float) -> float:
|
||
raw = (os.environ.get(name) or "").strip()
|
||
if not raw:
|
||
return default
|
||
try:
|
||
value = float(raw)
|
||
except ValueError:
|
||
return default
|
||
return value if value > 0 else default
|
||
|
||
|
||
def _int_env(name: str, default: int) -> int:
|
||
raw = (os.environ.get(name) or "").strip()
|
||
if not raw:
|
||
return default
|
||
try:
|
||
value = int(raw)
|
||
except ValueError:
|
||
return default
|
||
return value if value > 0 else default
|
||
|
||
|
||
def chars_per_token() -> float:
|
||
return _float_env("KANSHO_TOKEN_CHARS_PER_TOKEN", DEFAULT_CHARS_PER_TOKEN)
|
||
|
||
|
||
def safety_margin() -> float:
|
||
return min(0.4, max(0.05, _float_env("KANSHO_JOURNAL_SAFETY_MARGIN", DEFAULT_SAFETY_MARGIN)))
|
||
|
||
|
||
def min_context_tokens() -> int:
|
||
return _int_env("KANSHO_JOURNAL_MIN_CONTEXT_TOKENS", JOURNAL_MIN_CONTEXT_TOKENS)
|
||
|
||
|
||
def target_completion_tokens() -> int:
|
||
return _int_env("KANSHO_JOURNAL_MAX_COMPLETION_TOKENS", JOURNAL_TARGET_COMPLETION_TOKENS)
|
||
|
||
|
||
def day_message_safety_cap() -> int:
|
||
return _int_env("KANSHO_JOURNAL_DAY_MAX_MESSAGES", 500)
|
||
|
||
|
||
def estimate_tokens(text: str, *, extra: int = 0) -> int:
|
||
"""Conservative character ratio. Not a model tokenizer."""
|
||
n = len(text or "")
|
||
if n <= 0:
|
||
return extra
|
||
return max(1, math.ceil(n / chars_per_token())) + extra
|
||
|
||
|
||
def usable_context_tokens(window: ModelWindow) -> int:
|
||
margin = safety_margin()
|
||
usable = math.floor(window.context_length * (1.0 - margin))
|
||
return max(0, usable)
|
||
|
||
|
||
def reserved_output_tokens(window: ModelWindow) -> int:
|
||
wanted = target_completion_tokens()
|
||
supported = window.max_completion_tokens
|
||
if supported is None or supported <= 0:
|
||
raise JournalBudgetError(
|
||
ERROR_OUTPUT_LIMIT_UNSUPPORTED,
|
||
diagnostics={"model": window.model, "max_completion_tokens": supported},
|
||
)
|
||
reserved = min(wanted, int(supported))
|
||
if reserved < 256:
|
||
raise JournalBudgetError(
|
||
ERROR_OUTPUT_LIMIT_UNSUPPORTED,
|
||
diagnostics={
|
||
"model": window.model,
|
||
"max_completion_tokens": supported,
|
||
"reserved_output_tokens": reserved,
|
||
},
|
||
)
|
||
return reserved
|
||
|
||
|
||
def plan_journal_budget(window: ModelWindow, *, purpose: str) -> JournalBudget:
|
||
required = min_context_tokens()
|
||
if window.context_length < required:
|
||
raise JournalBudgetError(
|
||
ERROR_MODEL_CONTEXT_TOO_SMALL,
|
||
diagnostics={
|
||
"model": window.model,
|
||
"effective_context_window": window.context_length,
|
||
"required_context_window": required,
|
||
"source": window.source,
|
||
},
|
||
)
|
||
reserved = reserved_output_tokens(window)
|
||
usable = usable_context_tokens(window)
|
||
available = usable - reserved - CHAT_FORMAT_OVERHEAD_TOKENS
|
||
if available < 512:
|
||
raise JournalBudgetError(
|
||
ERROR_PROMPT_BUDGET_EXCEEDED,
|
||
diagnostics={
|
||
"model": window.model,
|
||
"effective_context_window": window.context_length,
|
||
"reserved_output_tokens": reserved,
|
||
"available_input_tokens": available,
|
||
},
|
||
)
|
||
compression = "disabled" if purpose in JOURNAL_PURPOSES else "provider_default"
|
||
return JournalBudget(
|
||
model=window.model,
|
||
purpose=purpose,
|
||
effective_context_window=window.context_length,
|
||
reserved_output_tokens=reserved,
|
||
safety_margin=safety_margin(),
|
||
available_input_tokens=available,
|
||
chars_per_token=chars_per_token(),
|
||
context_compression=compression,
|
||
extras={"metadata_source": window.source, "cached_metadata": window.cached},
|
||
)
|
||
|
||
|
||
def assert_input_fits(budget: JournalBudget, text: str) -> JournalBudget:
|
||
estimated = estimate_tokens(text, extra=CHAT_FORMAT_OVERHEAD_TOKENS)
|
||
budget.estimated_input_tokens = estimated
|
||
total = estimated + budget.reserved_output_tokens
|
||
ceiling = math.floor(budget.effective_context_window * (1.0 - budget.safety_margin))
|
||
if estimated > budget.available_input_tokens or total > ceiling:
|
||
budget.budget_ok = False
|
||
budget.abort_reason = ERROR_PROMPT_BUDGET_EXCEEDED
|
||
raise JournalBudgetError(
|
||
ERROR_PROMPT_BUDGET_EXCEEDED,
|
||
diagnostics={
|
||
**budget.as_diagnostics(),
|
||
"estimated_input_tokens": estimated,
|
||
"required_tokens": total,
|
||
"usable_tokens": ceiling,
|
||
},
|
||
)
|
||
budget.budget_ok = True
|
||
budget.abort_reason = None
|
||
return budget
|
||
|
||
|
||
def merge_usage(diagnostics: dict[str, Any], usage: dict[str, Any] | None, model: str | None = None) -> dict[str, Any]:
|
||
payload = dict(diagnostics)
|
||
data = usage or {}
|
||
if model:
|
||
payload["actual_model"] = model
|
||
payload["prompt_tokens"] = data.get("prompt_tokens")
|
||
payload["completion_tokens"] = data.get("completion_tokens")
|
||
payload["total_tokens"] = data.get("total_tokens")
|
||
cost = data.get("cost")
|
||
if cost is None:
|
||
cost = data.get("total_cost")
|
||
payload["cost"] = cost
|
||
return payload
|
||
|
||
|
||
def sum_usages(parts: list[dict[str, Any] | None]) -> dict[str, Any]:
|
||
prompt = 0
|
||
completion = 0
|
||
total = 0
|
||
cost = 0.0
|
||
saw_tokens = False
|
||
saw_cost = False
|
||
for item in parts:
|
||
data = item or {}
|
||
if data.get("prompt_tokens") is not None:
|
||
prompt += int(data.get("prompt_tokens") or 0)
|
||
saw_tokens = True
|
||
if data.get("completion_tokens") is not None:
|
||
completion += int(data.get("completion_tokens") or 0)
|
||
saw_tokens = True
|
||
if data.get("total_tokens") is not None:
|
||
total += int(data.get("total_tokens") or 0)
|
||
saw_tokens = True
|
||
extra = data.get("cost")
|
||
if extra is None:
|
||
extra = data.get("total_cost")
|
||
if extra is not None:
|
||
cost += float(extra)
|
||
saw_cost = True
|
||
payload: dict[str, Any] = {}
|
||
if saw_tokens:
|
||
payload["prompt_tokens"] = prompt
|
||
payload["completion_tokens"] = completion
|
||
payload["total_tokens"] = total or (prompt + completion)
|
||
if saw_cost:
|
||
payload["cost"] = cost
|
||
return payload
|