m4
aae8d0a379
In-progress work committed to unblock the config import/export plan: - prompts: FILE_REFERENCES section for workspace-relative file citation - backends: resolve quoted virtual absolute paths onto the sandbox workspace - middleware: read_file_images middleware; message_budget extensions - image_gen/model_registry: image model 'enabled' flag refactor - memory/launch, gateway/background_runs, tools/image follow-ons - scripts: dev_backend.sh, release.sh - tests for the above
608 lines
23 KiB
Python
608 lines
23 KiB
Python
"""Tests for message-only context budgeting (design doc 6.5, 8.1, 8.3).
|
|
|
|
Every budget comes from the frozen ``ResolvedModelConfig.budget`` of the
|
|
run snapshot — the middleware's own ``snapshot_role`` selects which frozen
|
|
configuration sizes the budget (section 6.1 role mapping). Runs without an
|
|
explicit ``runtime_snapshot_id`` lazily create a local snapshot bound to
|
|
their thread (section 8.1); a bootstrap registry fails closed with
|
|
``MODEL_REGISTRY_NOT_READY``. There is no interim 32K/profile fallback.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from contextlib import contextmanager
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
from langchain_core.messages import AIMessage, HumanMessage, ToolMessage
|
|
from langchain_openai import ChatOpenAI
|
|
|
|
from EvoScientist.middleware.message_budget import (
|
|
_snapshot_message_budget,
|
|
count_message_text_tokens,
|
|
create_message_budget_middleware,
|
|
)
|
|
from EvoScientist.model_registry.errors import (
|
|
MODEL_CONFIG_OUTSIDE_SNAPSHOT,
|
|
MODEL_REGISTRY_NOT_READY,
|
|
SNAPSHOT_NOT_FOUND,
|
|
ModelRegistryError,
|
|
)
|
|
from EvoScientist.model_registry.runtime import SnapshotRuntime
|
|
from EvoScientist.model_registry.store import ModelRuntimeStore
|
|
from tests.registry_fixtures import make_active_store, make_snapshot
|
|
|
|
|
|
@contextmanager
|
|
def _patched_config(configurable: dict | None):
|
|
"""Patch ``langgraph.config.get_config`` to expose ``configurable``."""
|
|
import langgraph.config as _lg_cfg
|
|
|
|
if configurable is None:
|
|
with patch.object(
|
|
_lg_cfg,
|
|
"get_config",
|
|
side_effect=RuntimeError("Called get_config outside of a runnable context"),
|
|
):
|
|
yield
|
|
else:
|
|
with patch.object(
|
|
_lg_cfg,
|
|
"get_config",
|
|
return_value={"configurable": configurable},
|
|
):
|
|
yield
|
|
|
|
|
|
@pytest.fixture
|
|
def store(tmp_path):
|
|
return make_active_store(tmp_path / "model-runtime")
|
|
|
|
|
|
@pytest.fixture
|
|
def runtime(store):
|
|
return SnapshotRuntime(store)
|
|
|
|
|
|
def _request(messages, tools=()):
|
|
return SimpleNamespace(messages=messages, tools=list(tools))
|
|
|
|
|
|
def _configurable_for(snapshot, **overrides):
|
|
configurable = {
|
|
"runtime_snapshot_id": snapshot.snapshot_id,
|
|
"thread_id": snapshot.thread_id,
|
|
}
|
|
configurable.update(overrides)
|
|
return configurable
|
|
|
|
|
|
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
|
|
|
|
|
|
# =============================================================================
|
|
# Snapshot-driven budget recomputation (section 6.5)
|
|
# =============================================================================
|
|
|
|
|
|
class TestSnapshotMessageBudget:
|
|
"""The frozen fixture registry resolves to:
|
|
|
|
resolved_input_limit = 1048576 - 32768 = 1015808
|
|
reserves: system 4096, tools 8192, attachments 4096
|
|
"""
|
|
|
|
def test_base_mode_deducts_only_system_reserve(self, store):
|
|
snapshot = make_snapshot(store)
|
|
budget = _snapshot_message_budget(
|
|
snapshot, "primary", has_tools=False, has_attachments=False
|
|
)
|
|
message_budget = 1015808 - 4096
|
|
assert budget.input_limit == 1015808
|
|
assert budget.hard_tokens == int(message_budget * 0.90)
|
|
assert budget.soft_tokens == int(budget.hard_tokens * 0.70)
|
|
assert budget.keep_tokens == int(budget.hard_tokens * 0.35)
|
|
assert budget.reserved_tokens == 4096
|
|
|
|
def test_tools_and_attachments_deduct_their_reserves(self, store):
|
|
snapshot = make_snapshot(store)
|
|
tools_only = _snapshot_message_budget(
|
|
snapshot, "primary", has_tools=True, has_attachments=False
|
|
)
|
|
full = _snapshot_message_budget(
|
|
snapshot, "primary", has_tools=True, has_attachments=True
|
|
)
|
|
assert tools_only.hard_tokens == int((1015808 - 4096 - 8192) * 0.90)
|
|
assert full.hard_tokens == int((1015808 - 4096 - 8192 - 4096) * 0.90)
|
|
assert tools_only.reserved_tokens == 4096 + 8192
|
|
assert full.reserved_tokens == 4096 + 8192 + 4096
|
|
|
|
class TestSnapshotModeMiddleware:
|
|
def test_budget_uses_frozen_reserves_not_model_profile(self, store, runtime):
|
|
"""A misleading compile-time profile must not affect snapshot mode."""
|
|
model = MagicMock()
|
|
model.profile = {"max_input_tokens": 32_768}
|
|
middleware = create_message_budget_middleware(
|
|
model, MagicMock(), runtime=runtime
|
|
)
|
|
snapshot = make_snapshot(store)
|
|
|
|
with _patched_config(_configurable_for(snapshot)):
|
|
budget = middleware._budget_for_request(_request([HumanMessage("hi")]))
|
|
|
|
assert budget.input_limit == 1015808
|
|
assert budget.has_tools is True # conservative construction-time flag
|
|
assert budget.has_attachments is False
|
|
assert budget.hard_tokens == int((1015808 - 4096 - 8192) * 0.90)
|
|
|
|
def test_recomputed_per_call_with_current_attachments(self, store, runtime):
|
|
middleware = create_message_budget_middleware(
|
|
MagicMock(), MagicMock(), runtime=runtime
|
|
)
|
|
snapshot = make_snapshot(store)
|
|
plain = _request([HumanMessage(content="plain")])
|
|
attached = _request(
|
|
[HumanMessage(content=[{"type": "image", "url": "https://example.test/a"}])]
|
|
)
|
|
|
|
with _patched_config(_configurable_for(snapshot)):
|
|
first = middleware._budget_for_request(plain)
|
|
second = middleware._budget_for_request(attached)
|
|
third = middleware._budget_for_request(plain)
|
|
|
|
assert first.has_attachments is False
|
|
assert second.has_attachments is True
|
|
assert second.hard_tokens == int((1015808 - 4096 - 8192 - 4096) * 0.90)
|
|
# No caching across calls: the third call recomputes from scratch.
|
|
assert third == first
|
|
|
|
def test_has_tools_is_conservative_not_request_derived(self, store, runtime):
|
|
"""Tool mode comes from construction, not the request's bound tools."""
|
|
middleware = create_message_budget_middleware(
|
|
MagicMock(), MagicMock(), has_tools=True, runtime=runtime
|
|
)
|
|
snapshot = make_snapshot(store)
|
|
|
|
with _patched_config(_configurable_for(snapshot)):
|
|
budget = middleware._budget_for_request(_request([HumanMessage("hi")]))
|
|
|
|
# The request binds no tools at all, yet the tools reserve applies.
|
|
assert budget.has_tools is True
|
|
assert budget.hard_tokens == int((1015808 - 4096 - 8192) * 0.90)
|
|
|
|
def test_has_tools_false_when_agent_has_no_toolset(self, store, runtime):
|
|
middleware = create_message_budget_middleware(
|
|
MagicMock(), MagicMock(), has_tools=False, runtime=runtime
|
|
)
|
|
snapshot = make_snapshot(store)
|
|
|
|
with _patched_config(_configurable_for(snapshot)):
|
|
budget = middleware._budget_for_request(
|
|
_request([HumanMessage("hi")], tools=[{"name": "read_file"}])
|
|
)
|
|
|
|
assert budget.has_tools is False
|
|
assert budget.hard_tokens == int((1015808 - 4096) * 0.90)
|
|
|
|
def test_unknown_snapshot_role_rejected(self, runtime):
|
|
with pytest.raises(ValueError, match="Unknown model role"):
|
|
create_message_budget_middleware(
|
|
MagicMock(), MagicMock(), snapshot_role="bogus", runtime=runtime
|
|
)
|
|
|
|
def test_snapshot_binding_is_verified(self, store, runtime):
|
|
middleware = create_message_budget_middleware(
|
|
MagicMock(), MagicMock(), runtime=runtime
|
|
)
|
|
snapshot = make_snapshot(store)
|
|
|
|
with (
|
|
_patched_config(_configurable_for(snapshot, thread_id="other")),
|
|
pytest.raises(ModelRegistryError) as excinfo,
|
|
):
|
|
middleware._budget_for_request(_request([HumanMessage("hi")]))
|
|
assert excinfo.value.code == SNAPSHOT_NOT_FOUND
|
|
|
|
def test_summarizer_model_uses_snapshot_primary(self, store, runtime):
|
|
middleware = create_message_budget_middleware(
|
|
MagicMock(), MagicMock(), runtime=runtime
|
|
)
|
|
snapshot = make_snapshot(store)
|
|
|
|
with _patched_config(_configurable_for(snapshot)):
|
|
model = middleware.model
|
|
|
|
# Every role maps to the snapshot's frozen primary (section 6.1).
|
|
assert isinstance(model, ChatOpenAI)
|
|
assert model.model_name == "glm-5.2"
|
|
|
|
def test_summarizer_model_cached_per_snapshot(self, store, runtime):
|
|
middleware = create_message_budget_middleware(
|
|
MagicMock(), MagicMock(), runtime=runtime
|
|
)
|
|
snapshot = make_snapshot(store)
|
|
|
|
with _patched_config(_configurable_for(snapshot)):
|
|
assert middleware.model is middleware.model
|
|
|
|
|
|
# =============================================================================
|
|
# Local snapshot entry (section 8.1): no explicit snapshot ID
|
|
# =============================================================================
|
|
|
|
|
|
class TestLocalSnapshotEntry:
|
|
def test_missing_snapshot_lazily_creates_one_bound_to_the_thread(
|
|
self, store, runtime
|
|
):
|
|
middleware = create_message_budget_middleware(
|
|
MagicMock(), MagicMock(), has_tools=False, runtime=runtime
|
|
)
|
|
|
|
with _patched_config({"thread_id": "cron-thread-1"}):
|
|
budget = middleware._budget_for_request(_request([HumanMessage("hi")]))
|
|
|
|
assert budget.input_limit == 1015808
|
|
# The lazy snapshot is persisted, bound to the run's own thread.
|
|
row = store.find_active_run_snapshot(
|
|
deployment_id="local",
|
|
thread_id="cron-thread-1",
|
|
run_request_id="auto:cron-thread-1",
|
|
)
|
|
assert row is not None
|
|
|
|
def test_lazy_creation_is_shared_across_middleware_calls(self, store, runtime):
|
|
"""The same thread converges on one snapshot (idempotent create)."""
|
|
middleware = create_message_budget_middleware(
|
|
MagicMock(), MagicMock(), runtime=runtime
|
|
)
|
|
|
|
with _patched_config({"thread_id": "cron-thread-2"}):
|
|
first = middleware._snapshot()
|
|
middleware._budget_for_request(_request([HumanMessage("hi")]))
|
|
second = middleware._snapshot()
|
|
|
|
assert first.snapshot_id == second.snapshot_id
|
|
|
|
def test_missing_snapshot_and_thread_id_fails_closed(self, runtime):
|
|
middleware = create_message_budget_middleware(
|
|
MagicMock(), MagicMock(), runtime=runtime
|
|
)
|
|
|
|
with (
|
|
_patched_config({}),
|
|
pytest.raises(ModelRegistryError) as excinfo,
|
|
):
|
|
middleware._budget_for_request(_request([HumanMessage("hi")]))
|
|
assert excinfo.value.code == MODEL_CONFIG_OUTSIDE_SNAPSHOT
|
|
|
|
def test_bootstrap_registry_fails_with_not_ready(self, tmp_path):
|
|
"""No 32K fallback: bootstrap registry → MODEL_REGISTRY_NOT_READY."""
|
|
bootstrap_runtime = SnapshotRuntime(
|
|
ModelRuntimeStore(config_dir=tmp_path / "model-runtime")
|
|
)
|
|
middleware = create_message_budget_middleware(
|
|
MagicMock(), MagicMock(), runtime=bootstrap_runtime
|
|
)
|
|
|
|
with (
|
|
_patched_config({"thread_id": "cli-thread"}),
|
|
pytest.raises(ModelRegistryError) as excinfo,
|
|
):
|
|
middleware._budget_for_request(_request([HumanMessage("hi")]))
|
|
assert excinfo.value.code == MODEL_REGISTRY_NOT_READY
|
|
|
|
|
|
# =============================================================================
|
|
# Cutoff behavior (budget-driven summarization trigger)
|
|
# =============================================================================
|
|
|
|
|
|
def test_budget_middleware_uses_message_threshold_and_safe_tool_cutoff(tmp_path):
|
|
# A small context window keeps the trigger thresholds reachable in a test.
|
|
store = make_active_store(
|
|
tmp_path / "model-runtime-small",
|
|
context_window_tokens=32768,
|
|
max_output_tokens=8192,
|
|
min_effective_input_tokens=1024,
|
|
fixed_system_reserve_tokens=1024,
|
|
fixed_tools_reserve_tokens=2048,
|
|
fixed_attachments_reserve_tokens=1024,
|
|
)
|
|
runtime = SnapshotRuntime(store)
|
|
middleware = create_message_budget_middleware(
|
|
MagicMock(), MagicMock(), has_tools=False, runtime=runtime
|
|
)
|
|
snapshot = make_snapshot(store)
|
|
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)
|
|
|
|
from EvoScientist.middleware.message_budget import _ACTIVE_BUDGET
|
|
|
|
with _patched_config(_configurable_for(snapshot)):
|
|
# Simulate the per-call budget activation that wrap_model_call sets.
|
|
token = _ACTIVE_BUDGET.set(middleware._budget_for_request(_request(messages)))
|
|
try:
|
|
assert middleware._should_summarize(messages, total) is True
|
|
cutoff = middleware._determine_cutoff_index(messages)
|
|
finally:
|
|
_ACTIVE_BUDGET.reset(token)
|
|
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_active_budget_outside_model_call_fails_loudly(runtime):
|
|
middleware = create_message_budget_middleware(
|
|
MagicMock(), MagicMock(), runtime=runtime
|
|
)
|
|
with pytest.raises(RuntimeError, match="outside a model call"):
|
|
middleware._active_budget()
|
|
|
|
|
|
# =============================================================================
|
|
# Compression token stats on the summarization event
|
|
# =============================================================================
|
|
|
|
|
|
class TestCompressionTokenStats:
|
|
"""A freshly created ``_summarization_event`` records the estimated
|
|
prompt size before and after compaction plus the active budget, so the
|
|
WebUI can show what compression actually did (character estimates — the
|
|
provider-measured sizes keep flowing through the usage pipeline)."""
|
|
|
|
def _event_response(self):
|
|
event = {
|
|
"cutoff_index": 1,
|
|
"summary_message": HumanMessage(content="summary"),
|
|
"file_path": None,
|
|
}
|
|
return SimpleNamespace(
|
|
command=SimpleNamespace(update={"_summarization_event": event})
|
|
)
|
|
|
|
@staticmethod
|
|
def _fresh_event_response():
|
|
return TestCompressionTokenStats._event_response(None)
|
|
|
|
def test_augment_attaches_estimates_and_budget(self, store, runtime):
|
|
from EvoScientist.middleware.message_budget import (
|
|
_LAST_TOTAL_TOKENS,
|
|
_augment_summarization_event,
|
|
)
|
|
|
|
middleware = create_message_budget_middleware(
|
|
MagicMock(), MagicMock(), has_tools=False, runtime=runtime
|
|
)
|
|
snapshot = make_snapshot(store)
|
|
with _patched_config(_configurable_for(snapshot)):
|
|
budget = middleware._budget_for_request(_request([HumanMessage("hi")]))
|
|
|
|
token = _LAST_TOTAL_TOKENS.set(5_000)
|
|
try:
|
|
response = self._event_response()
|
|
_augment_summarization_event(response, budget, [1_250])
|
|
finally:
|
|
_LAST_TOTAL_TOKENS.reset(token)
|
|
|
|
event = response.command.update["_summarization_event"]
|
|
assert event["estimated_tokens_before"] == 5_000
|
|
assert event["estimated_tokens_after"] == 1_250
|
|
assert event["budget"] == {
|
|
"hard_tokens": budget.hard_tokens,
|
|
"soft_tokens": budget.soft_tokens,
|
|
"keep_tokens": budget.keep_tokens,
|
|
}
|
|
|
|
def test_augment_ignores_responses_without_a_new_event(self, store, runtime):
|
|
from EvoScientist.middleware.message_budget import (
|
|
_augment_summarization_event,
|
|
)
|
|
|
|
middleware = create_message_budget_middleware(
|
|
MagicMock(), MagicMock(), has_tools=False, runtime=runtime
|
|
)
|
|
snapshot = make_snapshot(store)
|
|
with _patched_config(_configurable_for(snapshot)):
|
|
budget = middleware._budget_for_request(_request([HumanMessage("hi")]))
|
|
|
|
plain = SimpleNamespace()
|
|
_augment_summarization_event(plain, budget, [10])
|
|
assert not hasattr(plain, "command")
|
|
|
|
no_event = SimpleNamespace(command=SimpleNamespace(update={"messages": []}))
|
|
_augment_summarization_event(no_event, budget, [10])
|
|
assert "_summarization_event" not in no_event.command.update
|
|
|
|
def test_should_summarize_stashes_total_tokens(self, store, runtime):
|
|
from EvoScientist.middleware.message_budget import (
|
|
_ACTIVE_BUDGET,
|
|
_LAST_TOTAL_TOKENS,
|
|
)
|
|
|
|
middleware = create_message_budget_middleware(
|
|
MagicMock(), MagicMock(), has_tools=False, runtime=runtime
|
|
)
|
|
snapshot = make_snapshot(store)
|
|
messages = [HumanMessage(content="a" * 400)]
|
|
|
|
with _patched_config(_configurable_for(snapshot)):
|
|
token = _ACTIVE_BUDGET.set(
|
|
middleware._budget_for_request(_request(messages))
|
|
)
|
|
try:
|
|
middleware._should_summarize(messages, 123_456)
|
|
assert _LAST_TOTAL_TOKENS.get() == 123_456
|
|
finally:
|
|
_ACTIVE_BUDGET.reset(token)
|
|
|
|
def test_wrap_model_call_counts_final_prompt_and_augments_event(
|
|
self, store, runtime
|
|
):
|
|
from deepagents.middleware.summarization import SummarizationMiddleware
|
|
|
|
middleware = create_message_budget_middleware(
|
|
MagicMock(), MagicMock(), has_tools=False, runtime=runtime
|
|
)
|
|
snapshot = make_snapshot(store)
|
|
compressed = [
|
|
HumanMessage(content="summary"),
|
|
HumanMessage(content="kept tail"),
|
|
]
|
|
response_holder = {}
|
|
|
|
def fake_base(_self, _request, handler):
|
|
# The base class's last handler call carries the exact prompt sent
|
|
# to the model after compaction.
|
|
modified = SimpleNamespace(messages=compressed, tools=[])
|
|
response_holder["handler_result"] = handler(modified)
|
|
return TestCompressionTokenStats._fresh_event_response()
|
|
|
|
with (
|
|
_patched_config(_configurable_for(snapshot)),
|
|
patch.object(SummarizationMiddleware, "wrap_model_call", fake_base),
|
|
):
|
|
response = middleware.wrap_model_call(
|
|
SimpleNamespace(messages=[HumanMessage("x" * 400)], tools=[]),
|
|
MagicMock(return_value="model-response"),
|
|
)
|
|
|
|
assert response_holder["handler_result"] == "model-response"
|
|
event = response.command.update["_summarization_event"]
|
|
assert event["estimated_tokens_after"] == count_message_text_tokens(
|
|
compressed
|
|
)
|
|
assert event["budget"]["hard_tokens"] > 0
|
|
|
|
async def test_awrap_model_call_counts_final_prompt_and_augments_event(
|
|
self, store, runtime
|
|
):
|
|
from deepagents.middleware.summarization import SummarizationMiddleware
|
|
|
|
middleware = create_message_budget_middleware(
|
|
MagicMock(), MagicMock(), has_tools=False, runtime=runtime
|
|
)
|
|
snapshot = make_snapshot(store)
|
|
compressed = [HumanMessage(content="summary")]
|
|
|
|
async def fake_base(_self, _request, handler):
|
|
modified = SimpleNamespace(messages=compressed, tools=[])
|
|
await handler(modified)
|
|
return TestCompressionTokenStats._fresh_event_response()
|
|
|
|
async def handler(_request):
|
|
return "model-response"
|
|
|
|
with (
|
|
_patched_config(_configurable_for(snapshot)),
|
|
patch.object(SummarizationMiddleware, "awrap_model_call", fake_base),
|
|
):
|
|
response = await middleware.awrap_model_call(
|
|
SimpleNamespace(messages=[HumanMessage("x" * 400)], tools=[]),
|
|
handler,
|
|
)
|
|
|
|
event = response.command.update["_summarization_event"]
|
|
assert event["estimated_tokens_after"] == count_message_text_tokens(
|
|
compressed
|
|
)
|
|
|
|
|
|
class TestContextUsageEmission:
|
|
"""Every model call streams its prompt size on the ``custom`` stream
|
|
mode so the UI can render a live context-occupancy indicator."""
|
|
|
|
def _budget(self, middleware, snapshot):
|
|
with _patched_config(_configurable_for(snapshot)):
|
|
return middleware._budget_for_request(_request([HumanMessage("hi")]))
|
|
|
|
def test_emit_writes_payload_with_writer(self, store, runtime):
|
|
from EvoScientist.middleware.message_budget import _emit_context_usage
|
|
|
|
middleware = create_message_budget_middleware(
|
|
MagicMock(), MagicMock(), has_tools=False, runtime=runtime
|
|
)
|
|
budget = self._budget(middleware, make_snapshot(store))
|
|
written = []
|
|
with patch(
|
|
"langgraph.config.get_stream_writer",
|
|
lambda: written.append,
|
|
):
|
|
_emit_context_usage(budget, 123)
|
|
|
|
assert written == [
|
|
{
|
|
"type": "evoscientist_context_usage",
|
|
"used_tokens": 123,
|
|
"reserved_tokens": budget.reserved_tokens,
|
|
"input_limit": budget.input_limit,
|
|
"hard_tokens": budget.hard_tokens,
|
|
"soft_tokens": budget.soft_tokens,
|
|
}
|
|
]
|
|
|
|
def test_emit_skips_outside_runnable_context(self, store, runtime):
|
|
from EvoScientist.middleware.message_budget import _emit_context_usage
|
|
|
|
middleware = create_message_budget_middleware(
|
|
MagicMock(), MagicMock(), has_tools=False, runtime=runtime
|
|
)
|
|
budget = self._budget(middleware, make_snapshot(store))
|
|
|
|
with patch(
|
|
"langgraph.config.get_stream_writer",
|
|
side_effect=RuntimeError("no runnable context"),
|
|
):
|
|
_emit_context_usage(budget, 123) # must not raise
|
|
|
|
def test_wrap_model_call_emits_final_prompt_size(self, store, runtime):
|
|
from deepagents.middleware.summarization import SummarizationMiddleware
|
|
|
|
middleware = create_message_budget_middleware(
|
|
MagicMock(), MagicMock(), has_tools=False, runtime=runtime
|
|
)
|
|
snapshot = make_snapshot(store)
|
|
sent = [HumanMessage(content="final prompt")]
|
|
written = []
|
|
|
|
def fake_base(_self, _request, handler):
|
|
handler(SimpleNamespace(messages=sent, tools=[]))
|
|
return TestCompressionTokenStats._fresh_event_response()
|
|
|
|
with (
|
|
_patched_config(_configurable_for(snapshot)),
|
|
patch.object(SummarizationMiddleware, "wrap_model_call", fake_base),
|
|
patch("langgraph.config.get_stream_writer", lambda: written.append),
|
|
):
|
|
middleware.wrap_model_call(
|
|
SimpleNamespace(messages=[HumanMessage("x" * 400)], tools=[]),
|
|
MagicMock(return_value="model-response"),
|
|
)
|
|
|
|
assert [event["used_tokens"] for event in written] == [
|
|
count_message_text_tokens(sent)
|
|
]
|
|
assert all(event["type"] == "evoscientist_context_usage" for event in written)
|