context.py

"""
Token accounting a inteligentní správa kontextu.

Dvě odpovědnosti:

1. TokenStats — thread-safe počítadlo tokenů. Přesná čísla bere z `usage`
   pole odpovědí (vLLM i HF router je vracejí), takže globální statistiky
   (/health, UI) i per-turn statistiky (patička v chatu) jsou skutečné,
   ne odhady. Cache hity se nepočítají (nic nestály).

2. compact_messages — deterministická kompakce kontextu, aby se dlouhé
   agentní smyčky vešly do context window a šetřily KV cache:
   - Fáze 1 (stárnutí): staré výsledky nástrojů (mimo posledních
     `keep_last_steps` bloků) se zkrátí na `aged_chars` znaků — přesné
     výpisy souborů/logů ztrácejí hodnotu, jakmile na ně agent zareagoval.
   - Fáze 2 (vypouštění): pokud se stále nevejde, vypouští NEJSTARŠÍ celé
     bloky (assistant + jeho tool odpovědi) hned za kotvou a nahradí je
     jednou souhrnnou poznámkou, aby model věděl, že historie byla zkrácena.
   - Vždy zůstává: systémové zprávy, první user zpráva (zadání úlohy)
     a posledních `keep_last_steps` bloků v plném znění.

   Kompakce nikdy nerozbije strukturu tool-callingu (tool zpráva bez svého
   assistant/tool_calls předchůdce) — vypouští se výhradně celé bloky.

Odhad tokenů je heuristika ~4 znaky/token (±15 %) — pro rozhodnutí „kdy
kompaktovat" bohatě stačí; přesná čísla dodává `usage` po každém volání.
"""

from __future__ import annotations

import json
import time
from threading import Lock

CHARS_PER_TOKEN = 4
MESSAGE_OVERHEAD_TOKENS = 4  # role, oddělovače, formát šablony


# ---------------------------------------------------------------- odhad tokenů

def _message_text(msg: dict) -> str:
    content = msg.get("content") or ""
    if isinstance(content, list):
        content = json.dumps(content, ensure_ascii=False)
    text = str(content)
    if msg.get("tool_calls"):
        text += json.dumps(msg["tool_calls"], ensure_ascii=False)
    return text


def estimate_tokens(messages: list[dict]) -> int:
    """Hrubý odhad počtu tokenů zpráv (~4 znaky/token + režie zprávy)."""
    total = 0
    for msg in messages:
        total += MESSAGE_OVERHEAD_TOKENS + len(_message_text(msg)) // CHARS_PER_TOKEN
    return total


# ---------------------------------------------------------------- statistiky

class TokenStats:
    """Thread-safe počítadlo skutečné spotřeby tokenů (z `usage` odpovědí).

    Používá se dvojmo: globální instance (celkové statistiky od startu app)
    a per-turn instance (patička jednoho chatu vč. sub-agentů).
    """

    def __init__(self):
        self._lock = Lock()
        self.calls = 0
        self.prompt_tokens = 0
        self.completion_tokens = 0
        self.by_source: dict[str, dict] = {}
        self.last_context_tokens = 0     # prompt_tokens posledního volání
        self.saved_tokens = 0            # ušetřeno kompakcí kontextu
        self.started_at = time.time()

    def add(self, source: str, usage) -> None:
        """Zaznamená jedno LLM volání. `usage` je objekt/dict s
        prompt_tokens/completion_tokens, nebo None (pak se nepočítá nic)."""
        if usage is None:
            return
        prompt = getattr(usage, "prompt_tokens", None)
        completion = getattr(usage, "completion_tokens", None)
        if prompt is None and isinstance(usage, dict):
            prompt = usage.get("prompt_tokens")
            completion = usage.get("completion_tokens")
        prompt = int(prompt or 0)
        completion = int(completion or 0)
        with self._lock:
            self.calls += 1
            self.prompt_tokens += prompt
            self.completion_tokens += completion
            self.last_context_tokens = prompt
            src = self.by_source.setdefault(
                source, {"calls": 0, "prompt_tokens": 0, "completion_tokens": 0})
            src["calls"] += 1
            src["prompt_tokens"] += prompt
            src["completion_tokens"] += completion

    def add_saved(self, tokens: int) -> None:
        with self._lock:
            self.saved_tokens += max(0, tokens)

    def snapshot(self) -> dict:
        with self._lock:
            return {
                "calls": self.calls,
                "prompt_tokens": self.prompt_tokens,
                "completion_tokens": self.completion_tokens,
                "total_tokens": self.prompt_tokens + self.completion_tokens,
                "last_context_tokens": self.last_context_tokens,
                "saved_by_compaction_tokens": self.saved_tokens,
                "by_source": {k: dict(v) for k, v in self.by_source.items()},
                "since": int(self.started_at),
            }


