273 lines
9.1 KiB
Python
273 lines
9.1 KiB
Python
"""Regression tests for complete-turn context budgeting."""
|
|
|
|
import pytest
|
|
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
|
|
|
|
from blueprints.conversation import context
|
|
from blueprints.conversation.context import (
|
|
IMAGE_PLACEHOLDER,
|
|
build_context,
|
|
message_tokens,
|
|
trim_history,
|
|
)
|
|
|
|
|
|
def call(*ids):
|
|
return AIMessage(
|
|
content="",
|
|
tool_calls=[
|
|
{"id": i, "name": "search", "args": {"path": "notes/budget.md"}}
|
|
for i in ids
|
|
],
|
|
)
|
|
|
|
|
|
def result(i, content="found"):
|
|
return ToolMessage(content=content, tool_call_id=i)
|
|
|
|
|
|
def cost(messages):
|
|
return sum(message_tokens(m) for m in messages)
|
|
|
|
|
|
def test_long_current_turn_keeps_question_and_all_exchanges():
|
|
question = HumanMessage(content="Compare my spending with my plan")
|
|
messages = [question]
|
|
for i in range(8):
|
|
messages += [call(str(i)), result(str(i), "data " * 1000)]
|
|
out = trim_history(messages, token_budget=3000)
|
|
assert out[0] is question
|
|
assert len(out) == 17
|
|
assert cost(out) <= 3000
|
|
assert [m.tool_call_id for m in out if isinstance(m, ToolMessage)] == [
|
|
str(i) for i in range(8)
|
|
]
|
|
assert messages[2].content == "data " * 1000
|
|
assert out[1].tool_calls == messages[1].tool_calls
|
|
|
|
|
|
def test_parallel_tool_results_remain_paired():
|
|
messages = [
|
|
HumanMessage(content="q"),
|
|
call("a", "b"),
|
|
result("a", "x" * 20000),
|
|
result("b", "y" * 20000),
|
|
]
|
|
out = trim_history(messages, token_budget=700)
|
|
assert len(out) == 4
|
|
assert out[1].tool_calls == messages[1].tool_calls
|
|
assert cost(out) <= 700
|
|
|
|
|
|
def test_drops_whole_turns_and_includes_bounded_excerpts():
|
|
messages = []
|
|
for i in range(12):
|
|
messages += [
|
|
HumanMessage(content=f"Question {i}"),
|
|
call(str(i)),
|
|
result(str(i), "data " * 1000),
|
|
AIMessage(content=f"Decision {i}"),
|
|
]
|
|
question = HumanMessage(content="What did we decide?")
|
|
messages.append(question)
|
|
out = trim_history(messages, token_budget=3000)
|
|
assert isinstance(out[0], SystemMessage)
|
|
assert "Decision" in out[0].content
|
|
assert "historical data" in out[0].content
|
|
assert cost(out) <= 3000
|
|
assert out[-1] is question
|
|
# Every retained result still has its original call preceding it.
|
|
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_earlier_images_replaced_current_image_preserved():
|
|
image = {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}}
|
|
messages = [
|
|
HumanMessage(content=[{"type": "text", "text": "old"}, image]),
|
|
AIMessage(content="photo"),
|
|
HumanMessage(content=[{"type": "text", "text": "new"}, image]),
|
|
]
|
|
out = trim_history(messages)
|
|
assert out[0].content[1] == {"type": "text", "text": IMAGE_PLACEHOLDER}
|
|
assert out[-1].content[1] == image
|
|
assert messages[0].content[1] == image
|
|
|
|
|
|
def test_short_history_unchanged():
|
|
messages = [HumanMessage(content="hi"), AIMessage(content="hello")]
|
|
assert trim_history(messages) == messages
|
|
|
|
|
|
def test_question_that_cannot_fit_fails_without_truncation():
|
|
messages = [HumanMessage(content="important request " * 1000)]
|
|
with pytest.raises(ValueError, match="current question"):
|
|
trim_history(messages, token_budget=100)
|
|
assert messages[0].content == "important request " * 1000
|
|
|
|
|
|
def test_system_and_tool_schemas_are_charged_to_budget(monkeypatch):
|
|
monkeypatch.setattr(context, "PROMPT_BUDGET_TOKENS", 2500)
|
|
monkeypatch.setattr(context, "CONTEXT_WINDOW_TOKENS", 10000)
|
|
schemas = [{"name": "search", "description": "tool docs " * 100}]
|
|
messages = [HumanMessage(content="q"), call("a"), result("a", "result " * 3000)]
|
|
out = build_context(messages, "system " * 100, schemas)
|
|
assert cost(out) + context.estimate_tokens(schemas) <= 2500
|
|
assert out[1] is messages[0]
|
|
|
|
|
|
def test_fixed_overhead_that_cannot_fit_fails(monkeypatch):
|
|
monkeypatch.setattr(context, "PROMPT_BUDGET_TOKENS", 100)
|
|
with pytest.raises(ValueError, match="tool definitions"):
|
|
build_context([HumanMessage(content="q")], "system" * 1000, [])
|
|
|
|
|
|
def test_response_and_safety_reserves_reduce_prompt_limit(monkeypatch):
|
|
monkeypatch.setattr(context, "CONTEXT_WINDOW_TOKENS", 8000)
|
|
monkeypatch.setattr(context, "PROMPT_BUDGET_TOKENS", 16000)
|
|
messages = [HumanMessage(content="q"), call("a"), result("a", "data " * 5000)]
|
|
out = build_context(messages, "system", [])
|
|
assert (
|
|
cost(out) + context.estimate_tokens([])
|
|
<= 8000 - context.RESPONSE_RESERVE_TOKENS - context.SAFETY_RESERVE_TOKENS
|
|
)
|
|
|
|
|
|
def test_cancelled_turn_does_not_replay_incomplete_exchange():
|
|
messages = [
|
|
HumanMessage(content="old"),
|
|
call("a", "b"),
|
|
result("a"),
|
|
HumanMessage(content="new"),
|
|
]
|
|
out = trim_history(messages)
|
|
assert out == [messages[0], messages[-1]]
|
|
|
|
|
|
def test_orphan_results_are_removed():
|
|
messages = [
|
|
result("orphan"),
|
|
HumanMessage(content="q"),
|
|
result("another orphan"),
|
|
AIMessage(content="answer"),
|
|
]
|
|
assert trim_history(messages) == [messages[1], messages[3]]
|
|
|
|
|
|
def test_large_assistant_answer_is_bounded_in_earlier_turn():
|
|
messages = [
|
|
HumanMessage(content="old"),
|
|
AIMessage(content="long " * 10000),
|
|
HumanMessage(content="new"),
|
|
]
|
|
out = trim_history(messages, token_budget=4000)
|
|
assert cost(out) <= 4000
|
|
assert len(out[1].content) < len(messages[1].content)
|
|
assert messages[1].content == "long " * 10000
|
|
|
|
|
|
def test_multiblock_attachment_shares_one_cap():
|
|
messages = [
|
|
HumanMessage(content=[{"type": "text", "text": "x" * 10000} for _ in range(5)]),
|
|
AIMessage(content="ok"),
|
|
HumanMessage(content="next"),
|
|
]
|
|
out = trim_history(messages, token_budget=5000)
|
|
assert cost(out) <= 5000
|
|
assert sum(len(p["text"]) for p in out[0].content) < 5000
|
|
|
|
|
|
async def test_saved_memories_are_bounded_without_deleting(monkeypatch):
|
|
from blueprints.conversation import memory
|
|
|
|
saved = ["newest preference", "big " * 5000, "older preference"]
|
|
|
|
async def get_memories(user_id):
|
|
return saved
|
|
|
|
monkeypatch.setattr(memory, "get_memories_for_user", get_memories)
|
|
prompt = await memory.build_system_prompt_with_memories("user")
|
|
block = prompt.split("USER MEMORIES", 1)[1]
|
|
assert "newest preference" in block
|
|
assert "older preference" in block
|
|
assert "omitted" in block
|
|
assert "big big" not in block
|
|
assert saved[1] == "big " * 5000
|
|
|
|
|
|
def test_agent_prompt_accounts_for_extra_tool_schemas(monkeypatch):
|
|
monkeypatch.setenv("TAVILY_API_KEY", "test-key")
|
|
from blueprints.conversation import agents
|
|
|
|
captured = {}
|
|
|
|
def create(**kwargs):
|
|
captured.update(kwargs)
|
|
|
|
monkeypatch.setattr(agents, "create_react_agent", create)
|
|
extra = {
|
|
"type": "function",
|
|
"function": {
|
|
"name": "extra",
|
|
"description": "extra docs",
|
|
"parameters": {"type": "object", "properties": {}},
|
|
},
|
|
}
|
|
agents.build_agent(extra_tools=[extra])
|
|
assert extra in captured["tools"]
|
|
question = HumanMessage(content="hello")
|
|
actual = captured["prompt"].invoke({"messages": [question]})
|
|
assert actual[-1] is question
|
|
# The actual prompt closure includes this schema rather than relying on a
|
|
# static reserve that would miss admin-only tools.
|
|
seen = {}
|
|
|
|
def build(messages, system, schemas):
|
|
seen["schemas"] = schemas
|
|
return messages
|
|
|
|
monkeypatch.setattr(agents, "build_context", build)
|
|
messages = [HumanMessage(content="hello")]
|
|
assert captured["prompt"].invoke({"messages": messages}) == messages
|
|
assert extra in seen["schemas"]
|
|
|
|
|
|
def test_heavy_current_turn_leaves_room_for_previous_decisions():
|
|
question = HumanMessage(content="Apply that plan to my spending")
|
|
messages = [
|
|
HumanMessage(content="Plan?"),
|
|
AIMessage(content="Keep dining under $300."),
|
|
question,
|
|
call("a"),
|
|
result("a", "spending " * 10000),
|
|
]
|
|
out = trim_history(messages, token_budget=2500)
|
|
assert cost(out) <= 2500
|
|
assert question in out
|
|
assert any(
|
|
"Keep dining under $300." in m.content
|
|
for m in out
|
|
if isinstance(m.content, str)
|
|
)
|
|
|
|
|
|
def test_large_tool_arguments_fail_explicitly_instead_of_breaking_pair():
|
|
exchange = call("a")
|
|
exchange.tool_calls[0]["args"] = {"query": "huge " * 10000}
|
|
messages = [HumanMessage(content="q"), exchange, result("a")]
|
|
with pytest.raises(ValueError, match="tool calls exceed"):
|
|
trim_history(messages, token_budget=300)
|
|
assert exchange.tool_calls[0]["args"]["query"] == "huge " * 10000
|
|
|
|
|
|
def test_unicode_tool_result_fits_estimated_budget():
|
|
messages = [HumanMessage(content="q"), call("a"), result("a", "猫😺预算" * 10000)]
|
|
out = trim_history(messages, token_budget=600)
|
|
assert cost(out) <= 600
|
|
assert out[0] is messages[0]
|
|
assert out[1].tool_calls == messages[1].tool_calls
|