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

Merged
ryan merged 3 commits from perf/stable-prompt-prefix into main 2026-10-12 23:06:08 -04:00
12 changed files with 400 additions and 69 deletions
+3
View File
@@ -32,6 +32,9 @@ TTFT_TIMEOUT_SECONDS=30
CONTEXT_WINDOW_TOKENS=32768 CONTEXT_WINDOW_TOKENS=32768
PROMPT_BUDGET_TOKENS=16000 PROMPT_BUDGET_TOKENS=16000
RESPONSE_RESERVE_TOKENS=4096 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 # Optional durable summary sidecar; defaults to CHECKPOINT_DB_PATH + .summaries.db
# Store on the same persistent shared volume as the checkpoints. # Store on the same persistent shared volume as the checkpoints.
# CONVERSATION_SUMMARY_DB_PATH=/app/checkpoints/conversation-summaries.db # CONVERSATION_SUMMARY_DB_PATH=/app/checkpoints/conversation-summaries.db
+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, "prefill_ms": round((first_chunk - start) * 1000) if first_chunk else None,
"ms": round((end - start) * 1000) if end else None, "ms": round((end - start) * 1000) if end else None,
"prompt_tokens": call.get("prompt_tokens"), "prompt_tokens": call.get("prompt_tokens"),
"cached_prompt_tokens": call.get("cached_prompt_tokens"),
"output_tokens": call.get("output_tokens"), "output_tokens": call.get("output_tokens"),
"model": call.get("model"), "model": call.get("model"),
"provider": call.get("provider"), "provider": call.get("provider"),
@@ -600,6 +601,10 @@ async def stream_query():
if usage: if usage:
call["prompt_tokens"] = usage.get("input_tokens") call["prompt_tokens"] = usage.get("input_tokens")
call["output_tokens"] = usage.get("output_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"] = ( call["reasoning_tokens"] = (
usage.get("output_token_details") or {} usage.get("output_token_details") or {}
).get("reasoning") ).get("reasoning")
+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.context import RESPONSE_RESERVE_TOKENS, build_context
from blueprints.conversation.memory import save_memory from blueprints.conversation.memory import save_memory
from blueprints.conversation.prompts import SIMBA_SYSTEM_PROMPT 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.rag.logic import query_vector_store
from blueprints.ynab.sync import ynab_configured from blueprints.ynab.sync import ynab_configured
from blueprints.ynab.tools import YNAB_TOOLS from blueprints.ynab.tools import YNAB_TOOLS
@@ -1160,16 +1160,24 @@ def build_agent(checkpointer=None, extra_tools=()):
""" """
extra_tools = list(extra_tools) extra_tools = list(extra_tools)
# Keep the guarded Google tools executable for admins even when no MCP 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 # are configured and app.admin_agent aliases app.main_agent. Neither routing
# binds their schemas for a regular user. Extra tools exist only on the admin # mode binds their schemas for a regular user. Extra tools exist only on the
# graph and are independently gated by the same invocation config. # admin graph and are independently gated by the same invocation config.
agent_tools = tools + extra_tools agent_tools = tools + extra_tools
router = ToolRouter(agent_tools, admin_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): def routed_model(messages, config: RunnableConfig):
selected = router.select( is_admin = (config or {}).get("configurable", {}).get("is_admin", False)
messages, selected = (
is_admin=(config or {}).get("configurable", {}).get("is_admin", False), 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] schemas = [convert_to_openai_tool(t) for t in selected]
fallback = openai_fallback.bind_tools(selected) fallback = openai_fallback.bind_tools(selected)
@@ -1188,6 +1196,6 @@ def build_agent(checkpointer=None, extra_tools=()):
return create_react_agent( return create_react_agent(
model=select_model, model=select_model,
tools=agent_tools + [router.discovery], tools=agent_tools + ([router.discovery] if dynamic else []),
checkpointer=checkpointer, checkpointer=checkpointer,
) )
+118 -49
View File
@@ -29,6 +29,10 @@ EARLIER_MESSAGE_MAX_CHARS = 4000
TOOL_RESULT_MAX_CHARS = 12000 TOOL_RESULT_MAX_CHARS = 12000
TRUNCATION_NOTE = "\n[truncated: earlier context]" TRUNCATION_NOTE = "\n[truncated: earlier context]"
IMAGE_PLACEHOLDER = "[image from an earlier message]" 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
# Upper bound on the summary ledger message (summary.STATE_TOKEN_LIMIT + wrapper).
SUMMARY_RESERVE_TOKENS = 2000
def estimate_tokens(value) -> int: def estimate_tokens(value) -> int:
@@ -340,46 +344,48 @@ def _summary(turns: list[list[BaseMessage]], limit: int = 2400) -> AIMessage | N
) )
def trim_history( def _stable_cut(costs: list[int], room: int, step: int) -> int:
messages: list[BaseMessage], """Return how many of the oldest prior turns to drop so the rest fit ``room``.
*,
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 Cut points sit where the cumulative cost from the oldest turn crosses a
explicitly if these mandatory inputs cannot fit, rather than lose the task. 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: remaining = sum(costs)
if not isinstance(summary_message, AIMessage) or summary_message.tool_calls: if remaining <= room:
raise ValueError( return 0
"Conversation summary must be assistant-role evidence without tool calls" total = 0
) for i, turn_cost in enumerate(costs):
summary_cost = message_tokens(summary_message) before = total
if summary_cost > token_budget: total += turn_cost
logging.info("Omitting optional conversation summary: insufficient budget") remaining -= turn_cost
return trim_history(messages, token_budget=token_budget) if total // step > before // step and remaining <= room:
try: return i + 1
trimmed = trim_history(messages, token_budget=token_budget - summary_cost) # No cut point leaves room for any prior turn (e.g. a heavy current turn).
except ValueError: # Fall back to keeping as many recent whole turns as fit.
# Historical evidence must never evict a current request that fits remaining = sum(costs)
# without it. Required current-turn metadata retains priority. for i, turn_cost in enumerate(costs):
logging.info( if remaining <= room:
"Omitting optional conversation summary to preserve current request" return i
) remaining -= turn_cost
return trim_history(messages, token_budget=token_budget) return len(costs)
if (
trimmed
and isinstance(trimmed[0], AIMessage) def _trim(
and isinstance(trimmed[0].content, str) messages: list[BaseMessage], token_budget: int, reserve: int = 0
and trimmed[0].content.startswith("Earlier conversation excerpts") ) -> tuple[list[BaseMessage], int]:
): """Trim to ``token_budget - reserve``; return the messages and dropped turns.
trimmed = trimmed[1:]
return [summary_message, *trimmed] ``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)] turns = [_complete_exchanges(turn) for turn in _turns(messages)]
if not turns: if not turns:
return [] return [], 0
original_current = turns[-1] original_current = turns[-1]
current = original_current current = original_current
current = [ current = [
@@ -421,19 +427,21 @@ def trim_history(
else: else:
high = mid - 1 high = mid - 1
current = compact(low) current = compact(low)
selected = []
used = cost(current) used = cost(current)
prior = turns[:-1] prior = [
# Leave a small share for excerpts of older turns only when necessary. [_shrink(m, EARLIER_MESSAGE_MAX_CHARS, earlier=True) for m in turn]
summary_reserve = min(1000, max(0, token_budget - used)) for turn in turns[:-1]
for turn in reversed(prior): ]
shrunk = [_shrink(m, EARLIER_MESSAGE_MAX_CHARS, earlier=True) for m in turn] costs = [cost(turn) for turn in prior]
turn_cost = cost(shrunk) if used + sum(costs) <= token_budget:
if used + turn_cost > token_budget - summary_reserve: cut = 0
break else:
selected.insert(0, shrunk) # Leave a small share for excerpts of older turns only when dropping.
used += turn_cost summary_reserve = min(1000, max(0, token_budget - used))
dropped = prior[: len(prior) - len(selected)] cut = _stable_cut(costs, token_budget - summary_reserve - used, step)
selected = prior[cut:]
used += sum(costs[cut:])
dropped = turns[:cut]
summary = _summary(dropped) summary = _summary(dropped)
result = [m for turn in selected for m in turn] + current result = [m for turn in selected for m in turn] + current
if summary and used + message_tokens(summary) <= token_budget: if summary and used + message_tokens(summary) <= token_budget:
@@ -451,7 +459,68 @@ 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
# Reserve a fixed ceiling, not the ledger's current size, so a shrinking
# ledger can't pull the cut back and change the retained prefix.
for reserve in dict.fromkeys(
(max(summary_cost, SUMMARY_RESERVE_TOKENS), summary_cost)
):
try:
with_room, _ = _trim(messages, token_budget, reserve=reserve)
break
except ValueError:
continue
else:
# 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
# The ledger supersedes excerpts of dropped turns; whenever it is included
# they are consistently absent, so this doesn't disturb the cached prefix.
if (
with_room
and isinstance(with_room[0], AIMessage)
and isinstance(with_room[0].content, str)
and with_room[0].content.startswith("Earlier conversation excerpts")
):
with_room = with_room[1:]
# The ledger changes every turn, so place it just before the current turn:
# system, tools and retained history then stay byte-identical for the
# prompt cache. Turns start at a HumanMessage, so the last one opens the
# current turn and no tool exchange is split.
start = max(i for i, m in enumerate(with_room) if isinstance(m, HumanMessage))
return [*with_room[:start], summary_message, *with_room[start:]]
def build_context( def build_context(
+11 -2
View File
@@ -1,3 +1,5 @@
from collections import Counter
from .models import UserMemory from .models import UserMemory
from .prompts import SIMBA_SYSTEM_PROMPT from .prompts import SIMBA_SYSTEM_PROMPT
from .context import estimate_tokens 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) memories = await get_memories_for_user(user_id)
if not memories: if not memories:
return SIMBA_SYSTEM_PROMPT return SIMBA_SYSTEM_PROMPT
selected = []
header = "\n\nUSER MEMORIES (facts the user has asked you to remember):\n" header = "\n\nUSER MEMORIES (facts the user has asked you to remember):\n"
omission = "\n[Additional saved memories omitted from this prompt.]" omission = "\n[Additional saved memories omitted from this prompt.]"
# Reserve wrapper and omission costs so the entire injected memory block fits. # Reserve wrapper and omission costs so the entire injected memory block fits.
used = estimate_tokens(header) + estimate_tokens(omission) used = estimate_tokens(header) + estimate_tokens(omission)
chosen = Counter()
for memory in rank_memories(memories, query): for memory in rank_memories(memories, query):
cost = estimate_tokens(f"- {memory}\n") cost = estimate_tokens(f"- {memory}\n")
if used + cost <= MEMORY_BUDGET_TOKENS: if used + cost <= MEMORY_BUDGET_TOKENS:
selected.append(memory) chosen[memory] += 1
used += cost 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) block = "\n".join(f"- {m}" for m in selected)
if len(selected) < len(memories): if len(selected) < len(memories):
block += omission block += omission
+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 json
import re import re
@@ -29,6 +33,7 @@ ADMIN_TOOL_NAMES = {
} }
CORE_TOOL_NAMES = {"get_current_date", "save_user_memory"} CORE_TOOL_NAMES = {"get_current_date", "save_user_memory"}
DISCOVERY_NAME = "discover_tool_groups" DISCOVERY_NAME = "discover_tool_groups"
TOOL_ROUTING_MODES = ("static", "dynamic")
def words(text): def words(text):
@@ -138,6 +143,16 @@ class ToolRouter:
) )
return discover_tool_groups 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): def select(self, messages, *, is_admin=False):
available = self.available_groups(is_admin) available = self.available_groups(is_admin)
human_indices = [ human_indices = [
+1
View File
@@ -15,6 +15,7 @@ x-raggr-common: &raggr-common
- JWT_SECRET_KEY=${JWT_SECRET_KEY} - JWT_SECRET_KEY=${JWT_SECRET_KEY}
- LLAMA_SERVER_URL=${LLAMA_SERVER_URL} - LLAMA_SERVER_URL=${LLAMA_SERVER_URL}
- LLAMA_MODEL_NAME=${LLAMA_MODEL_NAME} - LLAMA_MODEL_NAME=${LLAMA_MODEL_NAME}
- TOOL_ROUTING_MODE=${TOOL_ROUTING_MODE:-static}
- OIDC_ISSUER=${OIDC_ISSUER} - OIDC_ISSUER=${OIDC_ISSUER}
- OIDC_CLIENT_ID=${OIDC_CLIENT_ID} - OIDC_CLIENT_ID=${OIDC_CLIENT_ID}
- OIDC_CLIENT_SECRET=${OIDC_CLIENT_SECRET} - OIDC_CLIENT_SECRET=${OIDC_CLIENT_SECRET}
+16 -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 assert "Now continue" not in message.content
prompt = build_context(messages, "trusted system", [], summary_message=message) prompt = build_context(messages, "trusted system", [], summary_message=message)
assert prompt[0].content == "trusted system" 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] assert prompt[-1] is messages[-1]
@@ -162,3 +163,17 @@ async def test_deletion_is_scoped_and_survives_reopen(tmp_path):
await store.delete("owner", "c") await store.delete("owner", "c")
assert (await SummaryStore(path).load("owner", "c")).cursor == 0 assert (await SummaryStore(path).load("owner", "c")).cursor == 0
assert (await SummaryStore(path).load("other", "c")).cursor == 4 assert (await SummaryStore(path).load("other", "c")).cursor == 4
async def test_largest_ledger_fits_the_fixed_summary_reserve(tmp_path):
from blueprints.conversation.context import SUMMARY_RESERVE_TOKENS, message_tokens
messages = []
for i in range(100):
messages += [
HumanMessage(content=f"Must {i}: " + "猫😺" * 1000),
AIMessage(content=f"Decided {i}: " + "budget " * 1000),
]
messages.append(HumanMessage(content="current"))
state = await SummaryStore(tmp_path / "summary.db").update("u", "c", messages)
assert message_tokens(state.as_message()) <= SUMMARY_RESERVE_TOKENS
+80
View File
@@ -462,3 +462,83 @@ def test_focused_note_heading_index_uses_absolute_content_lines():
metadata = json.loads(projected.content)["data"]["markdown_compaction"] metadata = json.loads(projected.content)["data"]["markdown_compaction"]
assert metadata["headings"] == [{"line": 1001, "heading": "# Requested section"}] assert metadata["headings"] == [{"line": 1001, "heading": "# Requested section"}]
assert msg.artifact["data"]["markdown"] == "full original" 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 assert_valid_sequence(out):
seen = set()
for m in out:
if isinstance(m, AIMessage):
seen.update(c["id"] for c in m.tool_calls)
if isinstance(m, ToolMessage):
assert m.tool_call_id in seen
def test_summary_placed_before_current_turn_when_turns_dropped():
question = HumanMessage(content="next")
# The current turn has its own tool exchange the summary must not split.
messages = _conversation(12) + [question, call("now"), result("now")]
summary = AIMessage(content="Earlier: keep dining low")
out = trim_history(messages, token_budget=6000, summary_message=summary)
assert cost(out) <= 6000
assert out[-4:] == [summary, question, messages[-2], messages[-1]]
assert out.count(summary) == 1
# The ledger supersedes excerpts of dropped turns; don't send both.
assert not any(
isinstance(m.content, str)
and m.content.startswith("Earlier conversation excerpts")
for m in out
)
assert isinstance(out[0], HumanMessage)
# Without a ledger, the excerpts still cover the dropped turns.
plain = trim_history(messages, token_budget=6000)
assert plain[0].content.startswith("Earlier conversation excerpts")
assert_valid_sequence(out)
def test_changing_summary_leaves_retained_prefix_identical():
messages = _conversation(12) + [HumanMessage(content="next")]
short = AIMessage(content="Earlier: keep dining low")
long = AIMessage(content="Earlier: " + "constraint " * 400)
a = trim_history(messages, token_budget=6000, summary_message=short)
b = trim_history(messages, token_budget=6000, summary_message=long)
assert a[: a.index(short)] == b[: b.index(long)]
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
+31
View File
@@ -46,3 +46,34 @@ async def test_backwards_compatible_no_query_and_empty_memories(monkeypatch):
) )
fetch.return_value = ["Likes tea"] fetch.return_value = ["Likes tea"]
assert "Likes tea" in await memory.build_system_prompt_with_memories("user") 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")
+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 = ( usage = (
{ {
"input_tokens": prompt_tokens, "input_tokens": prompt_tokens,
@@ -71,6 +71,8 @@ def _end(run_id, prompt_tokens=None):
if prompt_tokens if prompt_tokens
else None else None
) )
if usage and cached_tokens is not None:
usage["input_token_details"] = {"cache_read": cached_tokens}
return { return {
"event": "on_chat_model_end", "event": "on_chat_model_end",
"run_id": run_id, "run_id": run_id,
@@ -207,7 +209,7 @@ async def test_logs_per_call_prompt_tokens(owner, monkeypatch, caplog):
script = [ script = [
(0, _start("m1")), (0, _start("m1")),
(0, _chunk("m1", "Hi")), (0, _chunk("m1", "Hi")),
(0, _end("m1", prompt_tokens=27262)), (0, _end("m1", prompt_tokens=27262, cached_tokens=25000)),
(0, _final("Hi")), (0, _final("Hi")),
] ]
@@ -224,6 +226,7 @@ async def test_logs_per_call_prompt_tokens(owner, monkeypatch, caplog):
assert summary["max_prompt_tokens"] == 27262 assert summary["max_prompt_tokens"] == 27262
(call,) = summary["model_calls"] (call,) = summary["model_calls"]
assert call["prompt_tokens"] == 27262 assert call["prompt_tokens"] == 27262
assert call["cached_prompt_tokens"] == 25000
assert call["prefill_ms"] is not None and call["ms"] is not None assert call["prefill_ms"] is not None and call["ms"] is not None
+98 -6
View File
@@ -136,11 +136,7 @@ def test_discovery_recovery_admin_gating_and_isolation(router):
) )
@pytest.mark.asyncio def import_agents():
@pytest.mark.parametrize("primary_fails", [False, True])
async def test_real_graph_binds_discovered_schemas_and_accounts_same_subset(
monkeypatch, primary_fails
):
# Constructors are stubbed before agents import: no tokens, services or network. # Constructors are stubbed before agents import: no tokens, services or network.
with ( with (
patch("langchain_openai.ChatOpenAI"), 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") "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 = [], [], [] bound_names, primary_names, accounted_names = [], [], []
class FakeModel: class FakeModel:
@@ -238,3 +243,90 @@ def test_schema_dict_tools_remain_routable_and_admin_gated():
router = ToolRouter([extra], admin_tools=[extra]) router = ToolRouter([extra], admin_tools=[extra])
assert extra not in router.select([HumanMessage(content="extra")]) assert extra not in router.select([HumanMessage(content="extra")])
assert extra in router.select([HumanMessage(content="extra")], is_admin=True) 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()