Keep the chat prompt prefix stable for llama-server prompt caching #77
@@ -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
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -29,6 +29,10 @@ 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
|
||||
# Upper bound on the summary ledger message (summary.STATE_TOKEN_LIMIT + wrapper).
|
||||
SUMMARY_RESERVE_TOKENS = 2000
|
||||
|
||||
|
||||
def estimate_tokens(value) -> int:
|
||||
@@ -340,46 +344,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 +427,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 +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(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 = [
|
||||
|
||||
@@ -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}
|
||||
|
||||
@@ -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]
|
||||
|
||||
|
||||
@@ -162,3 +163,17 @@ async def test_deletion_is_scoped_and_survives_reopen(tmp_path):
|
||||
await store.delete("owner", "c")
|
||||
assert (await SummaryStore(path).load("owner", "c")).cursor == 0
|
||||
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
|
||||
|
||||
@@ -462,3 +462,83 @@ 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 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
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user