diff --git a/docs/Tuning-Guide.md b/docs/Tuning-Guide.md index 49fb34f..ab2ce87 100644 --- a/docs/Tuning-Guide.md +++ b/docs/Tuning-Guide.md @@ -114,8 +114,9 @@ to look at when deciding. ### `llm.max_output_tokens` — default 16384 -The ceiling on a single response. It is sent with every request, for every -role. +The ceiling on a single response, for every role. The ceiling actually sent +is the smaller of this and the model's own `max_output_tokens` in LiteLLM, so +a local model can be held to less without affecting the frontier roles. - **Why it exists:** a response that hits the ceiling is truncated mid-sentence, or worse, mid-tool-call, where it can still parse. The agent @@ -128,9 +129,8 @@ role. prompt plus ceiling exceeds the context. - **Too low:** truncated plans and tool calls. - **Deciding:** executor turns are short, so 2–4k is plenty for a local - executor. Frontier planning wants the large default. This one knob covers - every role for now; a per-model setting is coming with the context budget - (see *Coming next*). + executor. Frontier planning wants the large default. Leave this high, and + set a local model's smaller ceiling as its `max_output_tokens` in LiteLLM. ### `tuning.executor.max_turns` — default 24 @@ -224,19 +224,105 @@ How long one `run_command` may run before its whole process group is killed. ### `tuning.tools.max_output_chars` — default 4000 How much of a tool's output (command output, file contents, directory -listings) the model sees. Anything past this is cut, with a note saying how -much was dropped. +listings) the model sees. Longer output keeps its first quarter and last three +quarters, with a note in between saying how much was left out. The end gets +the larger share because test runners and compilers print their summary of +what failed last; the start keeps the exit code. - **This is a context knob.** Every tool result stays in the transcript and is re-sent every turn. 4,000 characters is about 1,000 tokens, so 20 turns of maximal output is about 20k tokens, which is already most of a local model's context. - **Too low:** the model doesn't see the part of the output that matters. -- **Too high:** the transcript outgrows the context window sooner. -- **A known weakness:** truncation keeps the *start* of the output, but test - runners put their failure summary at the *end*. A long, failing pytest run - can lose exactly the lines the model needed. Keeping the head and the tail - would fix this, and is a candidate for the context-budget work. +- **Too high:** the transcript outgrows the context window sooner. The + context budget copes (see below), but by cutting older output harder. + +### `tuning.context` — fitting each request into the window + +Before any request is sent, the agent checks it against the model's context +window and makes it fit. A server doesn't always refuse a prompt that's too +long. Ollama silently drops its start, which is where the system prompt and +the plan are, so an unfitted request can fail in a way nobody sees. + +The window is each model's `max_input_tokens` in LiteLLM (see *Declaring a +local model's limits* below). A model in use without one stops the agent at +startup. From the window come the output ceiling and a safety margin, and +the rest is for the prompt. When the prompt is too big, three steps apply in +order, stopping at the first that fits: + +1. **Older tool results are cut** to about 600 characters, head and tail. + The most recent turns keep their output whole. +2. **The oldest turns are dropped,** tool calls together with their results. + A note at the start tells the model how many events were left out. The + log keeps everything; only the request is shortened. +3. **If the system prompt and the recent turns still don't fit,** the + request is not sent. The milestone stops and asks you. + +Tokens are estimated from characters (pessimistically, at 3 per token) and +the estimate is corrected from the prompt sizes the server reports. The +correction only ever makes the estimate larger. + +The executor also sends less in the first place: the approved milestone plan +goes in its system prompt, which is never cut, and it sees only what happened +since that plan. The planning conversation and earlier milestones stay in the +log. + +| Knob | Default | | +|---|---|---| +| `keep_recent_turns` | 6 | model turns that are never cut or dropped | +| `safety_margin` | 0.1 | fraction of the window held back for estimation error | + +- **`keep_recent_turns` too low:** the model loses sight of what it just did + and repeats itself. **Too high:** less room to cut before step 3 gives up; + on a small window, overflows come sooner. +- **`safety_margin` too low:** an underestimate reaches the server, and a + silently cutting server drops the start of the prompt. **Too high:** wasted + context. Lower it only once the correction from reported usage has settled, + which takes a few calls per model. +- **Deciding:** if milestones stop with "cannot fit", look first at + `tools.max_output_chars` and the model's window, not at these. + +### Declaring a local model's limits + +LiteLLM knows commercial models' limits. For a local model, declare them in +LiteLLM's config as the model's `model_info`: + +```yaml +model_list: + - model_name: local-fast + litellm_params: {model: openai/qwen3-8b-awq, api_base: http://inference-host:8080/v1} + model_info: + max_input_tokens: 32768 # the server's per-request context + max_output_tokens: 4096 # the most one response may use +``` + +`max_input_tokens` is the *whole* window, prompt plus response. The agent +takes the output ceiling out of it. That's how the commercial numbers work +too: Claude's prompt and response together must fit in its 200k. + +The number to use is what the **server** allocates per request, not the +model's advertised maximum: + +- **The model's ceiling:** `max_position_embeddings` in `config.json`, or + `.context_length` in a GGUF file. Some models reach further with RoPE + scaling; the model card says whether that needs switching on. +- **llama.cpp / llama-swap:** the `-c` / `--ctx-size` flag, shown as `n_ctx` + in the load log and in `GET /props`. With `--parallel N` each request gets + `n_ctx / N`, so run the agent's models with `--parallel 1`. +- **vLLM:** `--max-model-len`, reported by `GET /v1/models`. +- **Ollama:** 4096 unless `OLLAMA_CONTEXT_LENGTH` or a Modelfile's `num_ctx` + says otherwise. `GET /api/ps` shows the loaded model's `context_length`; + `ollama show` shows only the trained maximum. + +Then check the declaration against the server: + + uv run python scripts/probe_context.py local-fast + +It sends one prompt half again as long as the declared window and reports +whether the server refused it, cut it (and where), or took it all. A cut +below the declaration is the case to fix: prompts between the two would be +silently cut. Measured on the Jetson's `llama3.2`: cut at 4096, matching its +declaration. --- @@ -259,7 +345,9 @@ change one thing at a time. 2. **Context length and KV-cache precision** on the server. Take the largest context that fits beside the weights with 1–2 GB of headroom. Try an 8-bit KV cache before giving up context. -3. **`llm.max_output_tokens`**, so the prompt keeps most of that context. +3. **Declare the window** in LiteLLM and check it with the probe, then set + the model's `max_output_tokens` there so the prompt keeps most of the + window. 4. **Real milestones.** Run representative work and look at the evidence: turns used, escalations, delegation outcomes, and time taken. 5. **The loop knobs**: `failure_limit`, `max_turns`, `consult_budget`, the @@ -278,9 +366,9 @@ run enough times that the result would survive a coin flip. The next changes in the 4090-readiness work add knobs, and this guide gains a section for each: -- **A context budget.** Rendering a transcript that fits the model's context - window, with what to drop or summarize when it doesn't, and per-model output - ceilings. +- **Compaction.** A model-written summary replacing old turns, instead of + dropping them. The log already has an event for it; it needs a summarizing + model and a policy for what is safe to summarize. - **Enforced tool calls.** An explicit `finish` tool, so the server can be told a tool call is required and hold the model to the contract, where today it can fall back to prose. diff --git a/nisse-dev.example.yaml b/nisse-dev.example.yaml index 8f61011..1c0e278 100644 --- a/nisse-dev.example.yaml +++ b/nisse-dev.example.yaml @@ -15,8 +15,8 @@ project: demo llm: # A5: LiteLLM is the only LLM destination base_url: http://litellm.example:4000 token: ${NISSE_LLM_TOKEN} - # Ceiling on each response. Must fit inside a local model's context window - # alongside the prompt; see docs/Tuning-Guide.md. + # Ceiling on each response. The model's own max_output_tokens in LiteLLM + # caps it further, which is where to hold a local model to less. max_output_tokens: 16384 roles: # A11: role -> LiteLLM model alias @@ -54,4 +54,7 @@ tuning: wall_clock_s: 900 tools: command_timeout_s: 120 # one run_command before it is killed - max_output_chars: 4000 # tool output kept for the model; the rest is cut + max_output_chars: 4000 # tool output kept for the model, head and tail + context: # fitting each request into the model's window + keep_recent_turns: 6 # turns never cut or dropped + safety_margin: 0.1 # fraction of the window held back diff --git a/nisse-dev/src/nisse_dev/agent.py b/nisse-dev/src/nisse_dev/agent.py index 7aa0f6f..87cf354 100644 --- a/nisse-dev/src/nisse_dev/agent.py +++ b/nisse-dev/src/nisse_dev/agent.py @@ -15,6 +15,7 @@ from nisse_dev.approvals import PlanBoard from nisse_dev.cancel import Cancel from nisse_dev.catalog import ModelCatalog, select_client from nisse_dev.config import AgentConfig +from nisse_dev.context import BudgetedClient from nisse_dev.control import Orchestrator from nisse_dev.delegation import Caps, Delegation from nisse_dev.executor import LLMExecutor @@ -50,6 +51,10 @@ def discover_models(config: AgentConfig, oplog: OpLog) -> ModelCatalog: """ catalog = ModelCatalog.discover(config.llm_base_url, config.llm_token) catalog.validate(config.roles.pinned()) + # Every role's context window must be known before a request is built + # against it (A10): a guess that is too large reaches a server that may + # silently cut the prompt. + catalog.require_limits(config.roles.pinned()) for role in TOOL_DRIVING_ROLES: facts = catalog.facts(config.roles.resolve(role)) if facts.supports_tools is False: @@ -156,14 +161,33 @@ def build_agent( max_output_chars=tuning.tools.max_output_chars, ) - normalized = LiteLLMClient( - config.llm_base_url, config.llm_token, max_tokens=config.max_output_tokens + def limits_for(model: str) -> tuple[int, int | None]: + facts = catalog.facts(model) + assert facts.max_input_tokens is not None # require_limits checked + return facts.max_input_tokens, facts.max_output_tokens + + def budgeted(client: object) -> BudgetedClient: + # Every request, on either route, is fitted to its model's window. + return BudgetedClient( + client, + limits_for, + max_output_tokens=config.max_output_tokens, + keep_recent_turns=tuning.context.keep_recent_turns, + safety_margin=tuning.context.safety_margin, + ) + + normalized = budgeted( + LiteLLMClient( + config.llm_base_url, config.llm_token, max_tokens=config.max_output_tokens + ) ) # Delegations prefer the vendor-native route; fidelity is cheap # elsewhere (A5). Which route a given model gets is decided by # select_client, not by where the caller happens to sit. - passthrough = AnthropicPassthroughClient( - config.llm_base_url, config.llm_token, max_tokens=config.max_output_tokens + passthrough = budgeted( + AnthropicPassthroughClient( + config.llm_base_url, config.llm_token, max_tokens=config.max_output_tokens + ) ) transport = MatrixTransport(config.matrix) diff --git a/nisse-dev/src/nisse_dev/catalog.py b/nisse-dev/src/nisse_dev/catalog.py index 175684d..0bab212 100644 --- a/nisse-dev/src/nisse_dev/catalog.py +++ b/nisse-dev/src/nisse_dev/catalog.py @@ -35,12 +35,25 @@ class UnknownModel(LLMError): stale alias, or a key without access. Always a misconfiguration.""" +class UnknownLimits(LLMError): + """A model in use has no declared context window. LiteLLM knows the + limits of commercial models; for a local one they are whatever its server + allocates, and have to be declared in LiteLLM's ``model_info``. Guessing + is not safe: a server that silently truncates a prompt (Ollama does) + drops its start, which is where the system prompt and the plan are.""" + + @dataclass(frozen=True) class ModelFacts: alias: str # what we ask for, e.g. "remote-architecture" backing: str # what LiteLLM calls, e.g. "anthropic/claude-opus-5" provider: str # "anthropic", "openai", "ollama", ... supports_tools: bool | None # None when LiteLLM does not say + # The whole context window, prompt plus response, as LiteLLM declares it + # (model_info.max_input_tokens); None when it does not say. + max_input_tokens: int | None = None + # The most one response may use (model_info.max_output_tokens). + max_output_tokens: int | None = None @property def native_anthropic(self) -> bool: @@ -72,6 +85,8 @@ class ModelCatalog: backing=backing, provider=str(provider), supports_tools=info.get("supports_function_calling"), + max_input_tokens=_positive_int(info.get("max_input_tokens")), + max_output_tokens=_positive_int(info.get("max_output_tokens")), ) return cls(facts) @@ -136,6 +151,24 @@ class ModelCatalog: f"available: {', '.join(self.aliases()) or '(none)'}" ) + def require_limits(self, roles: Mapping[str, str]) -> None: + """Refuse a role mapping that uses a model with no declared context + window, naming every such model at once.""" + missing = { + role: alias + for role, alias in roles.items() + if self.facts(alias).max_input_tokens is None + } + if missing: + named = ", ".join( + f"{role}={alias!r}" for role, alias in sorted(missing.items()) + ) + raise UnknownLimits( + f"no context window is declared for: {named}. Set " + "model_info.max_input_tokens (and max_output_tokens) for these " + "in LiteLLM; docs/Tuning-Guide.md says how to find the values" + ) + def pinned(self, roles: Mapping[str, str]) -> dict[str, dict[str, Any]]: """The resolved mapping to record at task start (A11, A16). @@ -154,6 +187,12 @@ class ModelCatalog: } +def _positive_int(value: Any) -> int | None: + if isinstance(value, bool) or not isinstance(value, int | float): + return None + return int(value) if value > 0 else None + + def select_client( model: str, normalized: LLMClient, diff --git a/nisse-dev/src/nisse_dev/config.py b/nisse-dev/src/nisse_dev/config.py index 03829f0..80d9071 100644 --- a/nisse-dev/src/nisse_dev/config.py +++ b/nisse-dev/src/nisse_dev/config.py @@ -86,12 +86,19 @@ class ToolsTuning: max_output_chars: int = 4000 +@dataclass(frozen=True) +class ContextTuning: + keep_recent_turns: int = 6 + safety_margin: float = 0.1 + + @dataclass(frozen=True) class Tuning: executor: ExecutorTuning = field(default_factory=ExecutorTuning) expert: ExpertTuning = field(default_factory=ExpertTuning) delegation: DelegationTuning = field(default_factory=DelegationTuning) tools: ToolsTuning = field(default_factory=ToolsTuning) + context: ContextTuning = field(default_factory=ContextTuning) _TUNING_SECTIONS: dict[str, type] = { @@ -128,6 +135,9 @@ def _parse_tuning(raw: Any, errors: list[str]) -> Tuning: continue values[key] = number sections[name] = cls(**values) + margin = sections.get("context", ContextTuning()).safety_margin + if margin >= 1: + errors.append("tuning.context.safety_margin: must be less than 1") return Tuning(**sections) diff --git a/nisse-dev/src/nisse_dev/context.py b/nisse-dev/src/nisse_dev/context.py new file mode 100644 index 0000000..1151dd5 --- /dev/null +++ b/nisse-dev/src/nisse_dev/context.py @@ -0,0 +1,262 @@ +"""The context budget (spec A10: 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, + ) -> None: + self._inner = inner + self._limits_for = limits_for + 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] = {} + + @property + def last_usage(self) -> dict[str, int]: + return getattr(self._inner, "last_usage", {}) or {} + + 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) + produced = self._inner.complete( + fitted, tools, model, system=system, max_tokens=budget.max_output + ) + 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 diff --git a/nisse-dev/src/nisse_dev/executor.py b/nisse-dev/src/nisse_dev/executor.py index 958d5cf..6ec80fa 100644 --- a/nisse-dev/src/nisse_dev/executor.py +++ b/nisse-dev/src/nisse_dev/executor.py @@ -15,7 +15,7 @@ from dataclasses import dataclass from typing import Any, Literal from nisse_dev.cancel import Cancel, skip_calls -from nisse_dev.canonical import ToolCall, ToolResult +from nisse_dev.canonical import PlanSubmitted, StoredEvent, ToolCall, ToolResult from nisse_dev.delegation import Delegation from nisse_dev.expert import ASK_EXPERT_TOOL, BudgetExhausted, ExpertDesk from nisse_dev.llm import LLMClient, LLMError, RoleMap @@ -24,7 +24,7 @@ from nisse_dev.tools import EXECUTOR_TOOLS, MUTATING_TOOLS, ToolOutcome, Workspa EXECUTOR_SYSTEM = """You are nisse-dev's executor, doing mechanical work. -Follow the approved plan above exactly. It was written to be executed +Follow the approved plan below exactly. It was written to be executed verbatim; do not redesign it, expand its scope, or skip steps. Work by calling tools. Verify your work by running the project's tests. @@ -145,12 +145,34 @@ class LLMExecutor: except LLMError as exc: return ToolOutcome(f"the consultant did not answer: {exc}", success=False) + def _scope(self) -> tuple[str | None, int]: + """The approved plan, and where this milestone's events begin. + + The most recent plan in the log is the approved milestone plan: no + plan is submitted between approval and execution, or during review + responses. The executor sees that plan and what came after it, and + not the planning conversation or earlier milestones, which a small + context window cannot afford and the plan already distils. + """ + latest: StoredEvent | None = None + for stored in self._store.read_events(): + if isinstance(stored.event, PlanSubmitted): + latest = stored + if latest is None: + return None, 0 + assert isinstance(latest.event, PlanSubmitted) + return latest.event.plan, latest.id + def run(self, milestone: int) -> ExecutorOutcome: model = self._roles.resolve("executor") + plan, after = self._scope() system = ( f"{EXECUTOR_SYSTEM}\n\n{self._workspace.describe()}" f"\n\nYou are executing milestone {milestone}." ) + if plan is not None: + # In the system prompt, so the context budget never drops it. + system += f"\n\nThe approved plan:\n{plan}" consecutive_failures = 0 did_work = False challenged = False @@ -162,7 +184,7 @@ class LLMExecutor: return ExecutorOutcome("stopped", f"before turn {turn + 1}") try: produced = self._client.complete( - self._store.read_events(), + self._store.read_events(after=after), tools=self._tools(), model=model, system=system, diff --git a/nisse-dev/src/nisse_dev/llm.py b/nisse-dev/src/nisse_dev/llm.py index 7f1542d..8cec4bf 100644 --- a/nisse-dev/src/nisse_dev/llm.py +++ b/nisse-dev/src/nisse_dev/llm.py @@ -149,12 +149,15 @@ class LiteLLMClient: tools: list[dict[str, Any]] | None, model: str, system: str | None = None, + *, + max_tokens: int | None = None, ) -> list[CanonicalEvent]: body = self._renderer.render(events, tools, model, system=system) # A ceiling on this route too: without one a runaway turn is bounded # only by the provider's default, and the same truncation check has - # nothing to check against. - body["max_tokens"] = self._max_tokens + # nothing to check against. A caller that knows the model's limits + # (the context budget) passes a tighter one. + body["max_tokens"] = max_tokens or self._max_tokens response = _post_with_rate_limit_retry( self._http, "/chat/completions", body, self._sleep ) @@ -175,8 +178,8 @@ class LiteLLMClient: raise LLMError(f"malformed completion response: {exc}") from exc if choice.get("finish_reason") == "length": raise TruncatedResponse( - f"{requested_model} hit its output limit mid-answer; " - "raise llm.max_output_tokens" + f"{requested_model} hit its output limit mid-answer; raise " + "llm.max_output_tokens, or the model's max_output_tokens in LiteLLM" ) produced_by = data.get("model") or requested_model usage = data.get("usage") or {} @@ -258,9 +261,15 @@ class AnthropicPassthroughClient: tools: list[dict[str, Any]] | None, model: str, system: str | None = None, + *, + max_tokens: int | None = None, ) -> list[CanonicalEvent]: body = self._renderer.render( - events, tools, model, system=system, max_tokens=self._max_tokens + events, + tools, + model, + system=system, + max_tokens=max_tokens or self._max_tokens, ) response = _post_with_rate_limit_retry( self._http, "/v1/messages", body, self._sleep @@ -277,8 +286,8 @@ class AnthropicPassthroughClient: ) -> list[CanonicalEvent]: if data.get("stop_reason") == "max_tokens": raise TruncatedResponse( - f"{requested_model} hit the output limit of {self._max_tokens} " - f"tokens mid-answer; raise llm.max_output_tokens" + f"{requested_model} hit its output limit mid-answer; raise " + "llm.max_output_tokens, or the model's max_output_tokens in LiteLLM" ) content = data.get("content") if not isinstance(content, list): diff --git a/nisse-dev/src/nisse_dev/probe.py b/nisse-dev/src/nisse_dev/probe.py new file mode 100644 index 0000000..54ef411 --- /dev/null +++ b/nisse-dev/src/nisse_dev/probe.py @@ -0,0 +1,114 @@ +"""Check a model's declared context window against what its server does. + +LiteLLM's ``max_input_tokens`` for a local model is only a declaration; the +real limit is whatever the model server allocates. This sends one prompt +half again as long as the declared window and reads back what the server did +with it: + +- **It refused.** Over-long prompts fail loudly, which is the safe behaviour. +- **It cut the prompt.** The reported prompt size is where it cut, which is + the server's real window. Ollama does this, silently. +- **It took the whole prompt.** The server allows more than is declared, so + the declaration is conservative. + +Run it with ``scripts/probe_context.py``. It spends one small and one large +request on the model, so it refuses windows above a ceiling unless told +otherwise: commercial models' limits are already known to LiteLLM. +""" + +from dataclasses import dataclass + +import httpx + +# Of a probe's size, how much must come back for it to count as uncut. A +# chat template adds a little, and a tokenizer counts a little differently. +_UNCUT = 0.9 +# How close the server's cut must be to the declaration to count as a match. +_MATCH = 0.02 + + +@dataclass(frozen=True) +class ProbeResult: + verdict: str # "matches", "declared_too_high", "declared_conservative", ... + message: str + server_window: int | None = None + + +def interpret( + declared: int, sent: int, reported: int | None, error: str | None +) -> ProbeResult: + """What one over-long probe says about the server's window. + + ``sent`` is roughly how many tokens the probe held; ``reported`` is the + prompt size the server reported, if it answered; ``error`` its error + text, if it refused. + """ + if error is not None: + return ProbeResult( + "refuses", + f"the server refused a ~{sent}-token prompt ({error[:200]}). Over-long " + "prompts fail loudly here, which is safe; the declared window of " + f"{declared} was not contradicted.", + ) + assert reported is not None + if reported >= sent * _UNCUT: + return ProbeResult( + "declared_conservative", + f"the server took the whole ~{sent}-token prompt ({reported} reported), " + f"more than the declared {declared}. The declaration is safe; raise it " + "if you want the room.", + server_window=None, + ) + if abs(reported - declared) <= declared * _MATCH: + return ProbeResult( + "matches", + f"the server cut the prompt at {reported} tokens, matching the declared " + f"{declared}.", + server_window=reported, + ) + if reported < declared: + return ProbeResult( + "declared_too_high", + f"the server cut the prompt at {reported} tokens, but LiteLLM declares " + f"{declared}. Prompts between the two would be silently cut: set " + f"max_input_tokens to {reported}.", + server_window=reported, + ) + return ProbeResult( + "declared_conservative", + f"the server cut the prompt at {reported} tokens, above the declared " + f"{declared}. The declaration is safe; it could be raised to {reported}.", + server_window=reported, + ) + + +def _prompt(words: int) -> str: + return " ".join(f"item{i}" for i in range(words)) + "\n\nReply with just: ok" + + +def _send(http: httpx.Client, model: str, words: int) -> tuple[int | None, str | None]: + response = http.post( + "/chat/completions", + json={ + "model": model, + "messages": [{"role": "user", "content": _prompt(words)}], + "max_tokens": 5, + }, + ) + if response.status_code != 200: + return None, f"{response.status_code}: {response.text}" + usage = response.json().get("usage") or {} + return int(usage.get("prompt_tokens") or 0), None + + +def probe(http: httpx.Client, model: str, declared: int) -> ProbeResult: + """Measure how many tokens a word of the probe costs on this model, then + send a prompt half again as long as the declared window.""" + calibration_words = 200 + tokens, error = _send(http, model, calibration_words) + if error is not None or not tokens: + return ProbeResult("failed", f"a short calibration request failed: {error}") + per_word = tokens / calibration_words + sent = int(declared * 1.5) + reported, error = _send(http, model, int(sent / per_word)) + return interpret(declared, sent, reported, error) diff --git a/nisse-dev/src/nisse_dev/tools.py b/nisse-dev/src/nisse_dev/tools.py index 4468d47..a450f07 100644 --- a/nisse-dev/src/nisse_dev/tools.py +++ b/nisse-dev/src/nisse_dev/tools.py @@ -20,6 +20,8 @@ from dataclasses import dataclass from pathlib import Path from typing import Any +from nisse_dev.context import head_and_tail + # Tools that change or verify something. Reading a file or listing a # directory is a model looking around, not doing the work, so neither counts # as evidence that work was carried out. @@ -112,11 +114,7 @@ class Workspace: self._max_output_chars = max_output_chars def _truncate(self, text: str) -> str: - limit = self._max_output_chars - if len(text) <= limit: - return text - dropped = len(text) - limit - return f"{text[:limit]}\n[... {dropped} characters truncated]" + return head_and_tail(text, self._max_output_chars) def _resolve(self, path: str) -> Path: candidate = Path(path) diff --git a/nisse-dev/tests/test_agent_build.py b/nisse-dev/tests/test_agent_build.py index 655309e..4ba7548 100644 --- a/nisse-dev/tests/test_agent_build.py +++ b/nisse-dev/tests/test_agent_build.py @@ -26,6 +26,8 @@ TEST_CATALOG = ModelCatalog.from_payload( "model_info": { "litellm_provider": "anthropic", "supports_function_calling": True, + "max_input_tokens": 200000, + "max_output_tokens": 64000, }, } for alias in ("claude-opus-5", "claude-haiku-4-5") @@ -185,8 +187,8 @@ def test_delegation_uses_the_passthrough_route( agent = build_offline(tmp_path, origin, monkeypatch) delegation = agent.orchestrator._delegate - assert isinstance(delegation._client, AnthropicPassthroughClient) - assert isinstance(agent.orchestrator._executor._client, LiteLLMClient) + assert isinstance(delegation._client._inner, AnthropicPassthroughClient) + assert isinstance(agent.orchestrator._executor._client._inner, LiteLLMClient) agent.stop() @@ -194,8 +196,8 @@ def test_httpx_clients_are_not_shared_between_routes( tmp_path: Path, origin: Path, monkeypatch ) -> None: agent = build_offline(tmp_path, origin, monkeypatch) - normalized = agent.orchestrator._executor._client._http - passthrough = agent.orchestrator._delegate._client._http + normalized = agent.orchestrator._executor._client._inner._http + passthrough = agent.orchestrator._delegate._client._inner._http assert isinstance(normalized, httpx.Client) assert normalized is not passthrough # different base URLs and headers agent.stop() @@ -244,6 +246,7 @@ def test_a_tool_less_model_warns_but_still_boots(tmp_path: Path, monkeypatch) -> "model_info": { "litellm_provider": "openai", "supports_function_calling": False, + "max_input_tokens": 32768, }, } for alias in ("claude-opus-5", "claude-haiku-4-5") @@ -338,6 +341,7 @@ def test_tuning_reaches_every_component( assert delegation._caps.wall_clock_s == 60 assert executor._workspace._timeout_s == 300 assert executor._workspace._max_output_chars == 2000 - assert executor._client._max_tokens == 8192 - assert delegation._client._max_tokens == 8192 + assert executor._client._inner._max_tokens == 8192 + assert delegation._client._inner._max_tokens == 8192 + assert executor._client._max_output_tokens == 8192 agent.stop() diff --git a/nisse-dev/tests/test_config.py b/nisse-dev/tests/test_config.py index 03c27ad..264ce63 100644 --- a/nisse-dev/tests/test_config.py +++ b/nisse-dev/tests/test_config.py @@ -297,3 +297,12 @@ def test_sibling_paths_are_fine(tmp_path: Path) -> None: {**SECRETS, "NISSE_ENV": str(tmp_path / ".env")}, ) assert config.workspace.name == "project" + + +def test_context_tuning_is_read_and_checked() -> None: + config = build(tuning={"context": {"keep_recent_turns": 4, "safety_margin": 0.2}}) + assert config.tuning.context.keep_recent_turns == 4 + assert config.tuning.context.safety_margin == 0.2 + assert build().tuning.context.keep_recent_turns == 6 + with pytest.raises(ConfigError, match="safety_margin: must be less than 1"): + build(tuning={"context": {"safety_margin": 1}}) diff --git a/nisse-dev/tests/test_context.py b/nisse-dev/tests/test_context.py new file mode 100644 index 0000000..6bda0d0 --- /dev/null +++ b/nisse-dev/tests/test_context.py @@ -0,0 +1,294 @@ +"""The context budget (spec A10).""" + +import io +from pathlib import Path +from typing import Any + +import pytest + +from nisse_dev.canonical import ( + AssistantText, + CanonicalEvent, + GateOpened, + PlanSubmitted, + StoredEvent, + ToolCall, + ToolResult, + UserMessage, +) +from nisse_dev.catalog import ModelCatalog, UnknownLimits +from nisse_dev.context import ( + Budget, + BudgetedClient, + ContextOverflow, + estimate_tokens, + fit, + head_and_tail, +) +from nisse_dev.executor import LLMExecutor +from nisse_dev.llm import RoleMap +from nisse_dev.oplog import OpLog +from nisse_dev.state import StateStore + + +def stored(events: list[CanonicalEvent]) -> list[StoredEvent]: + return [StoredEvent(id=i + 1, ts="t", event=e) for i, e in enumerate(events)] + + +def turn(n: int, output: str) -> list[CanonicalEvent]: + """One executor turn: a command and its (possibly long) output.""" + return [ + ToolCall(call_id=f"c{n}", tool="run_command", arguments={"command": "make"}), + ToolResult(call_id=f"c{n}", result=output, success=True), + ] + + +def transcript(turns: int, output_chars: int) -> list[StoredEvent]: + events: list[CanonicalEvent] = [UserMessage(text="start milestone 1", source="t")] + for n in range(turns): + events += turn(n, f"output {n} " + "x" * output_chars) + return stored(events) + + +def budget(window: int, keep: int = 2) -> Budget: + return Budget( + window=window, max_output=500, keep_recent_turns=keep, safety_margin=0.1 + ) + + +def results(events: list[StoredEvent]) -> dict[str, str]: + return { + e.event.call_id: e.event.result + for e in events + if isinstance(e.event, ToolResult) + } + + +# --- keeping the end of long output --------------------------------------- + + +def test_head_and_tail_keeps_the_failure_summary_at_the_end() -> None: + output = ( + "exit code: 1\nstdout:\n" + + "collected 400 items\n" * 500 + + "FAILED tests/test_app.py::test_greet - AssertionError\n" + + "==== 1 failed, 399 passed ====\n" + ) + kept = head_and_tail(output, 1000) + assert len(kept) < 1100 + assert kept.startswith("exit code: 1") # the head keeps the exit code + assert "FAILED tests/test_app.py::test_greet" in kept + assert "1 failed, 399 passed" in kept + assert "characters omitted" in kept + + +def test_short_text_is_untouched() -> None: + assert head_and_tail("fine", 1000) == "fine" + + +# --- fitting, step by step ------------------------------------------------- + + +def test_a_transcript_that_fits_is_sent_unchanged() -> None: + events = transcript(turns=3, output_chars=100) + assert fit(events, "system", None, budget(100_000)) == events + + +def test_step_one_cuts_older_tool_results_and_keeps_recent_ones_whole() -> None: + events = transcript(turns=6, output_chars=3000) + fitted = fit(events, "system", None, budget(4000, keep=2)) + + kept = results(fitted) + assert len(kept) == 6 # nothing dropped yet + for older in ("c0", "c1", "c2", "c3"): + assert "characters omitted" in kept[older] + for recent in ("c4", "c5"): + assert "characters omitted" not in kept[recent] + + +def test_step_two_drops_the_oldest_turns_and_says_so() -> None: + events = transcript(turns=20, output_chars=3000) + fitted = fit(events, "system", None, budget(4000, keep=2)) + + first = fitted[0].event + assert isinstance(first, UserMessage) + assert "omitted to fit the context window" in first.text + kept = results(fitted) + assert "c0" not in kept # the oldest went + assert {"c18", "c19"} <= kept.keys() # the recent turns stayed + assert "characters omitted" not in kept["c19"] + # Every surviving result still has its call: a provider rejects an + # answer to a call it never saw. + calls = {e.event.call_id for e in fitted if isinstance(e.event, ToolCall)} + assert kept.keys() <= calls + assert estimate_tokens(fitted, "system", None) <= budget(4000).for_prompt + + +def test_step_three_refuses_what_cannot_fit() -> None: + # Even the two most recent turns are bigger than the whole window. + events = transcript(turns=3, output_chars=20_000) + with pytest.raises(ContextOverflow, match="last 2 turns"): + fit(events, "system", None, budget(2000, keep=2)) + + +def test_a_huge_system_prompt_alone_overflows() -> None: + events = transcript(turns=1, output_chars=10) + with pytest.raises(ContextOverflow): + fit(events, "s" * 50_000, None, budget(4096)) + + +def test_the_output_ceiling_and_margin_come_out_of_the_window() -> None: + b = Budget(window=4096, max_output=1024, keep_recent_turns=6, safety_margin=0.1) + assert b.for_prompt == int(4096 * 0.9) - 1024 + + +# --- the client wrapper ---------------------------------------------------- + + +class RecordingClient: + def __init__(self, reported_input: int = 0) -> None: + self.calls: list[dict[str, Any]] = [] + self.last_usage = {"input": reported_input, "output": 5} + + def complete(self, events, tools, model, system=None, *, max_tokens=None): + self.calls.append({"events": events, "max_tokens": max_tokens}) + return [AssistantText(text="ok", model=model)] + + +def test_the_ceiling_is_the_smaller_of_config_and_model() -> None: + inner = RecordingClient() + client = BudgetedClient(inner, lambda m: (4096, 1024), max_output_tokens=16384) + client.complete(transcript(1, 10), None, "llama3.2") + assert inner.calls[0]["max_tokens"] == 1024 + + client = BudgetedClient(inner, lambda m: (200_000, 64_000), max_output_tokens=16384) + client.complete(transcript(1, 10), None, "claude-haiku-4-5") + assert inner.calls[1]["max_tokens"] == 16384 + + +def test_the_wrapper_sends_the_fitted_events() -> None: + inner = RecordingClient() + client = BudgetedClient( + inner, lambda m: (4000, 500), keep_recent_turns=2, safety_margin=0.1 + ) + client.complete(transcript(20, 3000), None, "m", system="system") + sent = inner.calls[0]["events"] + assert "omitted to fit" in sent[0].event.text + + +def test_reported_usage_only_ever_raises_the_estimate() -> None: + events = transcript(1, 10) + estimate = estimate_tokens(events, None, None) + + under = RecordingClient(reported_input=estimate * 2) # we undercounted + client = BudgetedClient(under, lambda m: (100_000, None)) + client.complete(events, None, "m") + assert client._factors["m"] == pytest.approx(2.0, rel=0.05) + + over = RecordingClient(reported_input=estimate // 2) # we overcounted + under_then_over = BudgetedClient(over, lambda m: (100_000, None)) + under_then_over._factors["m"] = 2.0 + under_then_over.complete(events, None, "m") + assert under_then_over._factors["m"] == 2.0 # never loosened + + +def test_last_usage_passes_through() -> None: + client = BudgetedClient(RecordingClient(reported_input=7), lambda m: (100, None)) + assert client.last_usage["input"] == 7 + + +# --- the limits come from LiteLLM's catalog -------------------------------- + + +CATALOG = ModelCatalog.from_payload( + { + "data": [ + { + "model_name": "claude-haiku-4-5", + "litellm_params": {"model": "anthropic/claude-haiku-4-5"}, + "model_info": {"max_input_tokens": 200000, "max_output_tokens": 64000}, + }, + { + "model_name": "llama3.2", + "litellm_params": {"model": "ollama/llama3.2"}, + "model_info": {"max_input_tokens": 4096, "max_output_tokens": 1024}, + }, + { + "model_name": "undeclared", + "litellm_params": {"model": "ollama/qwen3"}, + "model_info": {"max_input_tokens": None}, + }, + ] + } +) + + +def test_the_catalog_reads_each_models_limits() -> None: + facts = CATALOG.facts("llama3.2") + assert (facts.max_input_tokens, facts.max_output_tokens) == (4096, 1024) + assert CATALOG.facts("undeclared").max_input_tokens is None + + +def test_a_role_on_a_model_without_a_window_is_refused() -> None: + CATALOG.require_limits({"executor": "llama3.2", "planner": "claude-haiku-4-5"}) + with pytest.raises(UnknownLimits, match="executor='undeclared'"): + CATALOG.require_limits({"executor": "undeclared", "planner": "llama3.2"}) + + +# --- the executor sees the approved plan and what came after --------------- + + +def test_the_executor_sees_only_its_milestone(tmp_path: Path) -> None: + store = StateStore(tmp_path / "state.db", oplog=OpLog("t", io.StringIO())) + for event in ( + UserMessage(text="use greet: a long planning conversation", source="t"), + PlanSubmitted(plan="arch plan", model="m"), + UserMessage(text="approve", source="t"), + PlanSubmitted(plan="milestone 1: write greet.py", model="m"), + GateOpened(gate="milestone-1", reason="plan submitted"), + UserMessage(text="start milestone 1", source="t"), + ): + store.append_event(event) + + class Workspace: + def describe(self) -> str: + return "workspace" + + seen: dict[str, Any] = {} + + class Client: + def complete(self, events, tools, model, system=None): + seen["events"] = [e.event for e in events] + seen["system"] = system + return [AssistantText(text="done", model=model)] + + roles = RoleMap(executor="m", planner="m", delegate="m", consultant="m") + LLMExecutor(Client(), store, roles, Workspace()).run(1) + + texts = [getattr(e, "text", None) for e in seen["events"]] + assert "start milestone 1" in texts + assert "use greet: a long planning conversation" not in texts + assert not any(isinstance(e, PlanSubmitted) for e in seen["events"]) + # The plan rides in the system prompt, which the budget never drops. + assert "The approved plan:\nmilestone 1: write greet.py" in seen["system"] + assert "arch plan" not in seen["system"] + + +def test_a_request_that_cannot_fit_stops_the_milestone(tmp_path: Path) -> None: + # Never sent: the executor stops and asks the operator, like any failed + # call, rather than handing a server a prompt it would silently cut. + store = StateStore(tmp_path / "state.db", oplog=OpLog("t", io.StringIO())) + store.append_event(UserMessage(text="start milestone 1", source="t")) + + class Workspace: + def describe(self) -> str: + return "workspace " + "w" * 20_000 # a system prompt too big to fit + + inner = RecordingClient() + client = BudgetedClient(inner, lambda m: (4096, 1024)) + roles = RoleMap(executor="m", planner="m", delegate="m", consultant="m") + outcome = LLMExecutor(client, store, roles, Workspace()).run(1) + + assert outcome.status == "escalate" + assert "fit" in outcome.note + assert inner.calls == [] diff --git a/nisse-dev/tests/test_probe.py b/nisse-dev/tests/test_probe.py new file mode 100644 index 0000000..ca1be0d --- /dev/null +++ b/nisse-dev/tests/test_probe.py @@ -0,0 +1,53 @@ +import json + +import httpx + +from nisse_dev.probe import interpret, probe + + +def test_a_cut_at_the_declared_window_matches() -> None: + result = interpret(declared=4096, sent=6144, reported=4096, error=None) + assert result.verdict == "matches" + assert result.server_window == 4096 + + +def test_a_cut_below_the_declaration_says_what_to_set() -> None: + # The dangerous case: prompts between the two would be silently cut. + result = interpret(declared=8192, sent=12288, reported=4096, error=None) + assert result.verdict == "declared_too_high" + assert "set max_input_tokens to 4096" in result.message + + +def test_a_cut_above_the_declaration_is_conservative() -> None: + result = interpret(declared=4096, sent=6144, reported=5000, error=None) + assert result.verdict == "declared_conservative" + + +def test_an_uncut_prompt_means_the_server_allows_more() -> None: + result = interpret(declared=4096, sent=6144, reported=6100, error=None) + assert result.verdict == "declared_conservative" + assert result.server_window is None + + +def test_a_refusal_is_the_safe_behaviour() -> None: + result = interpret(declared=4096, sent=6144, reported=None, error="400: too long") + assert result.verdict == "refuses" + + +def test_probe_calibrates_then_sends_half_again_the_window() -> None: + """Against a fake Ollama: 2 tokens a word, silently cut at 4096.""" + sizes: list[int] = [] + + def handler(request: httpx.Request) -> httpx.Response: + words = len(json.loads(request.content)["messages"][0]["content"].split()) + sizes.append(words * 2) + reported = min(words * 2, 4096) + return httpx.Response(200, json={"usage": {"prompt_tokens": reported}}) + + with httpx.Client( + base_url="http://litellm", transport=httpx.MockTransport(handler) + ) as http: + result = probe(http, "llama3.2", declared=4096) + + assert result.verdict == "matches" + assert sizes[1] >= 4096 * 1.4 # the probe really exceeded the window diff --git a/nisse-dev/tests/test_tools.py b/nisse-dev/tests/test_tools.py index 6edb28e..18b14f1 100644 --- a/nisse-dev/tests/test_tools.py +++ b/nisse-dev/tests/test_tools.py @@ -126,7 +126,7 @@ def test_stdin_is_closed_so_prompts_fail_fast(ws: Workspace) -> None: def test_long_output_is_truncated(ws: Workspace) -> None: outcome = ws.run_command("printf 'x%.0s' $(seq 1 20000)") assert len(outcome.result) < MAX_OUTPUT_CHARS + 200 - assert "truncated" in outcome.result + assert "characters omitted" in outcome.result def test_dispatch_unknown_tool_and_missing_argument(ws: Workspace) -> None: diff --git a/scripts/probe_context.py b/scripts/probe_context.py new file mode 100644 index 0000000..f23704b --- /dev/null +++ b/scripts/probe_context.py @@ -0,0 +1,74 @@ +"""Check a model's declared context window against what its server does. + + uv run python scripts/probe_context.py ALIAS [--config PATH] + +ALIAS is a LiteLLM model alias. The LiteLLM URL and token come from the agent +config (NISSE_CONFIG, ./nisse-dev.yaml, or --config), with its secrets from +the environment or .env. See docs/Tuning-Guide.md, "Declaring a local model's +limits", and nisse_dev/probe.py for what the verdicts mean. +""" + +import argparse +import os +import sys +from pathlib import Path + +import httpx + +from nisse_dev.catalog import ModelCatalog +from nisse_dev.config import AgentConfig, ConfigError +from nisse_dev.envfile import load_env_file +from nisse_dev.probe import probe + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + parser.add_argument("alias", help="LiteLLM model alias to probe") + parser.add_argument("--config", help="agent config file to take LiteLLM from") + parser.add_argument( + "--max-probe-tokens", + type=int, + default=64_000, + help="refuse to send a prompt larger than this (default: %(default)s)", + ) + args = parser.parse_args() + + load_env_file(Path(os.environ.get("NISSE_ENV", ".env")).expanduser()) + try: + config = ( + AgentConfig.from_file(Path(args.config)) + if args.config + else AgentConfig.load() + ) + except ConfigError as exc: + print(exc, file=sys.stderr) + return 2 + + catalog = ModelCatalog.discover(config.llm_base_url, config.llm_token) + declared = catalog.facts(args.alias).max_input_tokens + if declared is None: + print( + f"LiteLLM declares no max_input_tokens for {args.alias}; set one first", + file=sys.stderr, + ) + return 1 + if declared * 1.5 > args.max_probe_tokens: + print( + f"{args.alias} declares {declared} tokens; probing would send about " + f"{int(declared * 1.5)}. Raise --max-probe-tokens to allow it.", + file=sys.stderr, + ) + return 1 + + with httpx.Client( + base_url=config.llm_base_url.rstrip("/"), + headers={"Authorization": f"Bearer {config.llm_token}"}, + timeout=600, + ) as http: + result = probe(http, args.alias, declared) + print(f"{args.alias}: {result.message}") + return 0 if result.verdict in ("matches", "refuses", "declared_conservative") else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/integration/test_llm_live.py b/tests/integration/test_llm_live.py index db0ef5c..1f96412 100644 --- a/tests/integration/test_llm_live.py +++ b/tests/integration/test_llm_live.py @@ -3,7 +3,9 @@ import pytest from nisse_dev.canonical import AssistantText, StoredEvent, UserMessage +from nisse_dev.catalog import ModelCatalog from nisse_dev.config import AgentConfig +from nisse_dev.context import BudgetedClient from nisse_dev.llm import LiteLLMClient pytestmark = pytest.mark.integration @@ -25,3 +27,35 @@ def test_executor_model_completes_via_litellm( texts = [e.text for e in result if isinstance(e, AssistantText)] assert texts, f"no assistant text in {result!r}" assert "pong" in texts[0].lower() + + +def test_declared_limits_reach_a_budgeted_request( + litellm: LiteLLMClient, live_config: AgentConfig +) -> None: + """The whole chain: LiteLLM's model_info, the catalog, the budget, and a + real request through the proxy with the ceiling the budget chose.""" + catalog = ModelCatalog.discover(live_config.llm_base_url, live_config.llm_token) + roles = live_config.roles.pinned() + catalog.require_limits(roles) # every test role has a declared window + + model = live_config.roles.executor + facts = catalog.facts(model) + client = BudgetedClient( + litellm, + lambda m: ( + catalog.facts(m).max_input_tokens, + catalog.facts(m).max_output_tokens, + ), + max_output_tokens=live_config.max_output_tokens, + ) + assert client.budget_for(model).window == facts.max_input_tokens + + events = [ + StoredEvent( + id=1, ts="t", event=UserMessage(text="Reply with exactly: pong", source="t") + ) + ] + produced = client.complete(events, None, model, system="Be terse.") + texts = [e.text for e in produced if isinstance(e, AssistantText)] + assert texts and "pong" in texts[0].lower() + assert client.last_usage.get("input", 0) > 0