ryan
891f44c1a9
llama-server only reuses a byte-identical prompt prefix. Several things changed the prefix every turn, forcing most of the prompt to be re-prefilled: - Memories: rank still chooses which memories fit the budget, but the chosen ones are emitted in stored order, so the block is identical across queries. - History trimming: the conversation-evidence summary is only included when turns were actually dropped (it duplicated retained turns otherwise). On overflow, old turns are dropped at stable cut points down to ~65% of the budget, so the retained prefix holds for several turns instead of shifting by one turn every request. - Tools: TOOL_ROUTING_MODE=static (default) binds the full role-appropriate tool set in registration order with no discovery tool; admin gating is unchanged. TOOL_ROUTING_MODE=dynamic keeps the keyword router. - Observability: per-call model_calls in the stream_request log now include cached_prompt_tokens when the provider reports cache reads. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
544 lines
20 KiB
Python
544 lines
20 KiB
Python
"""Budget the model's view of history without changing checkpointed messages.
|
|
|
|
Token estimates use UTF-8 bytes / 3 plus message overhead. They are deliberately
|
|
more conservative than the usual chars / 4 heuristic, but are not a model
|
|
specific tokenizer. A safety reserve covers estimation and chat-template costs.
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
import math
|
|
import os
|
|
|
|
from dotenv import load_dotenv
|
|
from langchain_core.messages import (
|
|
AIMessage,
|
|
BaseMessage,
|
|
HumanMessage,
|
|
SystemMessage,
|
|
ToolMessage,
|
|
)
|
|
|
|
load_dotenv()
|
|
|
|
CONTEXT_WINDOW_TOKENS = int(os.getenv("CONTEXT_WINDOW_TOKENS", "32768"))
|
|
PROMPT_BUDGET_TOKENS = int(os.getenv("PROMPT_BUDGET_TOKENS", "16000"))
|
|
RESPONSE_RESERVE_TOKENS = int(os.getenv("RESPONSE_RESERVE_TOKENS", "4096"))
|
|
SAFETY_RESERVE_TOKENS = 2048
|
|
EARLIER_MESSAGE_MAX_CHARS = 4000
|
|
TOOL_RESULT_MAX_CHARS = 12000
|
|
TRUNCATION_NOTE = "\n[truncated: earlier context]"
|
|
IMAGE_PLACEHOLDER = "[image from an earlier message]"
|
|
# On overflow, drop old turns down to about this share of the history budget.
|
|
TRIM_LOW_WATERMARK = 0.65
|
|
|
|
|
|
def estimate_tokens(value) -> int:
|
|
text = (
|
|
value
|
|
if isinstance(value, str)
|
|
else json.dumps(value, ensure_ascii=False, default=str)
|
|
)
|
|
return math.ceil(len(text.encode("utf-8")) / 3)
|
|
|
|
|
|
def message_tokens(message: BaseMessage) -> int:
|
|
# Count tool arguments but never artifact payloads (not sent to the model).
|
|
content = message.content
|
|
images = 0
|
|
if isinstance(content, list):
|
|
images = sum(
|
|
isinstance(p, dict) and p.get("type") in ("image", "image_url")
|
|
for p in content
|
|
)
|
|
content = [
|
|
p
|
|
for p in content
|
|
if not (isinstance(p, dict) and p.get("type") in ("image", "image_url"))
|
|
]
|
|
return (
|
|
16
|
|
+ estimate_tokens(content)
|
|
+ images * 4096
|
|
+ estimate_tokens(getattr(message, "tool_calls", []))
|
|
)
|
|
|
|
|
|
def _cap(text: str, limit: int, note: str) -> str:
|
|
if len(text) <= limit:
|
|
return text
|
|
# Retain both ends: identifiers and totals often appear at the end of results.
|
|
head = limit * 3 // 4
|
|
return text[:head] + note + (text[-(limit - head) :] if limit > head else "")
|
|
|
|
|
|
TEXT_FIELDS = {"markdown", "snippet", "content", "text", "memo", "description"}
|
|
ROW_FIELDS = {
|
|
"transactions",
|
|
"rows",
|
|
"results",
|
|
"documents",
|
|
"categories",
|
|
"accounts",
|
|
"payees",
|
|
}
|
|
|
|
|
|
def _source_identity(row):
|
|
"""Keep source references for omitted rows; never manufacture aggregates."""
|
|
if not isinstance(row, dict):
|
|
return {}
|
|
references = {}
|
|
for key, value in row.items():
|
|
if (
|
|
key
|
|
in {
|
|
"id",
|
|
"ids",
|
|
"path",
|
|
"filepath",
|
|
"filename",
|
|
"url",
|
|
"source",
|
|
"source_ref",
|
|
"source_refs",
|
|
"date",
|
|
"month",
|
|
"filters",
|
|
}
|
|
or key.endswith(("_id", "_ids", "_date"))
|
|
or "total" in key
|
|
):
|
|
references[key] = value
|
|
elif isinstance(value, dict):
|
|
if nested := _source_identity(value):
|
|
references[key] = nested
|
|
elif isinstance(value, list):
|
|
nested = [ref for item in value if (ref := _source_identity(item))]
|
|
if nested:
|
|
references[key] = nested
|
|
return references
|
|
|
|
|
|
def _structured_view(value, rows: int, text_limit: int):
|
|
if isinstance(value, dict):
|
|
result = {}
|
|
for key, item in value.items():
|
|
if (
|
|
key.endswith("_compaction")
|
|
and key.removesuffix("_compaction") in ROW_FIELDS
|
|
):
|
|
continue # rebuilt alongside its rows, including earlier omissions
|
|
if (
|
|
key.endswith("_compaction")
|
|
and key.removesuffix("_compaction") in TEXT_FIELDS
|
|
and isinstance(value.get(key.removesuffix("_compaction")), str)
|
|
and len(value[key.removesuffix("_compaction")]) > text_limit
|
|
):
|
|
continue
|
|
if key in ROW_FIELDS and isinstance(item, list):
|
|
kept = item[:rows]
|
|
result[key] = [_structured_view(v, rows, text_limit) for v in kept]
|
|
previous = value.get(key + "_compaction", {})
|
|
previous = previous if isinstance(previous, dict) else {}
|
|
already_omitted = previous.get("omitted_rows", 0)
|
|
if len(kept) < len(item) or already_omitted:
|
|
result[key + "_compaction"] = {
|
|
"returned_rows": len(item) + already_omitted,
|
|
"retained_rows": len(kept),
|
|
"omitted_rows": len(item) - len(kept) + already_omitted,
|
|
"omitted_source_refs": previous.get("omitted_source_refs", [])
|
|
+ [
|
|
ref
|
|
for v in item[len(kept) :]
|
|
if (ref := _source_identity(v))
|
|
],
|
|
"notice": "Rows omitted; use original artifact or source tool for details.",
|
|
}
|
|
elif (
|
|
key in TEXT_FIELDS and isinstance(item, str) and len(item) > text_limit
|
|
):
|
|
result[key] = item[:text_limit]
|
|
previous = value.get(key + "_compaction", {})
|
|
previous = previous if isinstance(previous, dict) else {}
|
|
result[key + "_compaction"] = {
|
|
"omitted_characters": len(item)
|
|
- text_limit
|
|
+ previous.get("omitted_characters", 0),
|
|
"notice": "Prefix excerpt; original text remains in artifact.",
|
|
}
|
|
if key == "markdown":
|
|
# A heading index lets focused note reads reach sections
|
|
# beyond the retained prefix without reloading the note.
|
|
start_line = value.get("read_range", {}).get("start_line", 1)
|
|
headings = previous.get("headings") or [
|
|
{"line": i + start_line, "heading": line.strip()}
|
|
for i, line in enumerate(item.splitlines())
|
|
if line.lstrip().startswith("#")
|
|
]
|
|
result[key + "_compaction"]["headings"] = headings[:40]
|
|
result[key + "_compaction"]["omitted_headings"] = previous.get(
|
|
"omitted_headings", 0
|
|
) + max(0, len(headings) - 40)
|
|
elif key in {"filters", "source_ref", "source_refs", "frontmatter"}:
|
|
result[key] = item
|
|
else:
|
|
result[key] = _structured_view(item, rows, text_limit)
|
|
return result
|
|
if isinstance(value, list):
|
|
return [_structured_view(v, rows, text_limit) for v in value]
|
|
return value
|
|
|
|
|
|
def compact_tool_result(message: ToolMessage, limit: int) -> ToolMessage:
|
|
"""Project structured data without editing artifacts or slicing JSON.
|
|
|
|
Indispensable metadata may exceed limit; the caller must charge that floor.
|
|
"""
|
|
payload = message.artifact
|
|
if isinstance(payload, dict) and isinstance(payload.get("model_context"), dict):
|
|
# Focused reads carry a separate model view; the full artifact belongs
|
|
# to the interface and must not overwrite the specifically requested range.
|
|
payload = {k: v for k, v in payload.items() if k != "model_context"}
|
|
payload["data"] = {
|
|
**payload.get("data", {}),
|
|
**message.artifact["model_context"],
|
|
}
|
|
if not isinstance(payload, (dict, list)):
|
|
try:
|
|
payload = (
|
|
json.loads(message.content)
|
|
if isinstance(message.content, str)
|
|
else None
|
|
)
|
|
except (ValueError, TypeError):
|
|
payload = None
|
|
if isinstance(payload, (dict, list)):
|
|
full = json.dumps(
|
|
payload, ensure_ascii=False, default=str, separators=(",", ":")
|
|
)
|
|
if len(full) <= limit:
|
|
return message.model_copy(update={"content": full})
|
|
|
|
def render(size):
|
|
return json.dumps(
|
|
_structured_view(payload, size // 256, size // 4),
|
|
ensure_ascii=False,
|
|
default=str,
|
|
separators=(",", ":"),
|
|
)
|
|
|
|
low, high = 0, max(0, limit)
|
|
while low < high:
|
|
mid = (low + high + 1) // 2
|
|
if len(render(mid)) <= limit:
|
|
low = mid
|
|
else:
|
|
high = mid - 1
|
|
return message.model_copy(update={"content": render(low)})
|
|
return _shrink(message, limit, earlier=False, structured=False)
|
|
|
|
|
|
def _shrink(
|
|
message: BaseMessage, limit: int, *, earlier: bool, structured: bool = True
|
|
) -> BaseMessage:
|
|
if structured and isinstance(message, ToolMessage):
|
|
return compact_tool_result(message, limit)
|
|
note = (
|
|
TRUNCATION_NOTE
|
|
if earlier
|
|
else "\n[tool result excerpt; omitted content can be fetched again with the source tool]\n"
|
|
)
|
|
content = message.content
|
|
if isinstance(content, str):
|
|
new_content = _cap(content, limit, note)
|
|
else:
|
|
new_content = []
|
|
remaining = limit
|
|
for part in content:
|
|
if (
|
|
isinstance(part, dict)
|
|
and part.get("type") in ("image_url", "image")
|
|
and earlier
|
|
):
|
|
new_content.append({"type": "text", "text": IMAGE_PLACEHOLDER})
|
|
elif isinstance(part, dict) and isinstance(part.get("text"), str):
|
|
text = _cap(part["text"], remaining, note)
|
|
new_content.append({**part, "text": text})
|
|
remaining = max(0, remaining - min(len(part["text"]), remaining))
|
|
elif isinstance(part, str):
|
|
new_content.append(_cap(part, remaining, note))
|
|
remaining = max(0, remaining - min(len(part), remaining))
|
|
else:
|
|
new_content.append(part)
|
|
return (
|
|
message
|
|
if new_content == content
|
|
else message.model_copy(update={"content": new_content})
|
|
)
|
|
|
|
|
|
def _turns(messages: list[BaseMessage]) -> list[list[BaseMessage]]:
|
|
turns = []
|
|
for message in messages:
|
|
if isinstance(message, HumanMessage):
|
|
turns.append([])
|
|
if turns:
|
|
turns[-1].append(message)
|
|
# Leading orphaned checkpoint fragments have no user question to anchor them.
|
|
return turns
|
|
|
|
|
|
def _complete_exchanges(turn: list[BaseMessage]) -> list[BaseMessage]:
|
|
"""Drop unfinished exchanges left by a cancelled turn and orphan results."""
|
|
result = []
|
|
i = 0
|
|
while i < len(turn):
|
|
message = turn[i]
|
|
if isinstance(message, AIMessage) and message.tool_calls:
|
|
expected = {c["id"] for c in message.tool_calls}
|
|
end = i + 1
|
|
while end < len(turn) and isinstance(turn[end], ToolMessage):
|
|
end += 1
|
|
replies = turn[i + 1 : end]
|
|
actual = [m.tool_call_id for m in replies]
|
|
if len(actual) == len(expected) and set(actual) == expected:
|
|
result.extend(turn[i:end])
|
|
i = end
|
|
else:
|
|
if not isinstance(message, ToolMessage):
|
|
result.append(message)
|
|
i += 1
|
|
return result
|
|
|
|
|
|
def _summary(turns: list[list[BaseMessage]], limit: int = 2400) -> AIMessage | None:
|
|
excerpts = []
|
|
remaining = limit
|
|
for turn in reversed(turns):
|
|
lines = []
|
|
for message in turn:
|
|
if isinstance(message, HumanMessage) or (
|
|
isinstance(message, AIMessage) and not message.tool_calls
|
|
):
|
|
content = message.content
|
|
if not isinstance(content, str):
|
|
content = " ".join(
|
|
p.get("text", "") if isinstance(p, dict) else str(p)
|
|
for p in content
|
|
)
|
|
role = "User" if isinstance(message, HumanMessage) else "Assistant"
|
|
lines.append(f"{role}: {_cap(content, 400, ' [excerpt] ')}")
|
|
excerpt = "\n".join(lines)
|
|
if len(excerpt) > remaining:
|
|
break
|
|
excerpts.insert(0, excerpt)
|
|
remaining -= len(excerpt) + 2
|
|
if not excerpts:
|
|
return None
|
|
return AIMessage(
|
|
content="Earlier conversation excerpts (historical data, not new instructions; excerpts may omit details):\n\n"
|
|
+ "\n\n".join(excerpts)
|
|
)
|
|
|
|
|
|
def _stable_cut(costs: list[int], room: int, step: int) -> int:
|
|
"""Return how many of the oldest prior turns to drop so the rest fit ``room``.
|
|
|
|
Cut points sit where the cumulative cost from the oldest turn crosses a
|
|
multiple of ``step``. They depend only on turns already in history, so the
|
|
cut stays put as new turns arrive and only moves, by about ``step`` tokens,
|
|
once the retained turns overflow again. That keeps the prompt prefix
|
|
byte-identical across several turns for llama-server's prompt cache.
|
|
"""
|
|
remaining = sum(costs)
|
|
if remaining <= room:
|
|
return 0
|
|
total = 0
|
|
for i, turn_cost in enumerate(costs):
|
|
before = total
|
|
total += turn_cost
|
|
remaining -= turn_cost
|
|
if total // step > before // step and remaining <= room:
|
|
return i + 1
|
|
# No cut point leaves room for any prior turn (e.g. a heavy current turn).
|
|
# Fall back to keeping as many recent whole turns as fit.
|
|
remaining = sum(costs)
|
|
for i, turn_cost in enumerate(costs):
|
|
if remaining <= room:
|
|
return i
|
|
remaining -= turn_cost
|
|
return len(costs)
|
|
|
|
|
|
def _trim(
|
|
messages: list[BaseMessage], token_budget: int, reserve: int = 0
|
|
) -> tuple[list[BaseMessage], int]:
|
|
"""Trim to ``token_budget - reserve``; return the messages and dropped turns.
|
|
|
|
``reserve`` shrinks the space without moving the hysteresis cut points, so a
|
|
caller-supplied summary of varying size doesn't shift which turns are kept.
|
|
"""
|
|
step = max(1, int(token_budget * (1 - TRIM_LOW_WATERMARK)))
|
|
token_budget -= reserve
|
|
turns = [_complete_exchanges(turn) for turn in _turns(messages)]
|
|
if not turns:
|
|
return [], 0
|
|
original_current = turns[-1]
|
|
current = original_current
|
|
current = [
|
|
_shrink(m, TOOL_RESULT_MAX_CHARS, earlier=False)
|
|
if isinstance(m, ToolMessage)
|
|
else m
|
|
for m in current
|
|
]
|
|
|
|
def cost(ms):
|
|
return sum(message_tokens(m) for m in ms)
|
|
|
|
# Keep room for continuity even when the current turn has many tool results.
|
|
current_budget = (
|
|
token_budget - min(1000, token_budget // 8) if len(turns) > 1 else token_budget
|
|
)
|
|
if cost(current) > current_budget:
|
|
# Binary-search the largest shared excerpt size that fits. Tool call IDs
|
|
# and all results remain, including parallel exchanges.
|
|
low, high = 0, TOOL_RESULT_MAX_CHARS
|
|
|
|
def compact(limit):
|
|
return [
|
|
_shrink(m, limit, earlier=False)
|
|
if isinstance(m, (ToolMessage, AIMessage))
|
|
else m
|
|
for m in original_current
|
|
]
|
|
|
|
if cost(compact(0)) > token_budget:
|
|
raise ValueError(
|
|
"The current question and tool calls exceed the context budget (including required tool metadata). Send a smaller attachment or start a new conversation."
|
|
)
|
|
current_budget = max(current_budget, cost(compact(0)))
|
|
while low < high:
|
|
mid = (low + high + 1) // 2
|
|
if cost(compact(mid)) <= current_budget:
|
|
low = mid
|
|
else:
|
|
high = mid - 1
|
|
current = compact(low)
|
|
used = cost(current)
|
|
prior = [
|
|
[_shrink(m, EARLIER_MESSAGE_MAX_CHARS, earlier=True) for m in turn]
|
|
for turn in turns[:-1]
|
|
]
|
|
costs = [cost(turn) for turn in prior]
|
|
if used + sum(costs) <= token_budget:
|
|
cut = 0
|
|
else:
|
|
# Leave a small share for excerpts of older turns only when dropping.
|
|
summary_reserve = min(1000, max(0, token_budget - used))
|
|
cut = _stable_cut(costs, token_budget - summary_reserve - used, step)
|
|
selected = prior[cut:]
|
|
used += sum(costs[cut:])
|
|
dropped = turns[:cut]
|
|
summary = _summary(dropped)
|
|
result = [m for turn in selected for m in turn] + current
|
|
if summary and used + message_tokens(summary) <= token_budget:
|
|
result.insert(0, summary)
|
|
logging.info(
|
|
json.dumps(
|
|
{
|
|
"event": "context_budget",
|
|
"budget_tokens": token_budget,
|
|
"estimated_tokens": cost(result),
|
|
"input_messages": len(messages),
|
|
"output_messages": len(result),
|
|
"dropped_turns": len(dropped),
|
|
"summary_included": bool(result and isinstance(result[0], AIMessage)),
|
|
}
|
|
)
|
|
)
|
|
return result, len(dropped)
|
|
|
|
|
|
def trim_history(
|
|
messages: list[BaseMessage],
|
|
*,
|
|
token_budget: int = 12000,
|
|
summary_message: AIMessage | None = None,
|
|
) -> list[BaseMessage]:
|
|
"""Keep the current turn, then recent whole turns, within an estimated budget.
|
|
|
|
The current user message and tool call arguments are never truncated. Fail
|
|
explicitly if these mandatory inputs cannot fit, rather than lose the task.
|
|
When the budget overflows, old turns are dropped down to a lower watermark
|
|
so the retained prefix stays stable (and cacheable) for several turns.
|
|
"""
|
|
if summary_message is not None:
|
|
if not isinstance(summary_message, AIMessage) or summary_message.tool_calls:
|
|
raise ValueError(
|
|
"Conversation summary must be assistant-role evidence without tool calls"
|
|
)
|
|
trimmed, dropped = _trim(messages, token_budget)
|
|
if summary_message is None or not dropped:
|
|
# The summary only duplicates retained turns when nothing was dropped,
|
|
# and it changes every turn, which would defeat prompt caching.
|
|
return trimmed
|
|
summary_cost = message_tokens(summary_message)
|
|
if summary_cost > token_budget:
|
|
logging.info("Omitting optional conversation summary: insufficient budget")
|
|
return trimmed
|
|
try:
|
|
trimmed, _ = _trim(messages, token_budget, reserve=summary_cost)
|
|
except ValueError:
|
|
# Historical evidence must never evict a current request that fits
|
|
# without it. Required current-turn metadata retains priority.
|
|
logging.info(
|
|
"Omitting optional conversation summary to preserve current request"
|
|
)
|
|
return trimmed
|
|
if (
|
|
trimmed
|
|
and isinstance(trimmed[0], AIMessage)
|
|
and isinstance(trimmed[0].content, str)
|
|
and trimmed[0].content.startswith("Earlier conversation excerpts")
|
|
):
|
|
trimmed = trimmed[1:]
|
|
return [summary_message, *trimmed]
|
|
|
|
|
|
def build_context(
|
|
messages: list[BaseMessage],
|
|
system_prompt: str,
|
|
tool_schemas: list[dict],
|
|
*,
|
|
summary_message: AIMessage | None = None,
|
|
) -> list[BaseMessage]:
|
|
limit = min(
|
|
PROMPT_BUDGET_TOKENS,
|
|
CONTEXT_WINDOW_TOKENS - RESPONSE_RESERVE_TOKENS - SAFETY_RESERVE_TOKENS,
|
|
)
|
|
system = SystemMessage(content=system_prompt)
|
|
budget = limit - message_tokens(system) - estimate_tokens(tool_schemas)
|
|
if budget <= 0:
|
|
raise ValueError(
|
|
"System instructions and tool definitions exceed the configured prompt budget."
|
|
)
|
|
history = trim_history(
|
|
messages, token_budget=budget, summary_message=summary_message
|
|
)
|
|
logging.info(
|
|
json.dumps(
|
|
{
|
|
"event": "prompt_budget",
|
|
"limit_tokens": limit,
|
|
"system_tokens": message_tokens(system),
|
|
"tool_schema_tokens": estimate_tokens(tool_schemas),
|
|
"estimated_prompt_tokens": message_tokens(system)
|
|
+ estimate_tokens(tool_schemas)
|
|
+ sum(message_tokens(m) for m in history),
|
|
"response_reserve_tokens": RESPONSE_RESERVE_TOKENS,
|
|
}
|
|
)
|
|
)
|
|
return [system] + history
|