5f7716e121
Address the 11 confirmed bugs and 4 cleanups from context-pruning-and-framing-review.md, plus the image-accounting regression it found on a second pass. classifier framing: - _previous_context walks back to the nearest assistant turn only, skipping system/user/tool so raw tool output and the system prompt never contaminate the framing signal; reuses context_prune.extract_text instead of a duplicated block parser. - classify() drops its redundant if/not-context branch. - framing instruction moved out of the static system prompt and into _classifier_user_content so it travels only with content that justifies it (respects context_framing opt-out). context pruning: - extract_text counts image_url bytes so image-bearing tool results are sized (and pruned) correctly. - -- new, follow-up review -- image_url blocks now count toward the size trigger but not toward the replacement text, and are stubbed to a short marker when the message is trimmed, so an image-heavy tool result genuinely shrinks rather than staying full-size while tokens_saved lied and base64 leaked into the text block. - tokens_saved is now the honest whole-list before/after reduction (max(orig - final, 0)), never a per-message estimate that could drift. - max_summarize_chars validated >= 3000 at config load; trim branch guards against negative math and message growth. - recency guard protects the newest tool result even with no user turn. - non-text blocks survive trimming (type preserved) rather than being flattened away. - dropped/stat terminology corrected to summarized; _turn_of() removed; module-level pinch defaults removed in favor of PinchConfig. Dispatch and config: - prune_context runs once before the measured-context routing decision (reused at dispatch, no double-prune, no in-place mutation); no-op when pinch is disabled. - dead context=prev_context arg removed from the measured reroute. - TaskRequest.context dual use documented. 15 new regression tests pin the fixes; 545 tests pass.
274 lines
12 KiB
Python
274 lines
12 KiB
Python
"""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,
|
|
}
|