ryan/simbarag
1
0
Fork 0
Code Issues Pull Requests 5 Actions 1 Packages Projects Releases Wiki Activity

Keep the chat prompt prefix stable for llama-server prompt caching

Browse Source 
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>
This commit is contained in:
Ryan Chen
2026-10-12 23:01:46 -04:00
parent 030c5d6c21
commit 891f44c1a9
12 changed files with 338 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.env.example
+3
View File
@@ -32,6 +32,9 @@ TTFT_TIMEOUT_SECONDS=30
CONTEXT_WINDOW_TOKENS=32768
PROMPT_BUDGET_TOKENS=16000
RESPONSE_RESERVE_TOKENS=4096
# static: bind every role-appropriate tool on each call (stable, cacheable prompt
# prefix). dynamic: keyword-route tool groups per call with a discovery tool.
TOOL_ROUTING_MODE=static
# Optional durable summary sidecar; defaults to CHECKPOINT_DB_PATH + .summaries.db
# Store on the same persistent shared volume as the checkpoints.
# CONVERSATION_SUMMARY_DB_PATH=/app/checkpoints/conversation-summaries.db
blueprints/conversation/__init__.py
+5
View File
@@ -75,6 +75,7 @@ def _model_call_summary(call: dict) -> dict:
"prefill_ms": round((first_chunk - start) * 1000) if first_chunk else None,
"ms": round((end - start) * 1000) if end else None,
"prompt_tokens": call.get("prompt_tokens"),
"cached_prompt_tokens": call.get("cached_prompt_tokens"),
"output_tokens": call.get("output_tokens"),
"model": call.get("model"),
"provider": call.get("provider"),
@@ -600,6 +601,10 @@ async def stream_query():
if usage:
call["prompt_tokens"] = usage.get("input_tokens")
call["output_tokens"] = usage.get("output_tokens")
# Prompt-cache hits, when the provider reports them.
call["cached_prompt_tokens"] = (
usage.get("input_token_details") or {}
).get("cache_read")
call["reasoning_tokens"] = (
usage.get("output_token_details") or {}
).get("reasoning")
blueprints/conversation/agents.py
+16 -8
View File
@@ -29,7 +29,7 @@ from blueprints.conversation.artifacts_util import (
from blueprints.conversation.context import RESPONSE_RESERVE_TOKENS, build_context
from blueprints.conversation.memory import save_memory
from blueprints.conversation.prompts import SIMBA_SYSTEM_PROMPT
from blueprints.conversation.tool_routing import ToolRouter
from blueprints.conversation.tool_routing import TOOL_ROUTING_MODES, ToolRouter
from blueprints.rag.logic import query_vector_store
from blueprints.ynab.sync import ynab_configured
from blueprints.ynab.tools import YNAB_TOOLS
@@ -1160,16 +1160,24 @@ def build_agent(checkpointer=None, extra_tools=()):
"""
extra_tools = list(extra_tools)
# Keep the guarded Google tools executable for admins even when no MCP tools
# are configured and app.admin_agent aliases app.main_agent. The router never
# binds their schemas for a regular user. Extra tools exist only on the admin
# graph and are independently gated by the same invocation config.
# are configured and app.admin_agent aliases app.main_agent. Neither routing
# mode binds their schemas for a regular user. Extra tools exist only on the
# admin graph and are independently gated by the same invocation config.
agent_tools = tools + extra_tools
router = ToolRouter(agent_tools, admin_tools=extra_tools)
mode = os.getenv("TOOL_ROUTING_MODE", "static").strip().lower()
if mode not in TOOL_ROUTING_MODES:
raise ValueError(
f"TOOL_ROUTING_MODE must be one of {TOOL_ROUTING_MODES}, got {mode!r}"
)
dynamic = mode == "dynamic"
def routed_model(messages, config: RunnableConfig):
selected = router.select(
messages,
is_admin=(config or {}).get("configurable", {}).get("is_admin", False),
is_admin = (config or {}).get("configurable", {}).get("is_admin", False)
selected = (
router.select(messages, is_admin=is_admin)
if dynamic
else router.select_static(is_admin=is_admin)
)
schemas = [convert_to_openai_tool(t) for t in selected]
fallback = openai_fallback.bind_tools(selected)
@@ -1188,6 +1196,6 @@ def build_agent(checkpointer=None, extra_tools=()):
return create_react_agent(
model=select_model,
tools=agent_tools + [router.discovery],
tools=agent_tools + ([router.discovery] if dynamic else []),
checkpointer=checkpointer,
)
blueprints/conversation/context.py
+101 -49
View File
@@ -29,6 +29,8 @@ 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:
@@ -340,46 +342,48 @@ def _summary(turns: list[list[BaseMessage]], limit: int = 2400) -> AIMessage | N
)
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.
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``.
The current user message and tool call arguments are never truncated. Fail
explicitly if these mandatory inputs cannot fit, rather than lose the task.
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.
"""
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"
)
summary_cost = message_tokens(summary_message)
if summary_cost > token_budget:
logging.info("Omitting optional conversation summary: insufficient budget")
return trim_history(messages, token_budget=token_budget)
try:
trimmed = trim_history(messages, token_budget=token_budget - 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 trim_history(messages, token_budget=token_budget)
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]
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 []
return [], 0
original_current = turns[-1]
current = original_current
current = [
@@ -421,19 +425,21 @@ def trim_history(
else:
high = mid - 1
current = compact(low)
selected = []
used = cost(current)
prior = turns[:-1]
# Leave a small share for excerpts of older turns only when necessary.
summary_reserve = min(1000, max(0, token_budget - used))
for turn in reversed(prior):
shrunk = [_shrink(m, EARLIER_MESSAGE_MAX_CHARS, earlier=True) for m in turn]
turn_cost = cost(shrunk)
if used + turn_cost > token_budget - summary_reserve:
break
selected.insert(0, shrunk)
used += turn_cost
dropped = prior[: len(prior) - len(selected)]
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:
@@ -451,7 +457,53 @@ def trim_history(
}
)
)
return result
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(
blueprints/conversation/memory.py
+11 -2
View File
@@ -1,3 +1,5 @@
from collections import Counter
from .models import UserMemory
from .prompts import SIMBA_SYSTEM_PROMPT
from .context import estimate_tokens
@@ -41,16 +43,23 @@ async def build_system_prompt_with_memories(user_id: str, query: str = "") -> st
memories = await get_memories_for_user(user_id)
if not memories:
return SIMBA_SYSTEM_PROMPT
selected = []
header = "\n\nUSER MEMORIES (facts the user has asked you to remember):\n"
omission = "\n[Additional saved memories omitted from this prompt.]"
# Reserve wrapper and omission costs so the entire injected memory block fits.
used = estimate_tokens(header) + estimate_tokens(omission)
chosen = Counter()
for memory in rank_memories(memories, query):
cost = estimate_tokens(f"- {memory}\n")
if used + cost <= MEMORY_BUDGET_TOKENS:
selected.append(memory)
chosen[memory] += 1
used += cost
# Rank only chooses what fits; emit in stored order so the system prompt is
# byte-identical across queries and llama-server can reuse its KV cache.
selected = []
for memory in memories:
if chosen[memory]:
chosen[memory] -= 1
selected.append(memory)
block = "\n".join(f"- {m}" for m in selected)
if len(selected) < len(memories):
block += omission
blueprints/conversation/tool_routing.py
+16 -1
View File
@@ -1,4 +1,8 @@
"""Deterministic per-call routing; discovery grants live in tool messages, not globals."""
"""Tool selection per model call.
Static mode (the default) binds the full role-appropriate set. Dynamic mode
routes per call by keyword; discovery grants live in tool messages, not globals.
"""
import json
import re
@@ -29,6 +33,7 @@ ADMIN_TOOL_NAMES = {
}
CORE_TOOL_NAMES = {"get_current_date", "save_user_memory"}
DISCOVERY_NAME = "discover_tool_groups"
TOOL_ROUTING_MODES = ("static", "dynamic")
def words(text):
@@ -138,6 +143,16 @@ class ToolRouter:
)
return discover_tool_groups
def select_static(self, *, is_admin=False):
"""Every role-appropriate tool in registration order, without discovery.
The bound set never varies by query or turn, so the rendered tool block
stays byte-identical and llama-server can reuse its cached prefix.
"""
return [
t for t in self.tools if is_admin or tool_name(t) not in self.admin_names
]
def select(self, messages, *, is_admin=False):
available = self.available_groups(is_admin)
human_indices = [
docker-compose.yml
+1
View File
@@ -15,6 +15,7 @@ x-raggr-common: &raggr-common
- JWT_SECRET_KEY=${JWT_SECRET_KEY}
- LLAMA_SERVER_URL=${LLAMA_SERVER_URL}
- LLAMA_MODEL_NAME=${LLAMA_MODEL_NAME}
- TOOL_ROUTING_MODE=${TOOL_ROUTING_MODE:-static}
- OIDC_ISSUER=${OIDC_ISSUER}
- OIDC_CLIENT_ID=${OIDC_CLIENT_ID}
- OIDC_CLIENT_SECRET=${OIDC_CLIENT_SECRET}
tests/unit/test_context_summary.py
+2 -1
View File
@@ -102,7 +102,8 @@ async def test_tools_and_incomplete_exchanges_do_not_become_summary_instructions
assert "Now continue" not in message.content
prompt = build_context(messages, "trusted system", [], summary_message=message)
assert prompt[0].content == "trusted system"
assert prompt[1] == message
# Nothing was dropped, so the ledger would only duplicate retained turns.
assert message not in prompt
assert prompt[-1] is messages[-1]
tests/unit/test_context_trim.py
+49
View File
@@ -462,3 +462,52 @@ def test_focused_note_heading_index_uses_absolute_content_lines():
metadata = json.loads(projected.content)["data"]["markdown_compaction"]
assert metadata["headings"] == [{"line": 1001, "heading": "# Requested section"}]
assert msg.artifact["data"]["markdown"] == "full original"
def _conversation(turns, size=1000):
messages = []
for i in range(turns):
messages += [
HumanMessage(content=f"Question {i}"),
call(str(i)),
result(str(i), "data " * size),
AIMessage(content=f"Decision {i}"),
]
return messages
def test_summary_omitted_when_no_turns_dropped():
messages = _conversation(2, size=10) + [HumanMessage(content="next")]
summary = AIMessage(content="Earlier: keep dining low")
assert trim_history(messages, summary_message=summary) == trim_history(messages)
assert summary not in trim_history(messages, summary_message=summary)
def test_summary_prepended_when_turns_dropped():
messages = _conversation(12) + [HumanMessage(content="next")]
summary = AIMessage(content="Earlier: keep dining low")
out = trim_history(messages, token_budget=6000, summary_message=summary)
assert out[0] is summary
assert out[-1] is messages[-1]
assert cost(out) <= 6000
def test_trim_drops_to_low_watermark_and_prefix_stays_stable():
budget = 8000
messages, oldest, costs = [], [], []
for i in range(30):
messages.append(HumanMessage(content=f"Question {i}"))
out = trim_history(messages, token_budget=budget)
assert out[-1] is messages[-1]
assert cost(out) <= budget
oldest.append(next(m for m in out if isinstance(m, HumanMessage)).content)
costs.append(cost(out))
messages += [call(str(i)), result(str(i), "data " * 300), AIMessage("ok")]
cuts = [i for i in range(1, 30) if oldest[i] != oldest[i - 1]]
assert cuts, "history should overflow the budget"
# Each drop goes well below the budget instead of one turn under it...
for i in cuts:
assert costs[i] <= budget * (context.TRIM_LOW_WATERMARK + 0.1)
# ...so the retained prefix is unchanged for several turns in between.
assert all(b - a >= 3 for a, b in zip(cuts, cuts[1:]))
assert len(cuts) <= 5
tests/unit/test_memory_retrieval.py
+31
View File
@@ -46,3 +46,34 @@ async def test_backwards_compatible_no_query_and_empty_memories(monkeypatch):
)
fetch.return_value = ["Likes tea"]
assert "Likes tea" in await memory.build_system_prompt_with_memories("user")
@pytest.mark.asyncio
async def test_memory_block_is_identical_across_queries_when_all_fit(monkeypatch):
memories = ["Favorite restaurant is Nori", "Cat Simba takes insulin", "Gym is Y"]
monkeypatch.setattr(
memory, "get_memories_for_user", AsyncMock(return_value=memories)
)
prompts = {
await memory.build_system_prompt_with_memories("user", query=q)
for q in ["", "Simba insulin", "gym", "Nori restaurant"]
}
assert len(prompts) == 1
prompt = prompts.pop()
assert prompt.index("Nori") < prompt.index("insulin") < prompt.index("Gym")
@pytest.mark.asyncio
async def test_rank_chooses_memories_but_stored_order_is_kept(monkeypatch):
filler = "Favorite restaurant is Nori. " * 20
memories = ["Likes tea", filler, "Simba insulin dose is two units."]
monkeypatch.setattr(
memory, "get_memories_for_user", AsyncMock(return_value=memories)
)
monkeypatch.setattr(memory, "MEMORY_BUDGET_TOKENS", 110)
prompt = await memory.build_system_prompt_with_memories(
"user", query="Simba insulin"
)
assert filler not in prompt
# Insulin ranks first but is emitted after the newer memory.
assert prompt.index("Likes tea") < prompt.index("insulin")
tests/unit/test_stream_stall.py
+5 -2
View File
@@ -61,7 +61,7 @@ def _chunk(run_id, text):
}
def _end(run_id, prompt_tokens=None):
def _end(run_id, prompt_tokens=None, cached_tokens=None):
usage = (
{
"input_tokens": prompt_tokens,
@@ -71,6 +71,8 @@ def _end(run_id, prompt_tokens=None):
if prompt_tokens
else None
)
if usage and cached_tokens is not None:
usage["input_token_details"] = {"cache_read": cached_tokens}
return {
"event": "on_chat_model_end",
"run_id": run_id,
@@ -207,7 +209,7 @@ async def test_logs_per_call_prompt_tokens(owner, monkeypatch, caplog):
script = [
(0, _start("m1")),
(0, _chunk("m1", "Hi")),
(0, _end("m1", prompt_tokens=27262)),
(0, _end("m1", prompt_tokens=27262, cached_tokens=25000)),
(0, _final("Hi")),
]
@@ -224,6 +226,7 @@ async def test_logs_per_call_prompt_tokens(owner, monkeypatch, caplog):
assert summary["max_prompt_tokens"] == 27262
(call,) = summary["model_calls"]
assert call["prompt_tokens"] == 27262
assert call["cached_prompt_tokens"] == 25000
assert call["prefill_ms"] is not None and call["ms"] is not None
tests/unit/test_tool_routing.py
+98 -6
View File
@@ -136,11 +136,7 @@ def test_discovery_recovery_admin_gating_and_isolation(router):
)
@pytest.mark.asyncio
@pytest.mark.parametrize("primary_fails", [False, True])
async def test_real_graph_binds_discovered_schemas_and_accounts_same_subset(
monkeypatch, primary_fails
):
def import_agents():
# Constructors are stubbed before agents import: no tokens, services or network.
with (
patch("langchain_openai.ChatOpenAI"),
@@ -154,7 +150,16 @@ async def test_real_graph_binds_discovered_schemas_and_accounts_same_subset(
"utils.obsidian_service.ObsidianService", side_effect=ValueError("offline")
),
):
agents = importlib.import_module("blueprints.conversation.agents")
return importlib.import_module("blueprints.conversation.agents")
@pytest.mark.asyncio
@pytest.mark.parametrize("primary_fails", [False, True])
async def test_real_graph_binds_discovered_schemas_and_accounts_same_subset(
monkeypatch, primary_fails
):
monkeypatch.setenv("TOOL_ROUTING_MODE", "dynamic")
agents = import_agents()
bound_names, primary_names, accounted_names = [], [], []
class FakeModel:
@@ -238,3 +243,90 @@ def test_schema_dict_tools_remain_routable_and_admin_gated():
router = ToolRouter([extra], admin_tools=[extra])
assert extra not in router.select([HumanMessage(content="extra")])
assert extra in router.select([HumanMessage(content="extra")], is_admin=True)
def test_static_selection_is_stable_role_gated_and_has_no_discovery(router):
queries = ["Review spending", "Find Simba's vet records", "Help", "calendar"]
user = [router.select_static() for _ in queries]
assert all(selected == user[0] for selected in user)
assert [t.name for t in user[0]] == [
"get_current_date",
"save_user_memory",
"ynab_budget_summary",
"mealie_get_recipe",
"obsidian_read_note",
"simba_search",
"web_search",
]
admin = [t.name for t in router.select_static(is_admin=True)]
assert admin == [t.name for t in router.tools]
assert "discover_tool_groups" not in admin
def test_static_selection_gates_schema_dict_admin_tools():
extra = {
"type": "function",
"function": {
"name": "extra",
"description": "docs",
"parameters": {"type": "object", "properties": {}},
},
}
router = ToolRouter([make_tool("web_search"), extra], admin_tools=[extra])
assert extra not in router.select_static()
assert extra in router.select_static(is_admin=True)
@pytest.mark.asyncio
async def test_static_graph_binds_full_role_set_every_call(monkeypatch):
monkeypatch.delenv("TOOL_ROUTING_MODE", raising=False)
agents = import_agents()
bound_names = []
class FakeModel:
def bind_tools(self, tools):
bound_names.append([t.name for t in tools])
def respond(messages):
if len(bound_names) == 1:
return call("ynab_budget_summary")
return AIMessage(content="done")
return RunnableLambda(respond)
monkeypatch.setattr(
agents,
"tools",
[
make_tool("get_current_date"),
make_tool("ynab_budget_summary"),
make_tool("get_calendar_events"),
],
)
monkeypatch.setattr(agents, "llama_chat", None)
monkeypatch.setattr(agents, "openai_fallback", FakeModel())
monkeypatch.setattr(agents, "build_context", lambda messages, *a, **k: messages)
graph = agents.build_agent()
# A query matching no keyword group still sees every tool, with no
# discovery round trip, and the bound set is identical on every call.
result = await graph.ainvoke({"messages": [HumanMessage(content="Help")]})
assert result["messages"][-1].content == "done"
assert bound_names == [["get_current_date", "ynab_budget_summary"]] * 2
bound_names.clear()
await graph.ainvoke(
{"messages": [HumanMessage(content="Help")]},
config={"configurable": {"is_admin": True}},
)
assert bound_names[0] == [
"get_current_date",
"ynab_budget_summary",
"get_calendar_events",
]
def test_unknown_routing_mode_fails_explicitly(monkeypatch):
agents = import_agents()
monkeypatch.setenv("TOOL_ROUTING_MODE", "sometimes")
with pytest.raises(ValueError, match="TOOL_ROUTING_MODE"):
agents.build_agent()