nisse/nisse-dev/src/nisse_dev/context.py

"""The context budget (architecture L6: token budgeting is a render-time concern).

Every model call re-sends a transcript that only grows, and a model server
does not always refuse one that no longer fits: Ollama silently drops the
start of the prompt, which is where the system prompt and the plan are. So a
request is made to fit the model's declared window *before* it is sent, and
one that cannot be made to fit is never sent.

The log itself is never changed. Fitting works on the events about to be
rendered, in three steps, stopping at the first that fits:

1. Tool results older than the most recent turns are cut to their first and
   last lines, with a note saying how much was left out.
2. The oldest events are dropped, keeping tool calls with their results, and
   a note at the start says how many.
3. If even the recent turns and the system prompt do not fit, the call
   fails with ContextOverflow. The caller treats it like any failed call:
   the executor stops and asks the operator.

Token counts are estimated from characters, conservatively, and the estimate
is corrected from the input-token counts the server reports back. The
correction only ever makes the estimate larger: underestimating is the
dangerous direction.
"""

import json
from collections.abc import Callable, Sequence
from dataclasses import dataclass, replace
from typing import Any

from nisse_dev.canonical import (
    AssistantText,
    CanonicalEvent,
    DelegationStarted,
    PlanSubmitted,
    StoredEvent,
    ToolCall,
    ToolResult,
    UserMessage,
)
from nisse_dev.llm import DEFAULT_MAX_OUTPUT_TOKENS, LLMError

# English and code run at about four characters a token. Three is deliberately
# pessimistic, before any correction from what the server reports.
CHARS_PER_TOKEN = 3.0
# Role markers and formatting each rendered message costs beyond its text.
MESSAGE_OVERHEAD_TOKENS = 8
# How much of an older tool result survives step 1.
ELIDED_RESULT_CHARS = 600

# Events the model produced. A turn starts where these start.
_MODEL_SIDE = (AssistantText, ToolCall, PlanSubmitted, DelegationStarted)


class ContextOverflow(LLMError):
    """The request cannot be made to fit the model's context window."""


def head_and_tail(text: str, limit: int) -> str:
    """Keep the start and, mostly, the end of a long text.

    The end matters more: test runners and compilers put their summary of
    what failed last. The start keeps the exit code and the first lines.
    """
    if len(text) <= limit:
        return text
    head = limit // 4
    tail = limit - head
    dropped = len(text) - limit
    return f"{text[:head]}\n[... {dropped} characters omitted ...]\n{text[-tail:]}"


def _event_chars(event: CanonicalEvent) -> int:
    match event:
        case ToolCall(tool=tool, arguments=arguments):
            return len(tool) + len(json.dumps(arguments))
        case ToolResult(result=result):
            return len(result)
        case _:
            text = getattr(event, "text", None) or getattr(event, "plan", None)
            if text is None:
                text = json.dumps(getattr(event, "__dict__", {}), default=str)
            return len(str(text))


def estimate_tokens(
    events: Sequence[StoredEvent],
    system: str | None,
    tools: list[dict[str, Any]] | None,
    factor: float = 1.0,
) -> int:
    chars = sum(_event_chars(e.event) for e in events)
    chars += len(system or "") + len(json.dumps(tools or []))
    raw = chars / CHARS_PER_TOKEN + MESSAGE_OVERHEAD_TOKENS * (len(events) + 1)
    return int(raw * factor) + 1


def _turn_starts(events: Sequence[StoredEvent]) -> list[int]:
    starts = []
    previous_model_side = False
    for index, stored in enumerate(events):
        model_side = isinstance(stored.event, _MODEL_SIDE)
        if model_side and not previous_model_side:
            starts.append(index)
        previous_model_side = model_side
    return starts


def _consistent(events: list[StoredEvent]) -> list[StoredEvent]:
    """Drop results whose call is gone: a provider rejects a tool result
    with nothing to answer."""
    calls = {e.event.call_id for e in events if isinstance(e.event, ToolCall)}
    return [
        e
        for e in events
        if not isinstance(e.event, ToolResult) or e.event.call_id in calls
    ]


@dataclass(frozen=True)
class Budget:
    window: int  # prompt plus response, in tokens
    max_output: int  # the ceiling sent with the request
    keep_recent_turns: int
    safety_margin: float  # fraction of the window held back

    @property
    def for_prompt(self) -> int:
        return int(self.window * (1 - self.safety_margin)) - self.max_output


