465 lines
16 KiB
Python
465 lines
16 KiB
Python
"""Regression tests for complete-turn context budgeting."""
|
|
|
|
import pytest
|
|
from langchain_core.messages import AIMessage, HumanMessage, 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], AIMessage)
|
|
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_prompt_accounts_for_selected_tool_schemas(monkeypatch):
|
|
monkeypatch.setenv("TAVILY_API_KEY", "test-key")
|
|
from blueprints.conversation import agents
|
|
|
|
extra = {
|
|
"type": "function",
|
|
"function": {
|
|
"name": "extra",
|
|
"description": "extra docs",
|
|
"parameters": {"type": "object", "properties": {}},
|
|
},
|
|
}
|
|
seen = {}
|
|
|
|
def build(messages, system, schemas, **kwargs):
|
|
seen["schemas"] = schemas
|
|
return messages
|
|
|
|
monkeypatch.setattr(agents, "build_context", build)
|
|
messages = [HumanMessage(content="hello")]
|
|
assert (
|
|
agents._build_prompt({"messages": messages}, tool_schemas=[extra]) == 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
|
|
|
|
|
|
def test_eight_tool_finance_scenario_preserves_structured_evidence():
|
|
import copy
|
|
import json
|
|
|
|
question = HumanMessage(
|
|
content="Compare September spending with my plan; exclude transfers."
|
|
)
|
|
messages = [question]
|
|
kinds = [
|
|
"ynab.summary",
|
|
"ynab.transactions",
|
|
"ynab.spending_report",
|
|
"ynab.category_budget",
|
|
"search.documents",
|
|
"obsidian.note",
|
|
"ynab.transactions",
|
|
"ynab.spending_report",
|
|
]
|
|
for i, kind in enumerate(kinds):
|
|
data = {
|
|
"start_date": "2026-09-01",
|
|
"end_date": "2026-09-30",
|
|
"filters": {"exclude_transfers": True, "budget_id": "budget-123"},
|
|
"total": -1234.56,
|
|
"count": 30,
|
|
}
|
|
if kind == "obsidian.note":
|
|
data.update(
|
|
path="notes/budget.md", markdown="Keep dining under $300.\n" * 2000
|
|
)
|
|
else:
|
|
key = (
|
|
"results"
|
|
if kind == "search.documents"
|
|
else "transactions"
|
|
if kind == "ynab.transactions"
|
|
else "rows"
|
|
)
|
|
data[key] = [
|
|
{
|
|
"id": f"{i}-{n}",
|
|
"path": f"notes/{n}.md",
|
|
"date": "2026-09-15",
|
|
"amount": -41.152,
|
|
"snippet": "finance evidence " * 200,
|
|
}
|
|
for n in range(30)
|
|
]
|
|
artifact = {
|
|
"kind": kind,
|
|
"data": data,
|
|
"source_ref": {"system": "ynab", "id": "budget-123"},
|
|
}
|
|
messages += [
|
|
call(str(i)),
|
|
ToolMessage(
|
|
content="formatted report " * 5000,
|
|
artifact=artifact,
|
|
tool_call_id=str(i),
|
|
),
|
|
]
|
|
original = copy.deepcopy(messages)
|
|
out = trim_history(messages, token_budget=10000)
|
|
assert cost(out) <= 10000
|
|
assert len(out) == 17
|
|
assert out[0] is question
|
|
assert messages == original
|
|
for old, projected in zip(messages[2::2], out[2::2]):
|
|
assert projected.artifact == old.artifact
|
|
parsed = json.loads(projected.content)
|
|
assert parsed["source_ref"] == old.artifact["source_ref"]
|
|
for key in ("total", "start_date", "end_date", "filters", "count"):
|
|
assert parsed["data"][key] == old.artifact["data"][key]
|
|
if parsed["kind"] == "obsidian.note":
|
|
assert parsed["data"]["path"] == "notes/budget.md"
|
|
assert parsed["data"]["markdown_compaction"]["omitted_characters"] > 0
|
|
else:
|
|
data = parsed["data"]
|
|
key = next(k for k in ("rows", "transactions", "results") if k in data)
|
|
omitted = data[key + "_compaction"]
|
|
assert omitted["retained_rows"] + omitted["omitted_rows"] == 30
|
|
assert {r["id"] for r in data[key]} | {
|
|
r["id"] for r in omitted["omitted_source_refs"]
|
|
} == {f"{old.tool_call_id}-{n}" for n in range(30)}
|
|
|
|
|
|
def test_json_content_without_artifact_stays_parseable():
|
|
import json
|
|
|
|
payload = {
|
|
"total": 10,
|
|
"filters": {"query": "food"},
|
|
"rows": [{"id": str(i), "memo": "x" * 2000} for i in range(20)],
|
|
}
|
|
source = result("a", json.dumps(payload))
|
|
projected = context.compact_tool_result(source, 900)
|
|
parsed = json.loads(projected.content)
|
|
assert parsed["total"] == 10
|
|
assert parsed["rows_compaction"]["omitted_rows"] > 0
|
|
assert json.loads(source.content) == payload
|
|
|
|
|
|
def test_unknown_result_is_marked_excerpt():
|
|
projected = context.compact_tool_result(result("a", "unknown prose " * 1000), 100)
|
|
assert "tool result excerpt" in projected.content
|
|
|
|
|
|
def test_structural_metadata_floor_fails_instead_of_corrupting_json():
|
|
message = ToolMessage(
|
|
content="short", artifact={"source_ref": {"id": "x" * 2000}}, tool_call_id="a"
|
|
)
|
|
with pytest.raises(ValueError, match="required tool metadata"):
|
|
trim_history([HumanMessage(content="q"), call("a"), message], token_budget=400)
|
|
|
|
|
|
def test_summary_role_and_budget_are_enforced():
|
|
from langchain_core.messages import SystemMessage
|
|
|
|
with pytest.raises(ValueError, match="assistant-role"):
|
|
trim_history(
|
|
[HumanMessage(content="q")],
|
|
summary_message=SystemMessage(content="untrusted"),
|
|
)
|
|
question = HumanMessage(content="Current request " * 10)
|
|
assert trim_history(
|
|
[question], token_budget=100, summary_message=AIMessage(content="x" * 1000)
|
|
) == [question]
|
|
with pytest.raises(ValueError, match="without tool calls"):
|
|
trim_history([HumanMessage(content="q")], summary_message=call("a"))
|
|
|
|
|
|
def test_nested_source_refs_and_dates_survive_row_omission():
|
|
import json
|
|
|
|
payload = {
|
|
"rows": [
|
|
{
|
|
"metadata": {
|
|
"source_ids": ["one", "two"],
|
|
"start_date": "2026-09-01",
|
|
"path": "notes/plan.md",
|
|
},
|
|
"snippet": "x" * 1000,
|
|
}
|
|
]
|
|
}
|
|
message = ToolMessage(content="display", artifact=payload, tool_call_id="a")
|
|
parsed = json.loads(context.compact_tool_result(message, 0).content)
|
|
assert (
|
|
parsed["rows_compaction"]["omitted_source_refs"][0]["metadata"]
|
|
== payload["rows"][0]["metadata"]
|
|
)
|
|
|
|
|
|
def test_repeated_json_only_compaction_preserves_all_source_ids_and_counts():
|
|
import json
|
|
|
|
original = {
|
|
"rows": [{"id": f"source-{i}", "memo": "detail " * 1000} for i in range(30)],
|
|
"total": 300,
|
|
}
|
|
message = result("a", json.dumps(original))
|
|
first = context.compact_tool_result(message, 3000)
|
|
second = context.compact_tool_result(first, 1500)
|
|
for projected in (first, second):
|
|
parsed = json.loads(projected.content)
|
|
metadata = parsed["rows_compaction"]
|
|
assert metadata["retained_rows"] == len(parsed["rows"])
|
|
assert metadata["retained_rows"] + metadata["omitted_rows"] == 30
|
|
ids = {r["id"] for r in parsed["rows"]} | {
|
|
r["id"] for r in metadata["omitted_source_refs"]
|
|
}
|
|
assert ids == {f"source-{i}" for i in range(30)}
|
|
out = trim_history(
|
|
[HumanMessage(content="q"), call("a"), message], token_budget=1000
|
|
)
|
|
parsed = json.loads(out[-1].content)
|
|
assert parsed["rows_compaction"]["retained_rows"] == len(parsed["rows"])
|
|
ids = {r["id"] for r in parsed["rows"]} | {
|
|
r["id"] for r in parsed["rows_compaction"]["omitted_source_refs"]
|
|
}
|
|
assert len(ids) == 30
|
|
assert cost(out) <= 1000
|
|
|
|
|
|
def test_focused_note_heading_index_uses_absolute_content_lines():
|
|
import json
|
|
|
|
artifact = {
|
|
"kind": "obsidian.note",
|
|
"data": {"path": "late.md", "markdown": "full original"},
|
|
"model_context": {
|
|
"markdown": "# Requested section\n" + "detail\n" * 2000,
|
|
"read_range": {"start_line": 1001, "end_line": 3001, "total_lines": 4000},
|
|
},
|
|
}
|
|
msg = ToolMessage(content="focused read", artifact=artifact, tool_call_id="a")
|
|
projected = context.compact_tool_result(msg, 1500)
|
|
metadata = json.loads(projected.content)["data"]["markdown_compaction"]
|
|
assert metadata["headings"] == [{"line": 1001, "heading": "# Requested section"}]
|
|
assert msg.artifact["data"]["markdown"] == "full original"
|