Merge pull request 'Fit every model request to its context window' (#44) from context-budget into main
CI / checks (push) Successful in 57s
CI / checks (push) Successful in 57s
Reviewed-on: dongar/nisse#44 Reviewed-by: Don Garrett <dgarrett@acm.org>
This commit was merged in pull request #44.
This commit is contained in:
commit
389b71a5d1
17 files changed
+1082
-45
No files matched your search
+104
-16
@@ -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
|
||||
`<arch>.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.
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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,
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
@@ -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}})
|
||||
@@ -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 == []
|
||||
@@ -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
|
||||
@@ -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:
|
||||
|
||||
@@ -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())
|
||||
@@ -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
|
||||
Reference in new issue
Block a user