Kansho/backend/prompt_budget.py

"""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