209 lines
8.4 KiB
Python
209 lines
8.4 KiB
Python
"""Keeps the sub-agent's context inside a budget, cheaply first.
|
|
|
|
WHAT THIS IS NOT
|
|
|
|
The archived sub-agent stack did not compact. At a fraction of num_ctx it
|
|
forced the agent to produce a final answer, which is why its long jobs
|
|
ended abruptly and mid-thought. Compaction here is a genuine behavioural
|
|
change, not a port of that.
|
|
|
|
THE ORDER MATTERS, AND IT IS CHEAPEST-FIRST
|
|
|
|
1. Cap every tool result on the way IN. A single `ls -R /` or a failing
|
|
build's output can be megabytes, and once it is in the transcript it
|
|
is paid for on every subsequent turn, forever.
|
|
2. Drop old tool RESULTS, oldest first, leaving a one-line stub saying
|
|
what was dropped. The agent almost never needs the full output of a
|
|
command it ran twenty steps ago, but it very often needs to remember
|
|
that it ran it and what happened.
|
|
3. Only then summarise, with one model call.
|
|
|
|
Steps 1 and 2 are free. Step 3 costs a call AND invalidates the prefix
|
|
cache, which on a single-slot Ollama means the next turn re-prefills
|
|
everything - so it is deliberately the last resort rather than the
|
|
mechanism.
|
|
|
|
THE IMAGE RULE IS THE BIGGEST SAVING ONCE THE SCREEN IS USED
|
|
|
|
This model has vision, so a screenshot is a real image in the context,
|
|
not a filename. Keeping the last few costs more than everything else put
|
|
together. Exactly one is kept - the newest - and older ones become the
|
|
text the agent itself wrote about them, which is the part it was using
|
|
anyway.
|
|
|
|
WHAT IS NEVER TOUCHED
|
|
|
|
The system prompt, the first user message, and the task list. The first
|
|
two are the prefix the cache is built on; the third is the standard the
|
|
session is measured against, and a standard that can be summarised away
|
|
is not one.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Any
|
|
|
|
# Qwen3's tokeniser on dense English prose, the same figure budget.py
|
|
# uses in the api container. It is an estimate and is used as one: the
|
|
# thresholds below have slack for it being wrong by a fifth.
|
|
CHARS_PER_TOKEN = 3.2
|
|
|
|
# A single tool result, on the way in.
|
|
MAX_RESULT_CHARS = 4000
|
|
|
|
# Working-set target. Well under the 163840 the model is loaded with:
|
|
# the aim is a fast prefill on every step, not to fill the window. A
|
|
# 40k-token prompt re-prefilled on a handover costs perhaps forty
|
|
# seconds; a 160k one costs minutes.
|
|
TARGET_TOKENS = 24_000
|
|
CEILING_TOKENS = 40_000
|
|
|
|
# Recent tool results always survive step 2 verbatim.
|
|
KEEP_VERBATIM = 8
|
|
|
|
|
|
def estimate_tokens(messages: list[dict[str, Any]]) -> int:
|
|
total = 0
|
|
for message in messages:
|
|
for key in ("content", "thinking"):
|
|
value = message.get(key)
|
|
if isinstance(value, str):
|
|
total += int(len(value) / CHARS_PER_TOKEN)
|
|
if message.get("tool_calls"):
|
|
total += int(len(json.dumps(message["tool_calls"])) / CHARS_PER_TOKEN)
|
|
# Open WebUI's own estimator charges 1000 tokens for an image
|
|
# part and 4 per message; both are borrowed here rather than
|
|
# invented, so the two layers at least disagree consistently.
|
|
total += 1000 * len(message.get("images") or [])
|
|
total += 4
|
|
return total
|
|
|
|
|
|
def cap_result(text: str, limit: int = MAX_RESULT_CHARS) -> str:
|
|
"""Trim one tool result, and SAY that it was trimmed.
|
|
|
|
Silent truncation is the failure mode worth avoiding: an agent that
|
|
thinks it read all of a file behaves very differently from one that
|
|
knows it read the first 4000 characters of it. Same reasoning as
|
|
budget.cap() in the api container.
|
|
"""
|
|
if len(text) <= limit:
|
|
return text
|
|
kept = text[:limit].rsplit("\n", 1)[0] or text[:limit]
|
|
return (
|
|
f"{kept}\n\n[TRUNCATED after {len(kept):,} of {len(text):,} characters. "
|
|
"Re-run with a narrower command - grep, tail, or a specific path - if you "
|
|
"need the rest; do not assume what was cut.]"
|
|
)
|
|
|
|
|
|
def _stub_for(message: dict[str, Any]) -> str:
|
|
label = message.get("_label") or "a tool result"
|
|
size = len(message.get("content") or "")
|
|
return f"[dropped to save context: {label}, {size:,} characters]"
|
|
|
|
|
|
def compact(messages: list[dict[str, Any]], protect: int,
|
|
summarise) -> tuple[list[dict[str, Any]], list[str]]:
|
|
"""Bring `messages` under budget. Returns the new list and what was done.
|
|
|
|
`protect` is how many leading messages are untouchable - the system
|
|
prompt and the opening task. `summarise` is a callable taking a list
|
|
of messages and returning prose; it is only invoked if steps 1 and 2
|
|
were not enough, and it is the only thing here that costs a model call.
|
|
"""
|
|
notes: list[str] = []
|
|
if estimate_tokens(messages) <= TARGET_TOKENS:
|
|
return messages, notes
|
|
|
|
working = list(messages)
|
|
|
|
# --- images: keep exactly the newest -------------------------------
|
|
image_indices = [i for i, m in enumerate(working) if m.get("images")]
|
|
for index in image_indices[:-1]:
|
|
message = dict(working[index])
|
|
count = len(message.get("images") or [])
|
|
message.pop("images", None)
|
|
message["content"] = (
|
|
(message.get("content") or "").strip()
|
|
+ f"\n[screenshot dropped to save context - {count} image(s). "
|
|
"Take another with screenshot() if you need to look again.]"
|
|
).strip()
|
|
working[index] = message
|
|
if len(image_indices) > 1:
|
|
notes.append(f"dropped {len(image_indices) - 1} older screenshot(s)")
|
|
if estimate_tokens(working) <= TARGET_TOKENS:
|
|
return working, notes
|
|
|
|
# --- old tool results: stub them, oldest first ----------------------
|
|
tool_indices = [i for i, m in enumerate(working)
|
|
if m.get("role") == "tool" and i >= protect]
|
|
droppable = tool_indices[:-KEEP_VERBATIM] if len(tool_indices) > KEEP_VERBATIM else []
|
|
dropped = 0
|
|
for index in droppable:
|
|
if estimate_tokens(working) <= TARGET_TOKENS:
|
|
break
|
|
if len(working[index].get("content") or "") < 200:
|
|
continue # not worth the churn
|
|
message = dict(working[index])
|
|
message["content"] = _stub_for(message)
|
|
working[index] = message
|
|
dropped += 1
|
|
if dropped:
|
|
notes.append(f"summarised {dropped} old tool result(s) to one line each")
|
|
if estimate_tokens(working) <= CEILING_TOKENS:
|
|
return working, notes
|
|
|
|
# --- last resort: one model call ------------------------------------
|
|
tail_start = max(protect, len(working) - (KEEP_VERBATIM * 2))
|
|
head, tail = working[protect:tail_start], working[tail_start:]
|
|
if not head:
|
|
return working, notes
|
|
prose = summarise(head)
|
|
rebuilt = working[:protect] + [{
|
|
"role": "user",
|
|
"content": (
|
|
"PROGRESS SO FAR (an automatic summary of earlier steps in this "
|
|
"same session - it replaces them, and they are gone):\n\n" + prose
|
|
),
|
|
}] + tail
|
|
notes.append(f"summarised {len(head)} earlier messages into one block")
|
|
|
|
# Even after a summary the tail can be over budget - a handful of
|
|
# very large results is enough. Keep stubbing from the oldest, always
|
|
# leaving the last two intact, because those are what the agent is
|
|
# actually working from right now. Without this the ceiling is a
|
|
# suggestion rather than a ceiling, which is how a session ends with
|
|
# a context error instead of an answer.
|
|
tail_tools = [i for i, m in enumerate(rebuilt)
|
|
if m.get("role") == "tool" and i >= protect + 1]
|
|
extra = 0
|
|
for index in tail_tools[:-2]:
|
|
if estimate_tokens(rebuilt) <= CEILING_TOKENS:
|
|
break
|
|
if len(rebuilt[index].get("content") or "") < 200:
|
|
continue
|
|
message = dict(rebuilt[index])
|
|
message["content"] = _stub_for(message)
|
|
rebuilt[index] = message
|
|
extra += 1
|
|
if extra:
|
|
notes.append(f"stubbed {extra} more result(s) to stay under the ceiling")
|
|
return rebuilt, notes
|
|
|
|
|
|
SUMMARY_PROMPT = """\
|
|
Summarise the work so far in this VM session, for your own use in the steps ahead.
|
|
|
|
Write it as notes to yourself, not as a report to anyone else. Include, in this order:
|
|
|
|
1. What has actually been done, and what the machine said when it was done.
|
|
2. What was tried and did NOT work, with the error - so it is not tried again.
|
|
3. The current state of the machine: what is installed, what files exist, what is running.
|
|
4. Decisions taken and why, so they are not relitigated.
|
|
|
|
Be specific about names, paths, versions and exit codes; those are the details that get \
|
|
lost and then invented. Leave out everything you would not need again. No preamble.
|
|
"""
|