Budget conversation context while preserving current turns and tool exchanges
This commit is contained in:
+247
-53
@@ -1,78 +1,272 @@
|
||||
"""Tests for trimming the replayed history (blueprints/conversation/context.py)."""
|
||||
"""Regression tests for complete-turn context budgeting."""
|
||||
|
||||
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
|
||||
import pytest
|
||||
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage
|
||||
|
||||
from blueprints.conversation import context
|
||||
from blueprints.conversation.context import (
|
||||
EARLIER_MESSAGE_MAX_CHARS,
|
||||
HISTORY_WINDOW,
|
||||
IMAGE_PLACEHOLDER,
|
||||
TRUNCATION_NOTE,
|
||||
build_context,
|
||||
message_tokens,
|
||||
trim_history,
|
||||
)
|
||||
|
||||
BIG = "x" * (EARLIER_MESSAGE_MAX_CHARS * 3)
|
||||
|
||||
def call(*ids):
|
||||
return AIMessage(
|
||||
content="",
|
||||
tool_calls=[
|
||||
{"id": i, "name": "search", "args": {"path": "notes/budget.md"}}
|
||||
for i in ids
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
def _call(call_id: str, name: str = "simba_search") -> AIMessage:
|
||||
return AIMessage(content="", tool_calls=[{"id": call_id, "name": name, "args": {}}])
|
||||
def result(i, content="found"):
|
||||
return ToolMessage(content=content, tool_call_id=i)
|
||||
|
||||
|
||||
def _result(call_id: str, text: str = BIG) -> ToolMessage:
|
||||
return ToolMessage(content=text, tool_call_id=call_id, name="simba_search")
|
||||
def cost(messages):
|
||||
return sum(message_tokens(m) for m in messages)
|
||||
|
||||
|
||||
def test_earlier_tool_results_are_capped_current_turn_untouched():
|
||||
messages = [
|
||||
HumanMessage(content="first question"),
|
||||
_call("c1"),
|
||||
_result("c1"),
|
||||
AIMessage(content="first answer"),
|
||||
HumanMessage(content="second question"),
|
||||
_call("c2"),
|
||||
_result("c2"),
|
||||
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)
|
||||
]
|
||||
|
||||
out = trim_history(messages)
|
||||
|
||||
earlier, current = out[2], out[6]
|
||||
assert earlier.content == BIG[:EARLIER_MESSAGE_MAX_CHARS] + TRUNCATION_NOTE
|
||||
assert current.content == BIG
|
||||
# The checkpointed originals are never mutated.
|
||||
assert messages[2].content == BIG
|
||||
assert messages[2].content == "data " * 1000
|
||||
assert out[1].tool_calls == messages[1].tool_calls
|
||||
|
||||
|
||||
def test_earlier_attachments_and_images_are_shrunk():
|
||||
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": BIG}, image]),
|
||||
AIMessage(content="nice photo"),
|
||||
HumanMessage(content=[{"type": "text", "text": "now this one"}, image]),
|
||||
HumanMessage(content=[{"type": "text", "text": "old"}, image]),
|
||||
AIMessage(content="photo"),
|
||||
HumanMessage(content=[{"type": "text", "text": "new"}, image]),
|
||||
]
|
||||
|
||||
out = trim_history(messages)
|
||||
|
||||
earlier = out[0].content
|
||||
assert earlier[0]["text"].endswith(TRUNCATION_NOTE)
|
||||
assert earlier[1] == {"type": "text", "text": IMAGE_PLACEHOLDER}
|
||||
# The current turn keeps its image.
|
||||
assert out[2].content[1] == image
|
||||
assert out[0].content[1] == {"type": "text", "text": IMAGE_PLACEHOLDER}
|
||||
assert out[-1].content[1] == image
|
||||
assert messages[0].content[1] == image
|
||||
|
||||
|
||||
def test_window_never_starts_with_orphaned_tool_results():
|
||||
# Many parallel tool results so the window cut lands between an AIMessage's
|
||||
# tool_calls and their ToolMessages.
|
||||
messages = [HumanMessage(content="q"), _call("c0")]
|
||||
messages += [_result("c0", "r") for _ in range(HISTORY_WINDOW)]
|
||||
messages += [AIMessage(content="answer"), HumanMessage(content="next")]
|
||||
|
||||
out = trim_history(messages)
|
||||
|
||||
assert len(out) <= HISTORY_WINDOW
|
||||
assert not isinstance(out[0], ToolMessage)
|
||||
|
||||
|
||||
def test_short_history_is_returned_as_is():
|
||||
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
|
||||
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
|
||||
|
||||
@@ -589,3 +589,26 @@ def test_search_documents_artifact_empty():
|
||||
)
|
||||
def test_format_date_range(start, end, expected):
|
||||
assert format_date_range(start, end) == expected
|
||||
|
||||
|
||||
async def test_obsidian_note_line_range_retains_full_artifact(obsidian):
|
||||
with open(os.path.join(obsidian.vault_path, "range.md"), "w") as f:
|
||||
f.write("first\nsecond\nthird\nfourth")
|
||||
msg = await call(
|
||||
agents.obsidian_read_note,
|
||||
{"relative_path": "range.md", "start_line": 2, "max_lines": 2},
|
||||
)
|
||||
assert "Content lines 2-3 of 4:\nsecond\nthird\n\n" in msg.content
|
||||
assert "first" not in msg.content
|
||||
assert msg.artifact["data"]["markdown"] == "first\nsecond\nthird\nfourth"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"args", [{"start_line": 0}, {"max_lines": 0}, {"start_line": 99}]
|
||||
)
|
||||
async def test_obsidian_note_invalid_line_range(obsidian, args):
|
||||
with open(os.path.join(obsidian.vault_path, "range.md"), "w") as f:
|
||||
f.write("first\nsecond")
|
||||
msg = await call(agents.obsidian_read_note, {"relative_path": "range.md", **args})
|
||||
assert msg.artifact is None
|
||||
assert "positive integers" in msg.content or "beyond the end" in msg.content
|
||||
|
||||
Reference in New Issue
Block a user