# ---------------------------------------------------------------- kompakce

TRUNCATION_MARK = "\n…[zkráceno kompakcí kontextu]"


def _split_blocks(messages: list[dict]) -> tuple[list[dict], list[list[dict]]]:
    """Rozdělí zprávy na (kotva, bloky).

    Kotva = úvodní system zprávy + první user zpráva. Blok = souvislá skupina,
    kterou lze bezpečně vypustit celou: assistant s tool_calls + jeho tool
    odpovědi, nebo samostatná user/assistant zpráva.
    """
    anchor: list[dict] = []
    i = 0
    while i < len(messages) and messages[i].get("role") == "system":
        anchor.append(messages[i])
        i += 1
    if i < len(messages) and messages[i].get("role") == "user":
        anchor.append(messages[i])
        i += 1

    blocks: list[list[dict]] = []
    current: list[dict] = []
    for msg in messages[i:]:
        if msg.get("role") == "tool":
            # tool patří k předchozímu assistant tool_calls bloku
            current.append(msg)
            continue
        if current:
            blocks.append(current)
        current = [msg]
    if current:
        blocks.append(current)
    return anchor, blocks


def _age_tool_results(blocks: list[list[dict]], keep_last: int,
                      aged_chars: int) -> int:
    """Zkrátí tool výsledky ve starých blocích. Vrací počet ušetřených znaků."""
    saved_chars = 0
    for block in blocks[:-keep_last] if keep_last > 0 else blocks:
        for msg in block:
            if msg.get("role") != "tool":
                continue
            content = msg.get("content") or ""
            if isinstance(content, str) and len(content) > aged_chars:
                saved_chars += len(content) - aged_chars
                msg["content"] = content[:aged_chars] + TRUNCATION_MARK
    return saved_chars


def _drop_summary(dropped: list[list[dict]]) -> dict:
    tools = []
    for block in dropped:
        for msg in block:
            for tc in msg.get("tool_calls") or []:
                name = (tc.get("function") or {}).get("name")
                if name:
                    tools.append(name)
    tool_note = f" (nástroje: {', '.join(tools[:15])})" if tools else ""
    return {"role": "system",
            "content": f"🧹 [Kompakce kontextu: vypuštěno {len(dropped)} "
                       f"starších kroků konverzace{tool_note}. Shrnutí stavu "
                       f"si drž ve vlastních odpovědích.]"}


def compact_messages(messages: list[dict], budget_tokens: int,
                     keep_last_steps: int = 6,
                     aged_chars: int = 2000) -> tuple[list[dict], int]:
    """Zhutní kontext pod budget_tokens. Vrací (nové zprávy, ušetřené tokeny).

    Deterministická, bez LLM volání. Vstup nemodifikuje (pracuje s kopií).
    Pokud se ani po vypuštění všeho kromě kotvy a posledních bloků nevejde,
    vrací nejlepší možný výsledek (best effort).
    """
    before = estimate_tokens(messages)
    if before <= budget_tokens:
        return messages, 0

    work = [dict(m) for m in messages]
    anchor, blocks = _split_blocks(work)

    # Fáze 1: stárnutí starých tool výsledků
    _age_tool_results(blocks, keep_last_steps, aged_chars)
    flat = anchor + [m for b in blocks for m in b]
    if estimate_tokens(flat) <= budget_tokens:
        return flat, before - estimate_tokens(flat)

    # Fáze 2: vypouštění nejstarších bloků (poslední keep_last_steps nikdy)
    dropped: list[list[dict]] = []
    keep = max(1, keep_last_steps)
    while len(blocks) > keep:
        dropped.append(blocks.pop(0))
        flat = anchor + [_drop_summary(dropped)] + [m for b in blocks for m in b]
        if estimate_tokens(flat) <= budget_tokens:
            break
    if dropped:
        flat = anchor + [_drop_summary(dropped)] + [m for b in blocks for m in b]
    return flat, max(0, before - estimate_tokens(flat))