OSVauco/ml/token_budget.py

"""
ml/token_budget.py — Token-budsjett og context-trimming for OSVauco OPAX.

CG4-budget:
  - MAX_INPUT_TOKENS:  hardt tak på input per kall
  - trim_context():    kutter eldre turns hvis kontekst overstiger grensen
  - route_model():     velger Flash vs Pro basert på kompleksitet
  - TokenBudgetExceeded: kastes hvis ett enkelt kall overstiger tak

Maks-verdier (konservative, justerbare via env):
  OPAX_MAX_INPUT_TOKENS   default 32_000   (Gemini 2.5 Flash: 1M, Pro: 1M)
  OPAX_MAX_OUTPUT_TOKENS  default 4_096
  OPAX_CONTEXT_TURNS      default 10       (maks antall turns beholdt i historikk)
  OPAX_FLASH_THRESHOLD    default 8_000    (over dette: bruk Pro i heavy-mode)
"""

import os
import logging
from typing import List

from google.genai import types

logger = logging.getLogger(__name__)

MAX_INPUT_TOKENS  = int(os.environ.get("OPAX_MAX_INPUT_TOKENS",  "32000"))
MAX_OUTPUT_TOKENS = int(os.environ.get("OPAX_MAX_OUTPUT_TOKENS", "4096"))
MAX_CONTEXT_TURNS = int(os.environ.get("OPAX_CONTEXT_TURNS",     "10"))
FLASH_THRESHOLD   = int(os.environ.get("OPAX_FLASH_THRESHOLD",   "8000"))

# Grov estimering: 1 token ≈ 4 tegn for norsk/engelsk tekst
_CHARS_PER_TOKEN = 4


class TokenBudgetExceeded(Exception):
    """Kastes når et enkelt kall overstiger MAX_INPUT_TOKENS."""
    pass


def _estimate_tokens(text: str) -> int:
    """Grov token-estimering uten API-kall."""
    return max(1, len(text) // _CHARS_PER_TOKEN)


def trim_context(history: List[types.Content], system_prompt: str = "", max_tokens: int = MAX_INPUT_TOKENS) -> List[types.Content]:
    """
    Trimmer samtalehistorikk slik at total estimert token-mengde
    holder seg under max_tokens.

    Strategi:
      1. Behold alltid siste MAX_CONTEXT_TURNS turns
      2. Fjern eldre turns inntil vi er under grensen
      3. Kast TokenBudgetExceeded hvis SISTE melding alene er over grensen
    """
    if not history:
        return history

    # Behold maks MAX_CONTEXT_TURNS turns
    trimmed = history[-MAX_CONTEXT_TURNS:] if len(history) > MAX_CONTEXT_TURNS else list(history)

    # Beregn budsjett: trekk fra system prompt
    system_tokens = _estimate_tokens(system_prompt)
    available = max_tokens - system_tokens - MAX_OUTPUT_TOKENS

    # Sjekk om siste melding alene er for stor
    last_text = _get_text(trimmed[-1]) if trimmed else ""
    last_tokens = _estimate_tokens(last_text)
    if last_tokens > available:
        raise TokenBudgetExceeded(
            f"Siste melding er ~{last_tokens} tokens (maks {available}). "
            f"Del opp forespørselen i mindre deler."
        )

    # Fjern eldste turns til vi er under budsjett
    while trimmed:
        total = sum(_estimate_tokens(_get_text(c)) for c in trimmed)
        if total <= available:
            break
        trimmed.pop(0)
        logger.info(f"[token_budget] Trimmet 1 turn — gjenstående: {len(trimmed)}")

    if len(trimmed) < len(history):
        logger.info(
            f"[token_budget] Kontekst trimmet: {len(history)} → {len(trimmed)} turns "
            f"(~{sum(_estimate_tokens(_get_text(c)) for c in trimmed)} tokens)"
        )

    return trimmed


def route_model(message: str, mode: str, flash_model: str, pro_model: str) -> str:
    """
    Flash-first routing:
      - light-mode:  alltid Flash
      - heavy-mode:  Flash hvis under FLASH_THRESHOLD tokens, ellers Pro

    Dette reduserer Pro-forbruk dramatisk for enkle heavy-kall.
    """
    if mode != "heavy":
        return flash_model

    estimated = _estimate_tokens(message)
    if estimated <= FLASH_THRESHOLD:
        logger.info(f"[token_budget] heavy-kall rutes til Flash (~{estimated} tokens < {FLASH_THRESHOLD})")
        return flash_model

    logger.info(f"[token_budget] heavy-kall rutes til Pro (~{estimated} tokens >= {FLASH_THRESHOLD})")
    return pro_model


def enforce_output_limit(generate_config: dict) -> dict:
    """Legger til max_output_tokens i generate_config hvis ikke satt."""
    if "max_output_tokens" not in generate_config:
        generate_config["max_output_tokens"] = MAX_OUTPUT_TOKENS
    return generate_config


def _get_text(content: types.Content) -> str:
    """Henter ren tekst fra et Content-objekt."""
    if not content or not content.parts:
        return ""
    return " ".join(p.text or "" for p in content.parts if hasattr(p, "text"))


def budget_summary() -> dict:
    """Returnerer gjeldende budsjett-konfigurasjon (for /opax/build-status eller debug)."""
    return {
        "max_input_tokens":  MAX_INPUT_TOKENS,
        "max_output_tokens": MAX_OUTPUT_TOKENS,
        "max_context_turns": MAX_CONTEXT_TURNS,
        "flash_threshold":   FLASH_THRESHOLD,
        "chars_per_token":   _CHARS_PER_TOKEN,
    }