def fit(
    events: Sequence[StoredEvent],
    system: str | None,
    tools: list[dict[str, Any]] | None,
    budget: Budget,
    factor: float = 1.0,
) -> list[StoredEvent]:
    """The events to render: all of them if they fit, otherwise the result of
    the three steps in the module docstring."""
    limit = budget.for_prompt

    def fits(candidate: Sequence[StoredEvent]) -> bool:
        return estimate_tokens(candidate, system, tools, factor) <= limit

    kept = list(events)
    if fits(kept):
        return kept

    starts = _turn_starts(kept)
    recent_from = (
        starts[-budget.keep_recent_turns]
        if len(starts) >= budget.keep_recent_turns
        else 0
    )

    # Step 1: cut older tool results down.
    kept = [
        replace(
            stored,
            event=replace(
                stored.event,
                result=head_and_tail(stored.event.result, ELIDED_RESULT_CHARS),
            ),
        )
        if index < recent_from and isinstance(stored.event, ToolResult)
        else stored
        for index, stored in enumerate(kept)
    ]
    if fits(kept):
        return kept

    # Step 2: drop the oldest events, one turn at a time, never the recent
    # ones. The note keeps the transcript starting with a user message, which
    # some chat templates require.
    cuts = sorted({s for s in starts if 0 < s < recent_from} | {recent_from} - {0})
    for cut in cuts:
        remaining = _consistent(kept[cut:])
        note = StoredEvent(
            id=0,
            ts="",
            event=UserMessage(
                text=(
                    f"[{cut} earlier events omitted to fit the context window; "
                    "the log still has them]"
                ),
                source="context-budget",
            ),
        )
        candidate = [note, *remaining]
        if fits(candidate):
            return candidate

    # Step 3: nothing more can go.
    needed = estimate_tokens(kept[recent_from:], system, tools, factor)
    raise ContextOverflow(
        f"the system prompt and the last {budget.keep_recent_turns} turns need "
        f"about {needed} tokens, but only {limit} fit a {budget.window}-token "
        "window after the output ceiling and safety margin"
    )


class BudgetedClient:
    """An LLM client that fits every request to the model's window.

    Wraps a real client (either route) and takes each model's limits from
    ``limits_for``, which the agent fills from LiteLLM's model catalog.
    """

    def __init__(
        self,
        inner: Any,
        limits_for: Callable[[str], tuple[int, int | None]],
        *,
        max_output_tokens: int = DEFAULT_MAX_OUTPUT_TOKENS,
        keep_recent_turns: int = 6,
        safety_margin: float = 0.1,
        enforces_tool_choice: Callable[[str], bool] = lambda model: False,
    ) -> None:
        self._inner = inner
        self._limits_for = limits_for
        # Whether a model's server holds it to the offered tools when asked
        # (LiteLLM's supports_tool_choice). Asking one that doesn't can be
        # worse than not asking: qwen3-14b can't stop, and rambles until the
        # output ceiling.
        self._enforces_tool_choice = enforces_tool_choice
        self._max_output_tokens = max_output_tokens
        self._keep_recent_turns = keep_recent_turns
        self._safety_margin = safety_margin
        # Per model: how much the character estimate undercounted, as learned
        # from reported usage. Starts at 1 and only grows.
        self._factors: dict[str, float] = {}
        self.trimmed_requests = 0

    @property
    def last_usage(self) -> dict[str, int]:
        return getattr(self._inner, "last_usage", {}) or {}

    @property
    def last_model(self) -> str | None:
        return getattr(self._inner, "last_model", None)

    @property
    def last_cost(self) -> float:
        return float(getattr(self._inner, "last_cost", 0.0) or 0.0)

    @property
    def last_rate_limit_wait_s(self) -> float:
        return float(getattr(self._inner, "last_rate_limit_wait_s", 0.0) or 0.0)

    def budget_for(self, model: str) -> Budget:
        window, model_max_output = self._limits_for(model)
        ceiling = self._max_output_tokens
        if model_max_output is not None:
            ceiling = min(ceiling, model_max_output)
        return Budget(
            window=window,
            max_output=ceiling,
            keep_recent_turns=self._keep_recent_turns,
            safety_margin=self._safety_margin,
        )

    def complete(
        self,
        events: list[StoredEvent],
        tools: list[dict[str, Any]] | None,
        model: str,
        system: str | None = None,
    ) -> list[CanonicalEvent]:
        budget = self.budget_for(model)
        factor = self._factors.get(model, 1.0)
        fitted = fit(events, system, tools, budget, factor)
        # Counted so the benchmark can report how often a model's window was
        # the constraint: any request that had to be cut down or shortened.
        if fitted != list(events):
            self.trimmed_requests += 1
        produced = self._inner.complete(
            fitted,
            tools,
            model,
            system=system,
            max_tokens=budget.max_output,
            # Every caller that offers tools needs one called: the executor
            # finishes with finish, the planner with submit_plan or
            # ask_human, a delegation with hand_back.
            require_tool_call=bool(tools) and self._enforces_tool_choice(model),
        )
        reported = int(self.last_usage.get("input", 0))
        if reported > 0:
            estimated = estimate_tokens(fitted, system, tools)
            observed = reported / max(estimated, 1)
            self._factors[model] = max(factor, observed)
        return produced