m4
8a0ab17936
Pre-existing uncommitted work (runtime snapshots, message budget middleware) preserved as baseline.
80 lines
2.6 KiB
Python
80 lines
2.6 KiB
Python
"""Tests for message-only context budgeting."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
|
|
|
|
from EvoScientist.middleware.message_budget import (
|
|
ContextBudgetUnsatisfiableError,
|
|
MessageReservePolicy,
|
|
count_message_text_tokens,
|
|
create_message_budget_middleware,
|
|
)
|
|
|
|
|
|
def test_text_counter_excludes_attachment_payloads_and_counts_tool_results():
|
|
messages = [
|
|
HumanMessage(
|
|
content=[
|
|
{"type": "text", "text": "abcd"},
|
|
{"type": "image", "base64": "x" * 100_000},
|
|
{"type": "file", "data": "y" * 100_000},
|
|
]
|
|
),
|
|
ToolMessage(content="wxyz", tool_call_id="tool-1"),
|
|
]
|
|
|
|
assert count_message_text_tokens(messages) == 2
|
|
|
|
|
|
def test_reserve_policy_uses_fixed_overhead_without_request_token_counting():
|
|
policy = MessageReservePolicy()
|
|
|
|
assert policy.hard_budget(32_768) == 17_204
|
|
assert policy.soft_budget(32_768) == 12_042
|
|
assert policy.keep_budget(32_768) == 6_021
|
|
assert policy.hard_budget(32_768, has_tools=True) == 9_012
|
|
|
|
|
|
def test_budget_middleware_uses_message_threshold_and_safe_tool_cutoff():
|
|
model = MagicMock()
|
|
model.profile = {"max_input_tokens": 32_768}
|
|
middleware = create_message_budget_middleware(model, MagicMock())
|
|
messages = [
|
|
HumanMessage(content="a" * 30_000),
|
|
AIMessage(content="", tool_calls=[{"name": "read_file", "args": {}, "id": "tool-1"}]),
|
|
ToolMessage(content="b" * 30_000, tool_call_id="tool-1"),
|
|
HumanMessage(content="c" * 30_000),
|
|
]
|
|
|
|
total = count_message_text_tokens(messages)
|
|
|
|
assert middleware._should_summarize(messages, total) is True
|
|
cutoff = middleware._determine_cutoff_index(messages)
|
|
assert cutoff in {1, 3}
|
|
# A cutoff never leaves the tool response without its matching AI tool call.
|
|
if cutoff == 1:
|
|
assert isinstance(messages[cutoff], AIMessage)
|
|
|
|
|
|
def test_budget_rejects_tools_and_attachments_when_reserves_exceed_model_limit():
|
|
model = MagicMock()
|
|
model.profile = {
|
|
"max_input_tokens": 28_672,
|
|
"min_effective_input_tokens": 4_096,
|
|
}
|
|
middleware = create_message_budget_middleware(model, MagicMock())
|
|
request = SimpleNamespace(
|
|
messages=[
|
|
HumanMessage(content=[{"type": "image", "url": "https://example.test/a"}])
|
|
],
|
|
tools=[{"name": "read_file"}],
|
|
)
|
|
|
|
with pytest.raises(ContextBudgetUnsatisfiableError, match="CONTEXT_BUDGET_UNSATISFIABLE"):
|
|
middleware._budget_for_request(request)
|