Local LLM Model Router: NeuralWatt routing service + context-aware framing and pruning #6

Merged
alee merged 3 commits from neuralwatt-router-service into main 2026-08-24 02:30:47 +00:00
10 changed files with 1529 additions and 8 deletions

View File

@@ -0,0 +1,142 @@
# Spec: two design decisions deferred out of the context-pruning/framing fix pass
**Origin.** The fix pass for
[`context-pruning-and-framing-review.md`](../code_reviews/context-pruning-and-framing-review.md)
(reviewed in
[`context-pruning-and-framing-fixes-review.md`](../code_reviews/context-pruning-and-framing-fixes-review.md))
explicitly declined two items as design decisions rather than confirmed
bugs: the `TaskRequest.context` dual-use ambiguity (review finding #6,
partially mitigated but not resolved), and classifying once per session
instead of once per message (an existing item on the project's own "what's
NOT built yet" list). Both are forward-looking — no bug is being reported
here, no code changes accompany this document.
---
## 1. `TaskRequest.context` carries two unrelated meanings
### The problem
`context` was originally, and is still documented as, "assembled context
(docs/code) to send with the task" — a `/route`/`/dispatch` caller pastes
reference material alongside a task description. `chat_completions` now
also uses the same field to carry the prior conversational turn
(`_previous_context`), so a short follow-up like "Yes" inherits that turn's
complexity instead of being classified as trivial in isolation.
The current fix (moving the framing instruction from the static system
prompt into `_classifier_user_content`'s output, appended only when
`context` is non-empty) narrowed the blast radius — the instruction no
longer reaches every classify() call, only ones that actually supply
`context` — but it didn't resolve which of the two meanings a given
`context` value has. A `/route` caller pasting 800 lines of Django code
under `task="Refactor this"` gets the same "a short follow-up continues the
prior turn, classify by the CONTEXT's complexity" instruction that exists
for the conversational case.
### Why this might not need fixing
The instruction's general principle — "short task + large/complex context
implies a non-trivial task" — is arguably a reasonable heuristic for the
docs-paste case too, even though it was written for conversational
follow-ups. It has never been measured against real `/route`/`/dispatch`
traffic with `context` set. This project's own stated epistemics apply
directly here: `POST /outcome` is "the only ground truth," and several
sections of the README describe correcting an assumption only after
measuring it, not before. Speculating about which framing is "more
correct" without a measurement repeats the mistake the project has already
named and moved past.
### Options, if it turns out to matter
| Option | Sketch | Trade-off |
|---|---|---|
| A. Decouple the mechanisms | Keep `TaskRequest.context` as the public docs/code field, untouched. Give `chat_completions`'s conversational-continuation signal its own internal path — e.g. `classify()` gains a private `prior_turn: Optional[str]` parameter distinct from `context`, so the instruction is only ever built for genuine conversational continuations | Cleanest semantically; requires touching `classify()`'s signature and both call sites; `/route`/`/dispatch` behavior is provably unaffected |
| B. Tag the field | Add a `context_kind: Literal["reference", "prior_turn"] = "reference"` field to `TaskRequest`; only `chat_completions`'s internal calls set `"prior_turn"`; the instruction is only appended for that kind | Smaller diff than A; adds a field to the public request model that external callers never need to know about |
| C. Leave as-is, measure | No code change. Watch `route_decisions` (already logs `task_category`/`task_tier`/`source` per request) for `/route`/`/dispatch` calls that supply `context` and see whether their tier/category looks skewed relative to before this change | Zero engineering cost; consistent with the project's own "measure before correcting" pattern; only viable if `/route`/`/dispatch`-with-`context` traffic is common enough to be observable |
**Recommendation:** C first. This codebase already has the instrumentation
(`route_decisions`, `feedback.py` folding in outcomes) to tell whether this
is a real problem instead of a theoretical one, and the existing pattern in
this project — cost-vs-eco, tier-from-price, the whole classifier-model
swap — is "measure, then fix what the measurement shows," not "fix what
looks fishy." If `/route`/`/dispatch`-with-`context` traffic turns out to
be rare or nonexistent, this is not worth A's or B's added surface at all.
---
## 2. Classify once per session, not once per message
### The problem, restated from the README
> ~10s of local overhead on every message is a real tax for an interactive
> agent... Still unaddressed: classify once per session rather than per
> message, cache by prompt hash, or skip classification for short prompts.
With the local `mistral-nemo` classifier this is now ~1.7s/call
(§"The classifier is the latency floor"); with a cloud classifier
(measured against `deepseek-v4-flash`) it's ~1.0s. Either way, every single
turn in a long agent session pays this again, even though the session's
*category* (coding_general, debugging, etc.) rarely changes turn to turn —
what changes is mostly the token count, which `chat_completions` already
measures directly via `estimate_prompt_tokens` and doesn't need the
classifier for.
### The mechanism already half-exists
`route()`'s override branch (`dispatcher.py:673`) already skips
`classify()` entirely whenever `task_category`, `task_tier`, and
`required_context_tokens` are all supplied — this is exactly what the
measured-context reroute (fixed in review finding #4) uses today, just
within a single request. Session-level caching is the same mechanism
applied across requests: classify once, store `(task_category, task_tier)`
keyed by session, and on every later turn in that session call `route()`
with the cached category/tier plus a **freshly measured**
`required_context_tokens` (which is cheap — pure token counting, no model
call) — landing on the override branch and skipping the classifier
round-trip entirely.
### Design questions to settle before implementing
1. **Session identity.** `session_fingerprint`/`session_directory`
(dispatcher.py) already derive a session identity from message content
for observation purposes. Whether that's the right key for a
*classification* cache (vs. e.g. a client-supplied session id, if
opencode's protocol carries one) needs checking — a cache keyed on the
wrong signal either misses constantly (no benefit) or collides across
genuinely different sessions (wrong category persists into unrelated
work).
2. **Invalidation.** A session's task can genuinely change category mid-way
(debugging turns into a docs-writing turn turns into refactoring). Pure
"classify once, cache forever" risks staleness. Candidate triggers to
re-classify: a large jump in `required_context_tokens` between turns (a
proxy for "something new started"), a fixed number of turns (e.g.
re-classify every 20), or a TTL. This needs the same "measure before
deciding" treatment as everything else in this project — a cheap thing
to instrument via `route_decisions.source` (add a `"cached"` value
alongside `"classifier"`/`"override"`/`"fallback"`) and watch category
drift over real sessions before picking a policy.
3. **Interaction with escalation and retries.** `apply_escalation` currently
runs on every fresh classification. A cached category/tier bypasses it
entirely on cache-hit turns — need to decide whether escalation state
should also be cached per-session or re-evaluated each turn (it's cheap,
pure Python, so probably always re-evaluate rather than cache).
4. **Storage.** In-memory dict keyed by session identity is the obvious
starting point (matches the process lifetime of the dispatcher; a
restart just means the next turn in every active session re-classifies
once, which is a safe failure mode) — no new persistence layer needed
unless multi-process deployment becomes a requirement.
### Recommendation
Worth building — the latency case is strong and the mechanism is a small
extension of code that already exists (the override branch) rather than a
new one. But settle invalidation policy (§2) with a measurement pass first,
the same way `context_framing`'s default-on-and-measure and the classifier
model swap were each decided by running both and comparing, not by
argument. A reasonable first cut: cache with no re-classification, ship
behind a config flag (default off, matching every other new-and-unproven
knob in this project — `pinch.enabled`, `min_tool_proficiency`), watch
`route_decisions` on real sessions for category drift, then decide whether
any invalidation trigger is actually needed or whether "classify once,
never again" is good enough in practice.

View File

@@ -0,0 +1,125 @@
# Review: fixes for `context-pruning-and-framing-review.md`
**What it was reviewing:** the working-tree diff (7 files, unstaged, not yet
committed) implementing fixes for the 11 confirmed bugs and 4 cleanup items
from [`context-pruning-and-framing-review.md`](context-pruning-and-framing-review.md),
done by opencode driven through the router. Reviewed by reading the full
diff against each numbered finding, re-deriving the fix logic by hand, and
reproducing behavior directly against `context_prune.py` and `dispatcher.py`
rather than trusting the diff's comments. Full suite: 544/544 passing (up
from 529; 15 new regression tests, one per fixed finding).
## Verdict: 10 of 11 bugs correctly fixed and verified; 1 new regression found in the process
### Fixed and verified by direct reproduction
| # | Original finding | Fix | Verified |
|---|---|---|---|
| 1 | `_previous_context` returned raw tool output as classifier context | Now walks back to the nearest `role == "assistant"` message only, skipping `system`/`user`/`tool`, and reuses `context_prune.extract_text` instead of a second copy of the block parser (also fixes #13) | `test_previous_context_excludes_tool_role_messages` — read directly, correct |
| 2 | Same function returned the system prompt on a session's first turn | Same fix as #1 (system role is no longer eligible at all) | `test_previous_context_excludes_system_prompt` — correct |
| 3 | Long-tool-result trim went negative and grew the message when `max_summarize_chars < 3000` | Two layers: `config.py` now validates `max_summarize_chars >= 3000` at load; `context_prune.py` also defensively checks `trimmed > 0 and len(elided) < len(text)` before using the elided form, falling through to the placeholder otherwise, so a direct call bypassing config can't hit it either | Reproduced by hand: `max_summarize_chars=100`, 1200-char input → output is strictly shorter, `tokens_saved >= 0`, no negative marker. Matches `test_max_summarize_chars_below_3000_never_grows_or_goes_negative` |
| 4 | Routing/tier/cost decided on unpruned tokens; pinch's savings never reached the decision | `prune_context` now runs once inside the `wants_routing` block, before `measured = estimate_prompt_tokens(send_messages, ...)`; the already-pruned `send_messages` is reused at dispatch, and the passthrough path (which never goes through `wants_routing`) still prunes on its own. Confirmed this doesn't double-prune and doesn't touch `messages` in place (only `send_messages` is reassigned; `session_directory(messages)`, `_last_user_text(messages)`, and `_run_local_vision(messages, ...)` all still read the original, full conversation) | Read the full function end-to-end to confirm the no-double-prune and no-mutation properties; matches `test_pruned_list_measures_smaller_than_the_raw_conversation` and `test_pinch_disabled_leaves_the_measured_size_unchanged` |
| 5 | `extract_text` read `image_url` blocks as empty, undercounting image-bearing messages | Now contributes the block's `url` string; nested `tool_result.content` lists are flattened recursively too | Correct as far as counting goes — **but see the new finding below**, this fix combined with #11's fix produces a different bug |
| 7 | `context=prev_context` was dead on the measured-context reroute `route()` call | Parameter removed from that call, with a comment explaining the override branch never reads it | Confirmed by re-reading `route()`'s branch condition; no behavior change, just removes a misleading dead arg |
| 8 | `dropped` stat/docstring claimed removal; code always replaced content in place | `dropped` counter removed; both trim paths now fold into `summarized`; docstring updated to say "trimmed or summarized (never removed)" | `test_stats_use_summarized_and_length_is_always_preserved` — correct |
| 9 | `classifier.system_prompt` unconditionally described the "Context:"/"Message:" framing even when `context_framing: false` | The framing instruction was removed from the static system prompt entirely and is now appended to `_classifier_user_content`'s output only when `context` is actually present, and worded generically enough (doesn't reference the specific labels) to be correct under both framing modes | Read both `config.yaml` and `_classifier_user_content`; the instruction now travels with the content that justifies it rather than being unconditional boilerplate |
| 10 | Zero user-role messages collapsed `protected_from` to `len(messages)`, protecting nothing — even the newest tool result was prunable | New branch: when there's no user message, protect the trailing `keep_last_turns` *tool* results instead | Reproduced by hand with 5 tool results / `keep_last_turns=2`: the 2 most recent survive verbatim, the 3 older ones are trimmed. Matches `test_zero_user_messages_protect_newest_tool_result` |
| 12 | `_turn_of()` was dead code | Removed | Confirmed, no remaining references |
| 13 | `_previous_context` duplicated `extract_text`'s block-parsing logic | Now imports and calls `context_prune.extract_text` directly | Confirmed via the `#1` fix above |
| 14 | `classify()`'s outer `if not context` duplicated a check `_classifier_user_content` already makes | Outer branch removed; `classify()` now calls `_classifier_user_content` unconditionally | `test_classify_user_content_invariant_when_no_context` — correct |
| 15 | Pinch's 4 defaults declared independently in `context_prune.py`, `config.py`, and `config.yaml` | `context_prune.py`'s module-level `DEFAULT_*` constants removed; `prune_context`'s default args now read from `PinchConfig.model_fields[...].default`, one source of truth | Confirmed — `CHARS_PER_TOKEN` itself is still a separate mirrored constant (unaddressed, but this was always the minor half of #15) |
Finding #11 (structured content silently flattened to a string, dropping
non-text blocks) is **also fixed** in the sense that it was scoped —
`image_url` and other non-text blocks now survive trimming instead of being
discarded — but the fix interacts badly with #5's fix, below.
### New finding: image-bearing tool results aren't actually shrunk, and the stats lie about it
`context_prune.py``extract_text` (39-62, image_url branch) + the trim
branch (155-175) + `_with_text` (91-114)
Fixing #5 (count `image_url` bytes into the size estimate) and #11 (never
drop non-text blocks) independently make sense, but composed, they produce
a case neither fix's own test covers: a tool result whose `image_url` block
is *large enough to actually matter*.
`extract_text` joins every block — text **and** the image's raw `url`
string — into one combined string. That combined string is what gets
head+tail-sliced for the elided/placeholder replacement. `_with_text` then
writes that replacement into the message's `text` block only, and copies
every non-text block (the `image_url` one) through **completely
unchanged** — full size, untouched. Net effect for an image-bearing tool
result: the image (almost always the dominant contributor to size) is never
actually reduced, while the `tokens_saved` stat is computed from the
combined length including the image bytes, so it reports a large "savings"
that didn't happen. Reproduced directly:
```python
huge_b64 = "A" * 200000
messages = [
{"role": "user", "content": "describe this"},
{"role": "assistant", "content": None, "tool_calls": [{"id": "1"}]},
{"role": "tool", "name": "screenshot", "content": [
{"type": "text", "text": "here is the screenshot"},
{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{huge_b64}"}},
]},
{"role": "user", "content": "now what"},
]
pruned, stats = prune_context(messages, budget_tokens=10, keep_last_turns=0)
# stats == {'pruned': True, 'original_tokens': 66687, 'final_tokens': 67690,
# 'tokens_saved': 65671, 'summarized': 1}
```
`final_tokens` (67690) is **larger** than `original_tokens` (66687) — the
message got bigger, not smaller — while `tokens_saved` claims 65,671 tokens
saved. The image_url block comes out with its full 200,022-char url
untouched, and the text block now reads `"here is the screenshot
data:image/png;base64,AAAA...[187,022 chars trimmed...]AAAA"` — a raw
fragment of the base64 payload, sliced at an arbitrary byte boundary, now
sitting in a `text` field. This is worse than a no-op: it doesn't reduce
what's shipped upstream, it corrupts the text block with base64 noise, and
it actively misreports the stat the whole feature exists to produce.
The existing regression test for #11
(`test_structured_tool_result_keeps_image_block_when_trimmed`) doesn't
catch this because its fixture image is 4 characters (`"AAAA"`) — small
enough that the leakage and false accounting are present but invisible.
Any real screenshot or image tool result (tens of KB to MB of base64) would
hit this.
The fix likely belongs in `extract_text`/the trim logic together:
`image_url` blocks should count toward the size estimate (that part of #5
is correct — undercounting was the original bug) but should **not** be
included in the text that gets sliced for the elided/placeholder
replacement, and should be excluded from (or separately accounted in)
whatever text ends up in the `text` block. Whether an oversized image
itself should also become a trim candidate (a placeholder replacing the
`image_url` block, not just the text) is a design question worth deciding
explicitly rather than falling out of two unrelated fixes' composition.
### Not fully resolved (matches what was flagged as deferred)
Original finding #6`TaskRequest.context`'s dual purpose (docs/code paste
for `/route`/`/dispatch` vs. the new conversational-continuation use in
`chat_completions`) — got a real, useful partial fix: the framing
instruction is no longer unconditionally injected into every classify()
call (that part is #9, now fixed), and the field's dual use is now
documented on `TaskRequest.context` itself. But the underlying semantic
question is still open: a `/route` or `/dispatch` caller who passes
`context=<pasted docs/code>` (the field's original, documented purpose)
still gets the classifier told "a short follow-up continues the prior turn,
classify by the CONTEXT's complexity" — a rule written for conversational
continuation, appended regardless of which of the two meanings this
particular caller's `context` actually carries. This tracks with what was
flagged as a deferred design decision rather than a missed fix.
## Recommendation
Don't commit yet — the image/large-tool-result finding above is a real,
reproducible regression (not a pre-existing issue; it's new from this
round's own fixes), and it directly undermines the stat this feature exists
to produce. Route it back for another pass; everything else here is solid
and doesn't need to be touched again. The two deferred design items are
covered in a separate forward-looking spec rather than this after-the-fact
report.

View File

@@ -0,0 +1,208 @@
# Review: fbcc636, "context-aware framing and relevance-based context pruning"
**What it was reviewing:** the two features ported from the MIT-licensed
`alexrudloff/llmrouter` project — context-aware classifier framing
(`_previous_context` / `_classifier_user_content` in `dispatcher.py`) and
relevance-based context pruning (`context_prune.py`, "pinch"). Reviewed with
`/code-review` at xhigh effort (10 finder angles + 1-vote verify + a gap
sweep), then spot-verified directly: read every file involved, and
reproduced the two most severe findings against the actual `prune_context`
function rather than trusting the description.
## Verdict: full suite is green (529/529), but that's not the same as correct
`pytest` passes clean, including the 10 new `tests/test_context_prune.py`
cases. None of those cases exercise the shapes that break, though: a
`max_summarize_chars` below ~3000, a conversation with zero user messages,
tool content as `image_url` blocks, or the dispatcher-level interaction
between `_previous_context` and `route()`. The suite passing means the
happy path works, not that the port is safe on real agent traffic — which is
exactly the traffic this router's own README says is dominant (tool calls,
long sessions, ~92% cache hits on huge prompts).
Both features are **off by default** (`pinch.enabled: false`,
`classifier.context_framing: true` — framing is on, pruning is off), so
nothing here is live yet. But `context_framing` defaulting **on** means the
framing bugs (#1, #2 below) are already affecting every routed
`/v1/chat/completions` call today.
## Confirmed bugs — verified directly, not just reported
### 1. `_previous_context` feeds the classifier raw tool output, not "the prior turn" — and it's the common case, not an edge case
`dispatcher.py:1560-1588`
```python
for i in range(len(messages) - 2, -1, -1):
if messages[i].get("role") == "user":
continue
...
```
Only `role == "user"` is skipped while walking backward from `messages[-2]`.
`role == "tool"` is not. In the standard OpenAI tool-loop shape —
`assistant(tool_call) -> tool(result) -> user("ok fix it")``messages[-2]`
*is* the tool message, so `_previous_context` returns the raw tool-result
payload (file contents, grep output, JSON), truncated to 200 chars, as
`"Context:"` for the classifier. This isn't a malformed-conversation edge
case; it's what every agent client running a tool loop produces on its very
next turn. The function's own docstring says it exists to let the
classifier "inherit the [prior] turn's complexity" — it's inheriting
arbitrary tool output instead.
### 2. Same function returns the system prompt as "previous context" on a session's first turn
`dispatcher.py:1570-1572`
For `messages = [system, user]`, the backward scan starts at index 0 (the
system message), which isn't `role == "user"`, so it's not skipped — its
first 200 chars come back as the "prior turn." Every fresh session's first
message gets classified with `"Context: <fragment of the system prompt>"`
prepended, which for opencode is ~32K chars of tool definitions.
### 3. Pinch's long-tool-result trim goes negative and *inflates* the message once `max_summarize_chars < 3000`
`context_prune.py:145-153`; `config.py:293-299` (no validator on this field, unlike its two siblings)
```python
head = text[:1500]
tail = text[-1500:]
trimmed = len(text) - 3000
```
This assumes any text reaching this branch is longer than 3000 chars, but
the only guard to get here is `len(text) > max_summarize_chars`, and
`max_summarize_chars` has no lower-bound validator (`budget_tokens` and
`keep_last_turns` both do). Reproduced directly:
```
max_summarize_chars=100, tool content = 800 chars
-> pruned tool content = 1632 chars (grew)
-> stats["tokens_saved"] = -734 (negative)
-> marker literally reads "[...-2,200 chars trimmed...]"
```
Below the default (4000) this can't trigger, but there's nothing stopping
an operator from setting it lower, and when they do, pruning does the
opposite of its job.
### 4. Model/tier selection runs on unpruned tokens; pinch's savings never reach the decision that spends the money
`dispatcher.py` — routing at 1913-1954 vs. `prune_context()` at 2079
`route()` is called (twice — see #7) using `estimate_prompt_tokens(messages, ...)`
on the **full, unpruned** message list, and that's what drives tier
selection, the context-window hard filter, and the cost tiebreak.
`prune_context()` doesn't run until line 2079, well after `decision.selected`
is fixed — it only shrinks the payload actually sent to the model already
picked. So a long tool-heavy session can get routed to a pricier
large-context model based on its pre-pruned size, even though pinch would
have brought the real outgoing request in well under budget. This isn't
wrong on invalid input, it's an ordering bug: pinch's entire stated purpose
(reduce shipped tokens on long sessions) doesn't influence the one decision
where that would save money.
### 5. `extract_text` silently treats `image_url` content blocks as empty
`context_prune.py:39-53`, used for both `orig_tokens` and the per-message trim decision
Only `type == "text"` and `type == "tool_result"` blocks are read; anything
else (including `image_url`) contributes `""`. Two consequences: token
estimates can undercount a session that's actually huge (a tool result full
of base64 image data reads as 0 tokens, so `orig_tokens` may never cross
`budget_tokens` and pruning never triggers), and if pruning does trigger for
other reasons, that same message reads as `len(text) == 0 <= max_summarize_chars`,
so the "only replace if the placeholder is shorter" check (`23 < 0`) is
false and the giant blob is left completely untouched while smaller
genuine text results nearby get trimmed.
### 6. `TaskRequest.context` now has two incompatible meanings sharing one field
`dispatcher.py:151-153` (field docstring: `"Assembled context (docs/code) to send with the task"`) vs. the new use in `chat_completions` (prior conversation turn) vs. `config.yaml:395-398` (system prompt instructions written for the second meaning)
`/route` and `/dispatch` callers have always been able to pass `context` as
pasted docs/code (`dispatch_endpoint` splices it verbatim into a system
message). The classifier's system prompt was changed globally to say
"classify by the CONTEXT's complexity, treating a short message with
complex context as inheriting that complexity" — a rule written for the
conversational-continuation case, but it now applies unconditionally to
every existing `context=<pasted code>` caller too, since it's the same
field and the same prompt. Not obviously wrong, but untested for that
existing use and not called out anywhere as a behavior change to it.
## Confirmed via reproduction — edge cases in pruning itself
### 7. `context=prev_context` is dead weight on the second (measured-context) `route()` call
`dispatcher.py:1942-1954` vs. `route()`'s branch condition at line 673
`route()` only reads `req.context` inside the `classify()` branch, which is
skipped whenever `task_category`, `task_tier`, and `required_context_tokens`
are all supplied together — which the reroute at 1942 always does (it
copies the first decision's category/tier and sets
`required_context_tokens=measured`). So `context=prev_context` on that call
is passed and never read. Harmless today since the first `route()` call
already consumed it, but it means a future fix to `_previous_context` (#1/#2
above) would silently not apply here, and there's nothing marking the
parameter as inert.
### 8. `dropped` doesn't mean dropped
`context_prune.py:100-107` (docstring/stat name) vs. `133-144` (actual behavior)
The docstring and `config.yaml` both say short old tool results are
"dropped entirely," and the counter is literally named `dropped`, but the
code never removes a message from the list — it always replaces `content`
with a placeholder string in place. `len(pruned) == len(messages)` always,
regardless of `dropped`. Any code (or test) later written against the
documented contract — e.g. `assert len(pruned) == len(messages) - dropped`
— would be wrong on every request that drops anything.
### 9. Stale prompt instructions when the framing opt-out is used
`config.yaml:395-398`
The classifier's `system_prompt` unconditionally describes the
`"Context: <prior>\n---\nMessage: <current>"` label format, but that layout
is only actually produced when `classifier.context_framing: true`. Set it
`false` (the documented way to get the legacy `"task\n\n--- context ---\ncontext"`
layout) and the classifier still receives instructions describing a format
it will never see.
### 10. Zero user messages -> the recency guard protects nothing
`context_prune.py:109-115`
Reproduced: with no `role == "user"` message anywhere in the conversation,
`num_protected_turns` collapses to 0 and `protected_from = len(messages)`,
which no real index ever reaches — so the "always keep if `i >= protected_from`"
branch never fires. The single most recent tool result (the one the next
turn actually needs) becomes eligible for trimming, same as the oldest one:
```
messages = [system, assistant(tool_call), tool(20000 chars)]
-> the only tool result gets summarized down, despite being the newest
```
### 11. Trimming silently flattens structured content to a plain string
`context_prune.py:141, 152`
Both trim branches do `{**msg, "content": <str>}` unconditionally, even
when the original `content` was a list of blocks (`[{"type": "text", ...},
{"type": "image_url", ...}]`). A tool result that mixes text and an image
part loses the image permanently the first time it ages past the protected
window — not "trimmed," just gone, with no signal that anything
non-text was there.
## Cleanup — lower severity, no behavior change
| # | Location | Issue |
|---|---|---|
| 12 | `context_prune.py:73-80` | `_turn_of()` has zero call sites anywhere in the repo — leftover from an earlier design. |
| 13 | `dispatcher.py:1573-1585` | `_previous_context`'s content-block extraction duplicates `context_prune.extract_text()` almost verbatim, despite `dispatcher.py:58` already importing from that module (`prune_context` only). A future change to one won't propagate to the other. |
| 14 | `dispatcher.py:420-428` vs. `468-481` | `classify()`'s `if not context: user_content = task else: user_content = _classifier_user_content(...)` duplicates a check `_classifier_user_content` already makes internally (`if not context: return task`, line 477). The outer branch can be deleted; `classify()` can call `_classifier_user_content` unconditionally. |
| 15 | `context_prune.py:27-31`, `config.py:293-299`, `config.yaml:175-177` | Pinch's four defaults and `CHARS_PER_TOKEN` are each declared independently in two or three places. `context_prune.py`'s own comment admits `CHARS_PER_TOKEN` "mirrors dispatcher.CHARS_PER_TOKEN" rather than importing it. `dispatcher.py` always passes `cfg.pinch.*` explicitly, so the module-level defaults in `context_prune.py` are dead in production. |
## Take for next time
The port kept the right invariant (user/assistant/system messages are never
touched, only tool results) but re-derived the surrounding plumbing instead
of reusing what the codebase already had for it — `_previous_context` is a
second, slightly-different copy of `extract_text`'s block-parsing logic, and
it re-introduces exactly the bug `_last_user_text` next to it was written to
avoid (`_last_user_text` correctly scans for the *nearest* user message
rather than assuming position; `_previous_context` assumes `messages[-2]` is
meaningful). Both new-feature bugs that actually change routing behavior
today (#1, #2) are about that same unguarded assumption: agent traffic
doesn't end tidily on a fresh user turn, and this codebase already knows
that everywhere else it touches messages.

View File

@@ -281,6 +281,48 @@ class FreshnessConfig(StrictModel):
exclude_deprecated: bool exclude_deprecated: bool
class PinchConfig(StrictModel):
"""Optional relevance-based context pruning (Port of llmrouter's pinch).
Prunes the provider-bound conversation — not the classifier input — when it
exceeds ``budget_tokens``, so a long agent session ships fewer prompt tokens
upstream. User/assistant/system messages are always kept; only tool results
are summarized or dropped (they carry the bulk of a long session's tokens).
"""
enabled: bool = False
budget_tokens: int = 50000
# How many recent user turns (plus their assistant replies and tool results)
# are protected from pruning.
keep_last_turns: int = 4
# Tool results longer than this many characters are summarized in place.
max_summarize_chars: int = 4000
@field_validator("budget_tokens")
@classmethod
def budget_positive(cls, v: int) -> int:
if v <= 0:
raise ValueError("pinch.budget_tokens must be > 0")
return v
@field_validator("keep_last_turns")
@classmethod
def turns_positive(cls, v: int) -> int:
if v <= 0:
raise ValueError("pinch.keep_last_turns must be > 0")
return v
@field_validator("max_summarize_chars")
@classmethod
def summarize_chars_valid(cls, v: int) -> int:
if v < 3000:
raise ValueError(
"pinch.max_summarize_chars must be >= 3000 (below this, "
"summarization grows the message)"
)
return v
class DatabaseConfig(StrictModel): class DatabaseConfig(StrictModel):
path: str path: str
@@ -295,6 +337,11 @@ class ClassifierConfig(StrictModel):
model: str model: str
# Ceiling on the text handed to the classifier. 0 disables clamping. # Ceiling on the text handed to the classifier. 0 disables clamping.
max_input_chars: int = 8000 max_input_chars: int = 8000
# When a preceding turn is available as context, frame the classifier
# input as llmrouter does — "Context: <prev>\n---\nMessage: <task>" — so a
# short follow-up ("Yes", "Try now?") can inherit the complexity of the
# turn it continues instead of being classified in isolation as trivial.
context_framing: bool = True
timeout_seconds: int timeout_seconds: int
temperature: float = 0.0 temperature: float = 0.0
max_output_tokens: int = 1024 max_output_tokens: int = 1024
@@ -342,6 +389,7 @@ class RouterConfig(StrictModel):
local_vision: LocalVisionConfig = LocalVisionConfig() local_vision: LocalVisionConfig = LocalVisionConfig()
escalation: EscalationConfig escalation: EscalationConfig
iteration: IterationConfig = IterationConfig() iteration: IterationConfig = IterationConfig()
pinch: PinchConfig = PinchConfig()
freshness: FreshnessConfig freshness: FreshnessConfig
database: DatabaseConfig database: DatabaseConfig
classifier: ClassifierConfig classifier: ClassifierConfig

View File

@@ -157,6 +157,25 @@ iteration:
# IS a quality loss. Batch work does not care. # IS a quality loss. Batch work does not care.
max_attempts_interactive: 1 max_attempts_interactive: 1
pinch:
# Relevance-based context pruning (ported from the MIT-licensed llmrouter's
# "pinch"). This is an OPTIONAL, pre-dispatch stage: when a conversation
# exceeds budget_tokens, the provider-bound messages are trimmed BEFORE any
# paid token is sent upstream. User/assistant/system messages are always
# kept verbatim; only old TOOL RESULTS are shortened or dropped, because
# they carry the bulk of a long agent session's tokens and are least needed
# in full by the time the next turn is answered.
#
# It does NOT touch the classifier's input, and it defaults off — enable it
# only if long sessions are shipping more prompt tokens than you want to pay
# for. Tool results can only be dropped safely because tool outputs are
# idempotent enough for a placeholder; a wrong guess here loses context, so
# start conservative (large budget, small reduction).
enabled: false
budget_tokens: 50000
keep_last_turns: 4
max_summarize_chars: 4000
routing: routing:
# Access gating is prose-only in the NeuralWatt catalog ("Private preview # Access gating is prose-only in the NeuralWatt catalog ("Private preview
# (grant-gated)", "(Canary)"), so the poller parses it into access_level and # (grant-gated)", "(Canary)"), so the poller parses it into access_level and
@@ -349,6 +368,12 @@ classifier:
# never depends on what the classifier saw. 0 disables clamping. # never depends on what the classifier saw. 0 disables clamping.
max_input_chars: 8000 max_input_chars: 8000
# When the chat path supplies the previous turn as context (see the pinch /
# context notes), frame the classifier input as "Context: <prev> / Message:
# <current>" so a short follow-up inherits the prior turn's complexity
# instead of being classified in isolation as trivial.
context_framing: true
fallback_tier: 2 fallback_tier: 2
fallback_category: general_chat fallback_category: general_chat
response_format: "json" # ask Ollama to constrain output to valid JSON response_format: "json" # ask Ollama to constrain output to valid JSON

273
context_prune.py Normal file
View File

@@ -0,0 +1,273 @@
"""Relevance-based context pruning for the provider-bound conversation.
Ported from the MIT-licensed alexrudloff/llmrouter "pinch" module, reduced to
the pure, injectable decision core that 6krrt can test offline. Where llmrouter
embeds every candidate message and scores cosine relevance, this module keeps
the same SAFE invariants without requiring an embedding model on the request
path:
- user / assistant / system messages are ALWAYS kept verbatim,
- only TOOL RESULTS older than the protected window are trimmed or
summarized (never removed),
- message order and the tool/assistant pairing are preserved so the result
still parses as a valid conversation.
The tailoring is deliberate. Tool results are where a long agent session's
tokens actually live, they are the least likely to be needed in full by the
time a later turn is answered, and replacing one with a short placeholder is
reversible at the semantic level -- a wrong guess costs context, but it never
breaks the request. Trimming a user or assistant message, by contrast, can
change what the model is being asked, so those are never touched.
This module is pure: it takes messages and limits and returns pruned messages.
`dispatcher.py` owns reading the config and deciding when to call it.
"""
from __future__ import annotations
from config import PinchConfig
# Mirrors dispatcher.CHARS_PER_TOKEN.
CHARS_PER_TOKEN = 3
def estimate_tokens(text: str | None) -> int:
"""Crude characters-per-token estimate, consistent with the dispatcher."""
return len(text) // CHARS_PER_TOKEN if text else 0
def extract_text(message: dict) -> str:
"""Best-effort text of a message, whether content is str or content blocks.
``image_url`` blocks contribute their ``url`` so image-bearing messages are
not silently undercounted as empty. A ``tool_result`` whose ``content`` is
itself a list of blocks is flattened recursively the same way.
"""
content = message.get("content")
if isinstance(content, str):
return content
if isinstance(content, list):
out: list[str] = []
for part in content:
if not isinstance(part, dict):
continue
ptype = part.get("type")
if ptype == "text":
text = part.get("text")
if isinstance(text, str):
out.append(text)
elif ptype == "tool_result":
inner = part.get("content")
if isinstance(inner, list):
out.append(extract_text({"content": inner}))
elif isinstance(inner, str):
out.append(inner)
elif ptype == "image_url":
image_url = part.get("image_url") or {}
url = image_url.get("url") if isinstance(image_url, dict) else None
if isinstance(url, str):
out.append(url)
return " ".join(out)
return str(content) if content else ""
def _text_only(message: dict) -> str:
"""The prose text of a message, excluding image payloads.
Unlike :func:`extract_text` -- which includes ``image_url`` urls so they
count toward the size estimate that triggers pruning -- this returns only
the text blocks' content. A trimmed replacement must be built from prose,
never from raw base64, or an arbitrary byte slice of an image ends up
sitting in a ``text`` field. ``extract_text`` still sizes the message and
computes savings; this only shapes the replacement.
"""
content = message.get("content")
if isinstance(content, str):
return content
if isinstance(content, list):
out: list[str] = []
for part in content:
if not isinstance(part, dict):
continue
ptype = part.get("type")
if ptype == "text":
text = part.get("text")
if isinstance(text, str):
out.append(text)
elif ptype == "tool_result":
inner = part.get("content")
if isinstance(inner, list):
out.append(_text_only({"content": inner}))
elif isinstance(inner, str):
out.append(inner)
return " ".join(out)
return str(content) if content else ""
def _tool_name(message: dict) -> str | None:
"""A tool message's name, under either the OpenAI or Anthropic spelling."""
name = message.get("name") or message.get("tool_name") or message.get("function")
if isinstance(name, dict):
name = name.get("name")
return name if isinstance(name, str) else None
def _first_user_turn_indexes(messages: list[dict]) -> list[int]:
"""Indexes of messages that start a new user turn.
A user message starts a turn; the assistant reply and any tool results that
follow belong to that turn until the next user message.
"""
return [i for i, m in enumerate(messages) if m.get("role") == "user"]
def _with_text(message: dict, new_text: str) -> dict:
"""Return ``message`` with its tool-result content trimmed to ``new_text``.
When the content is a list of blocks, the text blocks become ``new_text``
and any ``image_url`` blocks have their payload shrunk to a stub. Trimming
must reduce the *actual* bytes shipped upstream, and an image is usually
the dominant contributor, so leaving it full-size while claiming a token
saving would be a lie. The ``image_url`` block type is preserved (so the
model still knows an image was present) but the multi-MB base64 url is
replaced by a short marker. When there is no text block, one carrying
``new_text`` is appended; a message with only an image collapses to just
the stub.
"""
content = message.get("content")
if isinstance(content, list):
blocks: list[dict] = []
replaced_text = False
for part in content:
if not isinstance(part, dict):
blocks.append(part)
continue
ptype = part.get("type")
if ptype == "text":
if not replaced_text:
blocks.append({**part, "text": new_text})
replaced_text = True
# Any subsequent text block is folded into ``new_text``.
elif ptype == "image_url":
# Shrink the image payload to a stub; keep the block type.
blocks.append({**part, "image_url": {"url": "[image omitted]"}})
else:
blocks.append(part)
if not replaced_text:
blocks.append({"type": "text", "text": new_text})
if blocks:
return {**message, "content": blocks}
return {**message, "content": new_text}
def prune_context(
messages: list[dict],
budget_tokens: int = PinchConfig.model_fields["budget_tokens"].default,
keep_last_turns: int = PinchConfig.model_fields["keep_last_turns"].default,
max_summarize_chars: int = PinchConfig.model_fields["max_summarize_chars"].default,
) -> tuple[list[dict], dict]:
"""Trim old tool results once a conversation exceeds ``budget_tokens``.
Returns ``(pruned_messages, stats)``. Keeps every user/assistant/system
message verbatim and never removes a message: tool results are replaced
*in place* with a short summary or placeholder, so ``len(pruned)`` always
equals ``len(messages)``. Order and role pairing are preserved, so the
result is a valid conversation with the same shape.
Only runs (and only mutates anything) when the estimate actually exceeds
the budget; otherwise the original list is returned untouched.
"""
orig_tokens = sum(estimate_tokens(extract_text(m)) for m in messages)
if orig_tokens <= budget_tokens:
return messages, {
"pruned": False,
"original_tokens": orig_tokens,
"final_tokens": orig_tokens,
"tokens_saved": 0,
}
user_indexes = _first_user_turn_indexes(messages)
if not user_indexes:
# No user turn at all: treat the conversation as a single ongoing turn
# and protect the trailing ``keep_last_turns`` tool results (the ones
# the next turn needs). ``keep_last_turns`` bounds how many are
# protected; a tiny conversation protects all of them.
tool_indexes = [i for i, m in enumerate(messages) if m.get("role") == "tool"]
protected_from = (
tool_indexes[max(len(tool_indexes) - keep_last_turns, 0)]
if tool_indexes
else len(messages)
)
else:
num_protected_turns = min(keep_last_turns, len(user_indexes))
if num_protected_turns <= 0:
# Nothing to protect: everything before is a candidate for trimming.
protected_from = len(messages)
else:
protected_from = user_indexes[len(user_indexes) - num_protected_turns]
pruned: list[dict] = []
summarized = 0
for i, msg in enumerate(messages):
role = msg.get("role")
if role in ("user", "assistant", "system") or i >= protected_from:
pruned.append(msg)
continue
# Only tool results are candidates here.
if role != "tool":
pruned.append(msg)
continue
content = msg.get("content")
if isinstance(content, str):
text = content
prose = content
else:
# `text` sizes the message (includes image bytes, so an image-heavy
# result still triggers pruning and still affects the length guard);
# `prose` shapes the replacement (never raw base64).
text = extract_text(msg)
prose = _text_only(msg)
head_len = 1500
tail_len = 1500
replaced = False
if len(text) > max_summarize_chars:
# Long tool result: keep head + tail so the shape survives, but
# only when the elision actually saves characters. A result that is
# not meaningfully longer than head+tail+marker would only grow (or
# produce a negative trim), so fall through to the placeholder path.
head = prose[:head_len]
tail = prose[-tail_len:]
trimmed = len(prose) - head_len - tail_len
marker = f"\n\n[{trimmed:,} chars trimmed...]\n\n" if trimmed > 0 else ""
elided = f"{head}{marker}{tail}"
if trimmed > 0 and len(elided) < len(prose):
summarized += 1
pruned.append(_with_text(msg, elided))
replaced = True
if not replaced:
# Short result (or a long one that cannot be elided to save space):
# replace with a short placeholder, but only when it is strictly
# shorter than the *combined* size — an image-only result may have
# empty prose yet large real size, and shrinking it still saves
# tokens (the image payload is stubbed by _with_text).
name = _tool_name(msg) or "tool"
placeholder = f"[{name}: result omitted]"
if len(placeholder) < len(text):
summarized += 1
pruned.append(_with_text(msg, placeholder))
else:
pruned.append(msg)
final_tokens = sum(estimate_tokens(extract_text(m)) for m in pruned)
return pruned, {
"pruned": True,
"original_tokens": orig_tokens,
"final_tokens": final_tokens,
# Savings computed from what the payload actually shrunk by (the whole
# list before vs after), so images + prose that really got trimmed are
# the only thing counted — never a per-message estimate that could
# drift from reality as #5/#11 interacted.
"tokens_saved": max(orig_tokens - final_tokens, 0),
"summarized": summarized,
}

View File

@@ -55,6 +55,7 @@ from pydantic import BaseModel, Field
import logs import logs
from capabilities import detect_capabilities, iter_image_url_values from capabilities import detect_capabilities, iter_image_url_values
from config import RouterConfig, load_config from config import RouterConfig, load_config
from context_prune import extract_text, prune_context
from routing import ( from routing import (
BATCH, BATCH,
INTERACTIVE, INTERACTIVE,
@@ -148,7 +149,16 @@ logs.configure(cfg.logging.level)
class TaskRequest(BaseModel): class TaskRequest(BaseModel):
task: str = Field(..., description="The task to route.") task: str = Field(..., description="The task to route.")
context: Optional[str] = Field( context: Optional[str] = Field(
None, description="Assembled context (docs/code) to send with the task." None,
description=(
"Assembled context (docs/code) to send with the task. "
"On dispatch_endpoint the chat path reuses this field to carry "
"the prior conversation turn (framing) passed to the classifier; "
"this dual use is safe because dispatch_endpoint only reads it "
"for the classify() branch, while the chat path rewrites the "
"Classification when task_category+task_tier+required_context "
"are all provided (the override branch never reads req.context)."
),
) )
latency_tolerance: Optional[Literal["interactive", "batch"]] = Field( latency_tolerance: Optional[Literal["interactive", "batch"]] = Field(
None, None,
@@ -416,9 +426,12 @@ def classify(task: str, context: Optional[str]) -> Classification:
f"Allowed values for task_category (use EXACTLY one of these strings):\n" f"Allowed values for task_category (use EXACTLY one of these strings):\n"
+ "\n".join(f" - {c}" for c in categories) + "\n".join(f" - {c}" for c in categories)
) )
user_content = task if not context else f"{task}\n\n--- context ---\n{context}" # _classifier_user_content handles None context internally (returns task),
# so the outer if/else is redundant — removed for clarity.
user_content = _classifier_user_content(
task, context, cfg.classifier.context_framing
)
user_content = clamp_for_classifier(user_content, cfg.classifier.max_input_chars) user_content = clamp_for_classifier(user_content, cfg.classifier.max_input_chars)
client = _classifier_client() client = _classifier_client()
started = time.perf_counter() started = time.perf_counter()
try: try:
@@ -457,6 +470,35 @@ def classify(task: str, context: Optional[str]) -> Classification:
return resp return resp
def _classifier_user_content(
task: str, context: Optional[str], framing: bool
) -> str:
"""Assemble the classifier's user message (pure, for testability).
``framing`` selects llmrouter's "Context: <prev> / Message: <current>"
layout so a short follow-up inherits its prior turn's complexity;
otherwise the legacy "task / --- context --- / context" layout is used.
In both layouts the framing instruction is appended when context is
present, telling the classifier to treat follow-ups by the prior
complexity.
"""
if not context:
return task
if framing:
user_content = (
f"Context: {context}\n---\nMessage: {task}"
)
else:
user_content = f"{task}\n\n--- context ---\n{context}"
user_content += (
"\n\nA short follow-up message ('Yes', 'Try now?', 'Go ahead') "
"continues the prior turn, so classify by the CONTEXT's "
"complexity, not in isolation."
)
return user_content
def _classify_once(client: OpenAI, system_prompt: str, user_content: str) -> Classification: def _classify_once(client: OpenAI, system_prompt: str, user_content: str) -> Classification:
"""One classifier round-trip. Raises on anything unusable.""" """One classifier round-trip. Raises on anything unusable."""
categories = cfg.proficiency.categories categories = cfg.proficiency.categories
@@ -1533,6 +1575,28 @@ def _last_user_text(messages: list[dict]) -> str:
return "" return ""
def _previous_context(messages: list[dict]) -> str:
"""The preceding assistant turn of the last user message.
Walks backwards from the message before the last user turn, returning the
text content of the nearest ``assistant`` message. Skips ``system``,
``user``, and ``tool`` roles so that tool output or the system prompt
never contaminates the framing signal. Uses ``context_prune.extract_text``
for the actual content-block parsing (same logic as the former inlined
block parser, but deduplicated). Returns the concatenated text of the
first assistant message found, truncated to 200 characters, or ``""``
when no such message exists.
"""
for i in range(len(messages) - 2, -1, -1):
msg = messages[i]
if msg.get("role") != "assistant":
continue
text = extract_text(msg)
if text:
return text[:200]
return ""
def _count_images(messages: list[dict]) -> int: def _count_images(messages: list[dict]) -> int:
"""The number of image_url parts across the whole conversation. """The number of image_url parts across the whole conversation.
@@ -1858,10 +1922,15 @@ def chat_completions(body: dict[str, Any], background: BackgroundTasks):
if wants_routing: if wants_routing:
latency = BATCH if requested == ROUTER_MODEL_BATCH else INTERACTIVE latency = BATCH if requested == ROUTER_MODEL_BATCH else INTERACTIVE
tools_present = caps.tools_present tools_present = caps.tools_present
# The turn before the last: a short follow-up inherits its complexity,
# so the classifier sees "Context: <prior>\n---\nMessage: <current>"
# instead of judging the follow-up alone.
prev_context = _previous_context(messages)
classify_started = time.perf_counter() classify_started = time.perf_counter()
decision = route( decision = route(
TaskRequest( TaskRequest(
task=_last_user_text(messages), task=_last_user_text(messages),
context=prev_context,
latency_tolerance=latency, latency_tolerance=latency,
tools_present=tools_present, tools_present=tools_present,
has_images=caps.has_images, has_images=caps.has_images,
@@ -1876,12 +1945,39 @@ def chat_completions(body: dict[str, Any], background: BackgroundTasks):
classified_src = decision.classification.source classified_src = decision.classification.source
classifier_ms = _ms(classify_started) classifier_ms = _ms(classify_started)
ctx_src = "classifier" ctx_src = "classifier"
measured = estimate_prompt_tokens(messages, tools=body.get("tools")) # When pinch is enabled, prune ONCE before the measured-context
# decision, so the window/tier/cost choice sees the size that will
# actually ship upstream rather than the raw conversation. The pruned
# list is reused at dispatch time (send_messages below), never pruned
# twice. When pinch is off this is a byte-for-byte no-op (full list,
# as today). Pinch trims only tool results, never user/assistant/system
# messages, so the classification turn above (which used
# `_previous_context` from the full messages) is undisturbed.
if cfg.pinch.enabled:
send_messages, pinch_stats = prune_context(
list(messages),
budget_tokens=cfg.pinch.budget_tokens,
keep_last_turns=cfg.pinch.keep_last_turns,
max_summarize_chars=cfg.pinch.max_summarize_chars,
)
logs.debug(
"pinch",
pruned=pinch_stats["pruned"],
saved=pinch_stats["tokens_saved"],
orig=pinch_stats["original_tokens"],
final=pinch_stats["final_tokens"],
)
else:
send_messages = messages
measured = estimate_prompt_tokens(send_messages, tools=body.get("tools"))
if measured > decision.classification.required_context_tokens: if measured > decision.classification.required_context_tokens:
ctx_src = "measured" ctx_src = "measured"
decision = route( decision = route(
TaskRequest( TaskRequest(
task=_last_user_text(messages), task=_last_user_text(messages),
# context omitted: the override branch (task_category +
# task_tier + required_context_tokens all provided)
# never reads req.context — it skips classify().
latency_tolerance=latency, latency_tolerance=latency,
tools_present=tools_present, tools_present=tools_present,
has_images=caps.has_images, has_images=caps.has_images,
@@ -2009,7 +2105,31 @@ def chat_completions(body: dict[str, Any], background: BackgroundTasks):
# top (needed by the passthrough/local-vision branches), only the working # top (needed by the passthrough/local-vision branches), only the working
# directory is derived here because only the observation path uses it. # directory is derived here because only the observation path uses it.
session_dir = session_directory(messages) session_dir = session_directory(messages)
upstream_body = {**body, "model": target} if not wants_routing:
# The routed path pruned send_messages once, before its measured-context
# decision, and the pruned list is reused here. A passthrough request
# never went through that path, so it prunes now — same gating as
# before (only when pinch is enabled).
send_messages = messages
if cfg.pinch.enabled:
# Optional relevance-based context pruning: when the conversation
# exceeds budget_tokens, trim old tool results BEFORE any paid token
# is sent upstream. User/assistant/system messages are always kept;
# only the classifier input is exempt (it already clamps to head+tail).
send_messages, pinch_stats = prune_context(
list(messages),
budget_tokens=cfg.pinch.budget_tokens,
keep_last_turns=cfg.pinch.keep_last_turns,
max_summarize_chars=cfg.pinch.max_summarize_chars,
)
logs.debug(
"pinch",
pruned=pinch_stats["pruned"],
saved=pinch_stats["tokens_saved"],
orig=pinch_stats["original_tokens"],
final=pinch_stats["final_tokens"],
)
upstream_body = {**body, "model": target, "messages": list(send_messages)}
streaming = bool(body.get("stream")) streaming = bool(body.get("stream"))
if streaming: if streaming:
# Without this the final chunk carries no usage and the observation # Without this the final chunk carries no usage and the observation
@@ -2054,7 +2174,10 @@ def chat_completions(body: dict[str, Any], background: BackgroundTasks):
retry_trail: list[str] = [] retry_trail: list[str] = []
while True: while True:
attempt_body = {**body, "model": current_model} attempt_body = {
**body, "model": current_model,
"messages": list(send_messages),
}
if current_max_tokens is not None: if current_max_tokens is not None:
attempt_body["max_tokens"] = current_max_tokens attempt_body["max_tokens"] = current_max_tokens
upstream_started = time.perf_counter() upstream_started = time.perf_counter()

View File

@@ -14,7 +14,11 @@ With the clamp, the same prompts classify correctly in ~2.2s.
from __future__ import annotations from __future__ import annotations
from dispatcher import clamp_for_classifier from dispatcher import (
_classifier_user_content,
_previous_context,
clamp_for_classifier,
)
def test_short_input_is_untouched(): def test_short_input_is_untouched():
@@ -61,3 +65,133 @@ def test_zero_disables_clamping():
def test_negative_is_treated_as_disabled_not_as_a_crash(): def test_negative_is_treated_as_disabled_not_as_a_crash():
text = "w" * 100_000 text = "w" * 100_000
assert clamp_for_classifier(text, -1) == text assert clamp_for_classifier(text, -1) == text
# --- context framing (ported from llmrouter) -------------------------------
def test_framing_without_context_is_just_the_task():
assert _classifier_user_content("Refactor this", None, True) == "Refactor this"
assert _classifier_user_content("Refactor this", None, False) == "Refactor this"
def test_framing_puts_prior_turn_first_when_enabled():
out = _classifier_user_content("Try now?", "Design a distributed system", True)
assert out.startswith("Context: Design a distributed system")
assert "\n---\nMessage: Try now?\n\n" in out
# The framing instruction is appended.
assert ("A short follow-up message ('Yes', 'Try now?', 'Go ahead') "
"continues the prior turn") in out
def test_legacy_framing_puts_task_first():
out = _classifier_user_content("Try now?", "Design a distributed system", False)
assert out.startswith("Try now?")
assert "--- context ---" in out
assert "Design a distributed system" in out
def test_framing_input_still_feeds_the_clamp():
# The framed input is what clamp_for_classifier trims, so a huge prior turn
# is still bounded to head + tail and the elision stays visible.
context = "context blob " * 50_000
framed = _classifier_user_content("Try now?", context, True)
assert len(framed) > 8000
clamped = clamp_for_classifier(framed, 8000)
assert "Context:" in clamped
assert "Message: Try now?" in clamped
assert "elided" in clamped
# --- previous-message extraction ------------------------------------------
def test_previous_context_is_the_message_before_last_user():
messages = [
{"role": "user", "content": "first"},
{"role": "assistant", "content": "partly done"},
{"role": "user", "content": "Try now?"},
]
assert _previous_context(messages) == "partly done"
def test_previous_context_skips_user_messages():
# Only the message before the LAST user is returned; earlier user messages
# are skipped, matching llmrouter ("short follow-up inherits from context").
messages = [
{"role": "user", "content": "earlier"},
{"role": "user", "content": "current"},
]
assert _previous_context(messages) == ""
def test_previous_context_empty_when_no_prior_message():
assert _previous_context([{"role": "user", "content": "only"}]) == ""
assert _previous_context([]) == ""
def test_previous_context_handles_content_blocks():
messages = [
{
"role": "assistant",
"content": [
{"type": "text", "text": "head"},
{"type": "tool_result", "content": "output"},
],
},
{"role": "user", "content": "Go"},
]
ctx = _previous_context(messages)
assert "head" in ctx
assert "output" in ctx
def test_previous_context_excludes_tool_role_messages():
# Standard tool loop: assistant(tool_call) → tool(result) → user("ok fix it")
# Should not include raw tool output as context.
messages = [
{"role": "assistant", "content": "Let me check", "tool_calls": [{"id": "t1", "type": "function", "function": {"name": "read_file", "arguments": "{}"}}]},
{"role": "tool", "tool_call_id": "t1", "content": "/home/alee/file.py\nline 1\nline 2\nline 3"},
{"role": "user", "content": "ok fix it"},
]
assert _previous_context(messages) == "Let me check"
assert "/home/alee/file.py" not in _previous_context(messages)
def test_previous_context_excludes_system_prompt():
# [system, user] on the first turn — should never return a system-prompt fragment.
messages = [
{"role": "system", "content": "You are a helpful coding assistant with tool access. You have access to read_file, write_file, and run_tests."},
{"role": "user", "content": "hello"},
]
assert _previous_context(messages) == ""
def test_previous_context_is_truncated():
messages = [
{"role": "assistant", "content": "z" * 10_000},
{"role": "user", "content": "Go"},
]
assert len(_previous_context(messages)) <= 200
def test_previous_context_returns_assistant_text_when_it_exists():
# [system, assistant("the real prior turn"), user("Yes")]
# Should return the assistant turn text, truncated to 200.
messages = [
{"role": "system", "content": "You are helpful."},
{"role": "assistant", "content": "I think you should refactor the auth module to use middleware instead of decorators."},
{"role": "user", "content": "Yes"},
]
ctx = _previous_context(messages)
assert "refactor" in ctx
assert str(ctx) == "I think you should refactor the auth module to use middleware instead of decorators."[:200]
# --- classify() caller invariant --------------------------------------------
def test_classify_user_content_invariant_when_no_context():
# Regression for the removed if/else in classify(): _classifier_user_content
# already handles None context (returns task), so removing the outer guard
# must preserve user_content == task.
assert _classifier_user_content("do something", None, True) == "do something"
assert _classifier_user_content("do something", None, False) == "do something"

View File

@@ -230,3 +230,61 @@ def test_nonpositive_local_timeout_is_rejected(raw):
cfg["local_vision"]["timeout_seconds"] = 0 cfg["local_vision"]["timeout_seconds"] = 0
with pytest.raises(ValueError, match="timeout_seconds"): with pytest.raises(ValueError, match="timeout_seconds"):
RouterConfig(**cfg) RouterConfig(**cfg)
def test_the_shipped_config_loads_the_pinch_section(raw):
loaded = RouterConfig(**raw)
assert loaded.pinch.enabled is False
assert loaded.pinch.budget_tokens > 0
assert loaded.pinch.keep_last_turns > 0
def test_pinch_defaults_when_absent(raw):
cfg = copy.deepcopy(raw)
cfg.pop("pinch")
loaded = RouterConfig(**cfg)
assert loaded.pinch.enabled is False
assert loaded.pinch.budget_tokens == 50000
assert loaded.pinch.keep_last_turns == 4
def test_nonpositive_pinch_budget_is_rejected(raw):
cfg = copy.deepcopy(raw)
cfg["pinch"]["budget_tokens"] = 0
with pytest.raises(ValueError, match="budget_tokens"):
RouterConfig(**cfg)
def test_nonpositive_pinch_keep_last_turns_is_rejected(raw):
cfg = copy.deepcopy(raw)
cfg["pinch"]["keep_last_turns"] = 0
with pytest.raises(ValueError, match="keep_last_turns"):
RouterConfig(**cfg)
def test_classifier_context_framing_loads(raw):
# Defaults on (see config.yaml); the legacy layout is opt-out.
assert RouterConfig(**raw).classifier.context_framing is True
cfg = copy.deepcopy(raw)
cfg["classifier"]["context_framing"] = False
assert RouterConfig(**cfg).classifier.context_framing is False
def test_nonminimum_pinch_max_summarize_chars_is_rejected(raw):
cfg = copy.deepcopy(raw)
cfg["pinch"]["max_summarize_chars"] = 100
with pytest.raises(ValueError, match="max_summarize_chars.*>=.*3000"):
RouterConfig(**cfg)
def test_pinch_max_summarize_chars_at_valid_min_loads(raw):
cfg = copy.deepcopy(raw)
cfg["pinch"]["max_summarize_chars"] = 3000
loaded = RouterConfig(**cfg)
assert loaded.pinch.max_summarize_chars == 3000
def test_pinch_max_summarize_chars_accepts_default(raw):
loaded = RouterConfig(**raw)
assert loaded.pinch.max_summarize_chars == 4000

385
tests/test_context_prune.py Normal file
View File

@@ -0,0 +1,385 @@
"""Pure relevance-based context pruning.
Ported from the MIT-licensed llmrouter "pinch" module, reduced to the safe
invariants 6krrt requires: user/assistant/system messages are always kept
verbatim, only old tool results are trimmed, and message order / role pairing
survives so the result still parses as a conversation.
"""
from __future__ import annotations
from context_prune import estimate_tokens, extract_text, prune_context
from dispatcher import estimate_prompt_tokens
def _user(text: str) -> dict:
return {"role": "user", "content": text}
def _assistant(text: str) -> dict:
return {"role": "assistant", "content": text}
def _tool(name: str, content: str) -> dict:
return {"role": "tool", "name": name, "content": content}
def test_short_conversation_is_untouched():
messages = [_user("hi"), _assistant("hello")]
out, stats = prune_context(messages, budget_tokens=50000)
assert out == messages
assert stats["pruned"] is False
def test_user_and_assistant_messages_always_kept():
messages = [
_user("big user " * 5000),
_assistant("big assistant " * 5000),
_user("tail"),
]
out, stats = prune_context(messages, budget_tokens=100)
assert stats["pruned"] is True
assert len(out) == len(messages)
assert out[0]["role"] == "user"
assert out[1]["role"] == "assistant"
assert out[2]["role"] == "user"
# All three survived verbatim; only their absolute content is untouched.
assert out[0]["content"] == messages[0]["content"]
assert out[2]["content"] == "tail"
def test_old_tool_results_beyond_window_are_trimmed():
messages = [
_tool("read", "huge tool output " * 2000), # old, should be trimmed
_assistant("step one"),
_user("continue"),
_assistant("ok"),
]
out, stats = prune_context(messages, budget_tokens=100, keep_last_turns=1)
assert stats["pruned"] is True
assert stats["tokens_saved"] > 0
# The old tool result got shortened (head + tail + elision marker).
tool_out = out[0]
assert tool_out["role"] == "tool"
assert "chars trimmed" in tool_out["content"]
assert out[1]["content"] == "step one"
assert out[2]["content"] == "continue"
assert out[3]["content"] == "ok"
def test_recent_turn_is_protected_from_pruning():
# keep_last_turns=1 protects the last user turn AND its tool results.
messages = [
_user("first"),
_tool("search", "old result " * 3000),
_assistant("first answer"),
_user("second"),
_tool("search", "recent result " * 3000),
_assistant("second answer"),
]
out, stats = prune_context(messages, budget_tokens=100, keep_last_turns=1)
assert stats["pruned"] is True
# The recent tool result (after the 2nd user) survives verbatim.
recent_tool = [m for m in out if m.get("role") == "tool"][-1]
assert recent_tool["content"] == messages[4]["content"]
# Old tool result got trimmed.
old_tool = [m for m in out if m.get("role") == "tool"][0]
assert "chars trimmed" in old_tool["content"]
def test_tool_result_role_pairing_preserved():
# The tool message keeps its name and position, so the API still parses.
messages = [
_user("q"),
_tool("search", "x" * 6000),
_assistant("a"),
]
out, stats = prune_context(messages, budget_tokens=50, keep_last_turns=0)
tool_out = [m for m in out if m.get("role") == "tool"][0]
assert tool_out["name"] == "search"
assert out[0]["role"] == "user"
assert out[1]["role"] == "tool"
assert out[2]["role"] == "assistant"
def test_short_tool_result_is_dropped_to_placeholder():
# A large user message pushes the conversation over budget, so the SHORT
# old tool result (under max_summarize_chars but longer than the marker) is
# dropped to a placeholder that is shorter than the result it replaces.
result_text = "result " * 40 # 280 chars > placeholder, < 4000 cap
messages = [
_user("big user " * 5000),
_tool("search", result_text),
_assistant("a"),
]
out, stats = prune_context(messages, budget_tokens=100, keep_last_turns=0)
assert stats["pruned"] is True
tool_out = [m for m in out if m.get("role") == "tool"][0]
assert "[search: result omitted]" in tool_out["content"]
assert len(tool_out["content"]) < len(result_text)
def test_tiny_tool_result_is_kept_when_placeholder_loses_tokens():
# A result shorter than the placeholder marker is kept verbatim: replacing
# it would ADD tokens, which is not the point of pruning.
messages = [
_user("big user " * 5000),
_tool("search", "hi"),
_assistant("a"),
]
out, stats = prune_context(messages, budget_tokens=100, keep_last_turns=0)
tool_out = [m for m in out if m.get("role") == "tool"][0]
assert tool_out["content"] == "hi"
def test_stats_report_tokens_saved():
messages = [_user("q"), _tool("search", "y" * 9000), _assistant("a")]
out, stats = prune_context(messages, budget_tokens=10, keep_last_turns=0)
assert stats["pruned"] is True
assert stats["original_tokens"] > stats["final_tokens"]
assert stats["tokens_saved"] > 0
assert stats["summarized"] >= 1
def test_extract_text_handles_content_blocks():
msg = {
"role": "tool",
"content": [
{"type": "text", "text": "hello"},
{"type": "tool_result", "content": "world"},
],
}
assert "hello" in extract_text(msg)
assert "world" in extract_text(msg)
def test_estimate_tokens_is_crude_but_consistent():
assert estimate_tokens("x" * 30) == 10
assert estimate_tokens("") == 0
assert estimate_tokens(None) == 0
def test_max_summarize_chars_below_3000_never_grows_or_goes_negative():
# Regression: a max_summarize_chars well below 3000 (here 100) with a
# moderately long tool result used to head/tail-trim into something LONGER
# than the input and a negative tokens_saved. It must never grow or negate.
content = "tool output " * 100 # 1200 chars
messages = [
_user("big user " * 5000),
_tool("read", content),
_assistant("a"),
]
out, stats = prune_context(
messages,
budget_tokens=100,
keep_last_turns=0,
max_summarize_chars=100,
)
tool_out = [m for m in out if m.get("role") == "tool"][0]
assert len(tool_out["content"]) < len(content) # strictly shorter
assert stats["tokens_saved"] >= 0
assert "[-" not in tool_out["content"] # no negative trim marker
def test_extract_text_counts_image_url_blocks():
# Regression: image_url blocks used to read as "", undercounting the
# message and letting pruning miss it. They must contribute their url.
msg = {
"content": [
{"type": "text", "text": "describe this"},
{"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}},
],
}
text = extract_text(msg)
assert text != ""
assert "describe this" in text
assert "AAAA" in text
def test_extract_text_flattens_nested_tool_result_content_list():
# Regression: a tool_result whose content is itself a list of blocks was
# str()'d as a repr instead of being flattened like a top-level list.
msg = {
"content": [
{
"type": "tool_result",
"content": [
{"type": "text", "text": "nested text"},
{"type": "image_url", "image_url": {"url": "http://x/img.png"}},
],
},
],
}
text = extract_text(msg)
assert "nested text" in text
assert "img.png" in text
def test_stats_use_summarized_and_length_is_always_preserved():
# Regression: the stat was named "dropped" but nothing is ever removed —
# content is replaced in place, so len(pruned) == len(messages) always.
messages = [
_user("q"),
_tool("search", "x" * 200),
_assistant("a"),
]
out, stats = prune_context(
messages,
budget_tokens=50,
keep_last_turns=0,
max_summarize_chars=5000,
)
assert "dropped" not in stats
assert "summarized" in stats
assert stats["summarized"] >= 1
assert len(out) == len(messages)
def test_zero_user_messages_protect_newest_tool_result():
# Regression: with no user message the recency guard collapsed to
# protected_from = len(messages), so even the newest tool result (the one
# the next turn needs) was trimmed. It must be protected.
messages = [
{"role": "system", "content": "sys"},
{"role": "assistant", "content": "call tool", "tool_calls": [{"id": "t1"}]},
{"role": "tool", "name": "read", "content": "x" * 20000},
]
out, stats = prune_context(messages, budget_tokens=100, keep_last_turns=4)
assert stats["pruned"] is True
tool_out = [m for m in out if m.get("role") == "tool"][0]
assert tool_out["content"] == messages[2]["content"]
def test_structured_tool_result_keeps_image_block_when_trimmed():
# Regression: trimming flattened structured content to a string, dropping
# the image_url part permanently. The non-text block must survive — kept
# as an image_url block (so the model knows an image was present), with
# its payload shrunk to a stub rather than shipped full-size.
messages = [
_user("big user " * 5000),
{
"role": "tool",
"name": "capture",
"content": [
{"type": "text", "text": "some result " * 300}, # > 3000 chars
{"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}},
],
},
_assistant("a"),
]
out, stats = prune_context(
messages,
budget_tokens=100,
keep_last_turns=0,
max_summarize_chars=100,
)
tool_out = [m for m in out if m.get("role") == "tool"][0]
# The content is still a structured list, not a flattened string.
assert isinstance(tool_out["content"], list)
text_blocks = [
b for b in tool_out["content"]
if isinstance(b, dict) and b.get("type") == "text"
]
assert text_blocks
# Non-empty prose and non-trivial image means it was shrunk (not kept full).
img_blocks = [
b for b in tool_out["content"]
if isinstance(b, dict) and b.get("type") == "image_url"
]
assert len(img_blocks) == 1
# The block survives (type preserved) but its payload is stubbed, so the
# actual shipped bytes drop rather than staying full-size (review #new).
assert img_blocks[0]["image_url"]["url"] == "[image omitted]"
assert img_blocks[0]["image_url"]["url"] != "data:image/png;base64,AAAA"
def test_large_image_tool_result_actually_shrinks_and_stats_are_honest():
# Regression (follow-up review): fixing image_url counting (old #5) and
# preserving non-text blocks (old #11) composed into a case where the text
# block was sliced from prose+base64 joined together and the image_url
# block was copied through UNCHANGED — so the dominant size was never
# reduced while tokens_saved reported a large "saving" that didn't happen,
# and a raw base64 fragment leaked into the text field. The payload must
# genuinely shrink, the image payload must be stubbed, base64 must not
# leak into the text block, and tokens_saved must equal the real shrink.
huge_b64 = "A" * 200000
messages = [
_user("describe this"),
{"role": "assistant", "content": None, "tool_calls": [{"id": "1"}]},
{
"role": "tool",
"name": "screenshot",
"content": [
{"type": "text", "text": "here is the screenshot"},
{"type": "image_url", "image_url": {"url": f"data:image/png;base64,{huge_b64}"}},
],
},
_user("now what"),
]
out, stats = prune_context(messages, budget_tokens=10, keep_last_turns=0)
assert stats["pruned"] is True
# The payload genuinely shrank — final below original.
assert stats["final_tokens"] < stats["original_tokens"]
# Savings equal the real reduction, not a per-message estimate.
assert stats["tokens_saved"] == (
stats["original_tokens"] - stats["final_tokens"]
)
tool_out = [m for m in out if m.get("role") == "tool"][0]
for part in tool_out["content"]:
if isinstance(part, dict) and part.get("type") == "image_url":
assert "AAAA" not in part["image_url"]["url"]
elif isinstance(part, dict) and part.get("type") == "text":
assert "base64" not in part.get("text", "")
def test_pruned_list_measures_smaller_than_the_raw_conversation():
"""Regression for review #4: routing must measure the PRUNED size.
chat_completions used to route on the raw conversation's token count and
only shrink the upstream payload later, so a long tool-heavy session could
be sent to a pricier large-context model based on its pre-pruned size even
though pinch would bring the outgoing request well under budget. This test
pins the invariant that enables the fix: after one prune, the router's own
``estimate_prompt_tokens`` sees a smaller count, so the window/tier/cost
decision reflects what actually ships upstream.
"""
messages = [
{"role": "user", "content": "inspect the repo"},
{"role": "assistant", "content": "reading files", "tool_calls": [{"id": "t1"}]},
{"role": "tool", "name": "read", "content": "file1: " + "x" * 4000},
{"role": "assistant", "content": "still going", "tool_calls": [{"id": "t2"}]},
{"role": "tool", "name": "read", "content": "file2: " + "y" * 4000},
{"role": "assistant", "content": "ok, fixing"},
{"role": "user", "content": "fix it"},
]
raw_tokens = estimate_prompt_tokens(messages)
pruned, stats = prune_context(
list(messages),
budget_tokens=100,
keep_last_turns=0,
max_summarize_chars=4000,
)
assert stats["pruned"] is True
pruned_tokens = estimate_prompt_tokens(pruned)
# The outgoing payload measures strictly smaller than the raw conversation,
# so routing on the pruned list picks the cheaper/smaller-window model.
assert pruned_tokens < raw_tokens
def test_pinch_disabled_leaves_the_measured_size_unchanged():
"""When pinch is off the routed decision measures the full conversation.
Mirrors the byte-for-byte no-op contract: with ``prune_context`` not run,
``estimate_prompt_tokens`` sees the raw messages exactly as before the
reordering fix.
"""
messages = [
{"role": "user", "content": "q"},
{"role": "assistant", "content": "a", "tool_calls": [{"id": "t1"}]},
{"role": "tool", "name": "read", "content": "z" * 4000},
{"role": "user", "content": "done"},
]
# unpinned: identical to calling estimate_prompt_tokens(messages) directly
assert estimate_prompt_tokens(messages) == estimate_prompt_tokens(list(messages))