wehub-resource-sync
1574 lines
56 KiB
Python
1574 lines
56 KiB
Python
"""Unit tests for omnigent.runtime.compaction."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from omnigent.entities import (
|
|
CompactionData,
|
|
ConversationItem,
|
|
FunctionCallData,
|
|
FunctionCallOutputData,
|
|
MessageData,
|
|
)
|
|
from omnigent.llms.context_window import resolve_effective_context_window
|
|
from omnigent.llms.errors import RetryableLLMError
|
|
from omnigent.llms.types import MessageOutput, OutputText, Response
|
|
from omnigent.runtime.compaction import (
|
|
_BINARY_CONTENT_CLEARED,
|
|
_TOOL_RESULT_CLEARED,
|
|
_is_summary_auth_error,
|
|
_pair_aware_drop_count,
|
|
_truncate_oldest,
|
|
compact,
|
|
compaction_to_history_items,
|
|
count_tokens,
|
|
summarize_history,
|
|
)
|
|
from omnigent.runtime.workflow import _route_bare_model_for_compaction
|
|
from omnigent.spec.types import CompactionConfig, LLMConfig
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# LLM client stubs
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _RaisesIfCalled:
|
|
"""
|
|
LLM client stub that fails the test if ``responses.create()`` is ever
|
|
called.
|
|
|
|
Use this for ``compact()`` calls where Layer 2 must NOT fire. If the
|
|
production code unexpectedly reaches ``summarize_history``, the
|
|
``AssertionError`` surfaces immediately rather than silently succeeding
|
|
via a ``MagicMock``.
|
|
"""
|
|
|
|
class responses:
|
|
"""Namespace mirroring the real client's ``responses`` attribute."""
|
|
|
|
@staticmethod
|
|
async def create(**kwargs: Any) -> None:
|
|
"""
|
|
Raise if called — Layer 2 must not have fired.
|
|
|
|
:param kwargs: Forwarded kwargs from the real API call.
|
|
:raises AssertionError: Always.
|
|
"""
|
|
raise AssertionError(
|
|
"llm_client.responses.create() was called unexpectedly. "
|
|
"Layer 2 must not fire in this test — check that count_tokens "
|
|
"is mocked below budget or that summarize_history is patched."
|
|
)
|
|
|
|
|
|
class _ReturnsTextClient:
|
|
"""
|
|
LLM client stub that returns a real ``Response`` containing a fixed text.
|
|
|
|
Use this for ``summarize_history`` tests where a real LLM response is
|
|
needed but the test must not hit the network.
|
|
|
|
:param text: The assistant text the stub will return, e.g.
|
|
``"Summary of earlier conversation context."``.
|
|
:param model: The model name to embed in the returned ``Response``, e.g.
|
|
``"openai/gpt-4o"``.
|
|
"""
|
|
|
|
def __init__(self, text: str, model: str = "test-model") -> None:
|
|
self._text = text
|
|
self._model = model
|
|
self.call_count = 0
|
|
|
|
class _Responses:
|
|
"""
|
|
Inner namespace mirroring ``client.responses``.
|
|
|
|
:param outer: The enclosing ``_ReturnsTextClient`` instance.
|
|
"""
|
|
|
|
def __init__(self, outer: _ReturnsTextClient) -> None:
|
|
self._outer = outer
|
|
|
|
async def create(self, **kwargs: Any) -> Response:
|
|
"""
|
|
Return a real ``Response`` with the configured text.
|
|
|
|
:param kwargs: Forwarded kwargs from the real API call.
|
|
:returns: A ``Response`` wrapping the configured text.
|
|
"""
|
|
self._outer.call_count += 1
|
|
return Response(
|
|
output=[MessageOutput(content=[OutputText(text=self._outer._text)])],
|
|
model=self._outer._model,
|
|
)
|
|
|
|
@property
|
|
def responses(self) -> _ReturnsTextClient._Responses:
|
|
"""
|
|
Return the ``responses`` namespace for this stub client.
|
|
|
|
:returns: The ``_Responses`` inner instance.
|
|
"""
|
|
return self._Responses(self)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Module-level helpers
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_conv_item(
|
|
item_id: str,
|
|
item_type: str,
|
|
data: Any,
|
|
response_id: str = "resp_001",
|
|
) -> ConversationItem:
|
|
"""
|
|
Build a ConversationItem for testing.
|
|
|
|
:param item_id: Unique identifier for the item.
|
|
:param item_type: Type string, e.g. "message", "function_call".
|
|
:param data: The item payload (MessageData, FunctionCallData, etc.).
|
|
:param response_id: Response/task identifier to associate with the item.
|
|
"""
|
|
return ConversationItem(
|
|
id=item_id,
|
|
type=item_type,
|
|
status="completed",
|
|
response_id=response_id,
|
|
created_at=1000,
|
|
data=data,
|
|
)
|
|
|
|
|
|
def _user_msg(item_id: str, text: str = "User message") -> ConversationItem:
|
|
"""
|
|
Build a user-role ConversationItem with a single input_text block.
|
|
|
|
:param item_id: Unique identifier for the item.
|
|
:param text: Text content of the user message.
|
|
"""
|
|
return _make_conv_item(
|
|
item_id,
|
|
"message",
|
|
MessageData(role="user", content=[{"type": "input_text", "text": text}]),
|
|
)
|
|
|
|
|
|
def _assistant_msg(item_id: str, text: str = "Assistant response") -> ConversationItem:
|
|
"""
|
|
Build an assistant-role ConversationItem with a single output_text block.
|
|
|
|
:param item_id: Unique identifier for the item.
|
|
:param text: Text content of the assistant response.
|
|
"""
|
|
return _make_conv_item(
|
|
item_id,
|
|
"message",
|
|
MessageData(
|
|
role="assistant",
|
|
content=[{"type": "output_text", "text": text}],
|
|
agent="test-model",
|
|
),
|
|
)
|
|
|
|
|
|
def _fc_item(item_id: str, call_id: str = "call_abc") -> ConversationItem:
|
|
"""
|
|
Build a function_call ConversationItem.
|
|
|
|
:param item_id: Unique identifier for the item.
|
|
:param call_id: Tool call identifier, e.g. "call_abc".
|
|
"""
|
|
return _make_conv_item(
|
|
item_id,
|
|
"function_call",
|
|
FunctionCallData(
|
|
agent="test-model",
|
|
name="my_tool",
|
|
arguments="{}",
|
|
call_id=call_id,
|
|
),
|
|
)
|
|
|
|
|
|
def _fco_item(
|
|
item_id: str,
|
|
call_id: str = "call_abc",
|
|
output: str = "tool result",
|
|
) -> ConversationItem:
|
|
"""
|
|
Build a function_call_output ConversationItem.
|
|
|
|
:param item_id: Unique identifier for the item.
|
|
:param call_id: Tool call identifier matching the originating function_call.
|
|
:param output: The tool output string.
|
|
"""
|
|
return _make_conv_item(
|
|
item_id,
|
|
"function_call_output",
|
|
FunctionCallOutputData(call_id=call_id, output=output),
|
|
)
|
|
|
|
|
|
def _user_msg_dict(text: str = "User message") -> dict[str, Any]:
|
|
"""
|
|
Build a user-role message dict for the messages list.
|
|
|
|
:param text: Text content of the user message.
|
|
"""
|
|
return {"role": "user", "content": [{"type": "input_text", "text": text}]}
|
|
|
|
|
|
def _assistant_msg_dict(text: str = "Assistant response") -> dict[str, Any]:
|
|
"""
|
|
Build an assistant-role message dict for the messages list.
|
|
|
|
:param text: Text content of the assistant response.
|
|
"""
|
|
return {"role": "assistant", "content": [{"type": "output_text", "text": text}]}
|
|
|
|
|
|
def _fc_dict(call_id: str = "call_abc", name: str = "my_tool") -> dict[str, Any]:
|
|
"""
|
|
Build a function_call dict for the messages list.
|
|
|
|
:param call_id: Tool call identifier.
|
|
:param name: Name of the tool being called.
|
|
"""
|
|
return {"type": "function_call", "id": call_id, "name": name, "arguments": "{}"}
|
|
|
|
|
|
def _fco_dict(call_id: str = "call_abc", output: str = "tool result") -> dict[str, Any]:
|
|
"""
|
|
Build a function_call_output dict for the messages list.
|
|
|
|
:param call_id: Tool call identifier matching the originating function_call.
|
|
:param output: The tool output string.
|
|
"""
|
|
return {"type": "function_call_output", "call_id": call_id, "output": output}
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Fixtures
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_compaction_under_threshold(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""Layer 1 always runs but returns early if token count is within budget."""
|
|
monkeypatch.setattr("omnigent.runtime.compaction.count_tokens", lambda msgs, model: 50)
|
|
messages = [_user_msg_dict("hi"), _assistant_msg_dict("hello")]
|
|
history = [_user_msg("msg_001", "hi"), _assistant_msg("msg_002", "hello")]
|
|
|
|
result = await compact(
|
|
messages,
|
|
history,
|
|
config=None,
|
|
context_window=100000,
|
|
system_token_budget=0,
|
|
model="openai/gpt-4o",
|
|
task_id="task_001",
|
|
# _RaisesIfCalled: Layer 2 must not fire (budget met after Layer 1).
|
|
# If summarize_history() is unexpectedly called, the test fails immediately.
|
|
llm_client=_RaisesIfCalled(),
|
|
)
|
|
|
|
# Layer 1 always applies clearing, but since budget is met, returns early.
|
|
# summary_metadata=None proves Layer 2 (summarization) never fired.
|
|
assert result.summary_metadata is None
|
|
# Messages content preserved — no tool result bodies were replaced.
|
|
assert result.messages[0]["content"][0]["text"] == "hi"
|
|
assert result.messages[1]["content"][0]["text"] == "hello"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_layer1_clears_tool_results_outside_window(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""
|
|
Layer 1 replaces function_call_output bodies outside the recent window
|
|
with _TOOL_RESULT_CLEARED, while preserving bodies inside the window.
|
|
|
|
``recent_window=2`` counting backward through [u3,fc3,fco3,a3] and [u2,fc2,fco2,a2]:
|
|
- i=11: a3 → groups=1; i=9: fc3 → groups=2 ≥ 2 → boundary=9.
|
|
- Items 0..8 outside window (eligible for clearing).
|
|
- Items 9..11 inside window (protected: fc3, fco3, a3).
|
|
"""
|
|
monkeypatch.setattr("omnigent.runtime.compaction.count_tokens", lambda msgs, model: 50)
|
|
|
|
history = [
|
|
_user_msg("msg_u1", "iter1"),
|
|
_fc_item("msg_fc1", "c1"),
|
|
_fco_item("msg_fco1", "c1"),
|
|
_assistant_msg("msg_a1"),
|
|
_user_msg("msg_u2", "iter2"),
|
|
_fc_item("msg_fc2", "c2"),
|
|
_fco_item("msg_fco2", "c2"),
|
|
_assistant_msg("msg_a2"),
|
|
_user_msg("msg_u3", "iter3"),
|
|
_fc_item("msg_fc3", "c3"),
|
|
_fco_item("msg_fco3", "c3"),
|
|
_assistant_msg("msg_a3"),
|
|
]
|
|
messages = [
|
|
_user_msg_dict("iter1"),
|
|
_fc_dict("c1"),
|
|
_fco_dict("c1", "tool result iter1"),
|
|
_assistant_msg_dict(),
|
|
_user_msg_dict("iter2"),
|
|
_fc_dict("c2"),
|
|
_fco_dict("c2", "tool result iter2"),
|
|
_assistant_msg_dict(),
|
|
_user_msg_dict("iter3"),
|
|
_fc_dict("c3"),
|
|
_fco_dict("c3", "tool result iter3"),
|
|
_assistant_msg_dict(),
|
|
]
|
|
|
|
result = await compact(
|
|
messages,
|
|
history,
|
|
config=CompactionConfig(trigger_threshold=0.8, recent_window=2),
|
|
context_window=100000,
|
|
system_token_budget=0,
|
|
model="openai/gpt-4o",
|
|
task_id="task_001",
|
|
# _RaisesIfCalled: token count is within budget so Layer 2 must not
|
|
# fire. Fails loudly if summarize_history() is unexpectedly reached.
|
|
llm_client=_RaisesIfCalled(),
|
|
)
|
|
|
|
# fco at index 2 (iter1, outside window) must be cleared.
|
|
assert result.messages[2]["output"] == _TOOL_RESULT_CLEARED, (
|
|
f"Expected iter1 tool result to be cleared (outside window), "
|
|
f"got: {result.messages[2]['output']!r}"
|
|
)
|
|
# fco at index 6 (iter2, outside window) must be cleared.
|
|
assert result.messages[6]["output"] == _TOOL_RESULT_CLEARED, (
|
|
f"Expected iter2 tool result to be cleared (outside window), "
|
|
f"got: {result.messages[6]['output']!r}"
|
|
)
|
|
# fco at index 10 (iter3 — inside window, boundary=9 so index 10 ≥ 9) must be preserved.
|
|
assert result.messages[10]["output"] == "tool result iter3", (
|
|
f"Expected iter3 tool result to be preserved (inside window, boundary=9), "
|
|
f"got: {result.messages[10]['output']!r}"
|
|
)
|
|
# summary_metadata=None confirms only Layer 1 fired (Layer 2 not triggered).
|
|
assert result.summary_metadata is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_layer1_never_touches_user_message_text(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""
|
|
Layer 1 (tool result clearing) must never modify user message text content,
|
|
even for messages outside the recent window.
|
|
"""
|
|
monkeypatch.setattr("omnigent.runtime.compaction.count_tokens", lambda msgs, model: 50)
|
|
|
|
history = [
|
|
_user_msg("msg_u1", "Important user text outside window"),
|
|
_fc_item("msg_fc1", "c1"),
|
|
_fco_item("msg_fco1", "c1"),
|
|
_assistant_msg("msg_a1"),
|
|
_user_msg("msg_u2", "Another user message inside window"),
|
|
_assistant_msg("msg_a2"),
|
|
]
|
|
messages = [
|
|
_user_msg_dict("Important user text outside window"),
|
|
_fc_dict("c1"),
|
|
_fco_dict("c1", "tool output"),
|
|
_assistant_msg_dict(),
|
|
_user_msg_dict("Another user message inside window"),
|
|
_assistant_msg_dict(),
|
|
]
|
|
|
|
result = await compact(
|
|
messages,
|
|
history,
|
|
config=CompactionConfig(trigger_threshold=0.8, recent_window=1),
|
|
context_window=100000,
|
|
system_token_budget=0,
|
|
model="openai/gpt-4o",
|
|
task_id="task_001",
|
|
# _RaisesIfCalled: budget is met after Layer 1 so Layer 2 must not
|
|
# fire. Fails loudly if summarize_history() is unexpectedly reached.
|
|
llm_client=_RaisesIfCalled(),
|
|
)
|
|
|
|
# User text at index 0 (outside window) must be preserved.
|
|
# Failure here means Layer 1 modified non-tool-result content.
|
|
assert result.messages[0]["content"][0]["text"] == "Important user text outside window"
|
|
# User text at index 4 (inside window) must also be preserved.
|
|
assert result.messages[4]["content"][0]["text"] == "Another user message inside window"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_layer1_clears_binary_content_and_preserves_file_id(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""
|
|
Layer 1 clears image/file block data outside the recent window,
|
|
preserves file_id, and leaves text blocks in the same message untouched.
|
|
"""
|
|
monkeypatch.setattr("omnigent.runtime.compaction.count_tokens", lambda msgs, model: 50)
|
|
|
|
# User message with image block (outside window) + text block
|
|
image_msg = {
|
|
"role": "user",
|
|
"content": [
|
|
{"type": "image", "data": "base64IMAGEDATA==", "file_id": "file_abc123"},
|
|
{"type": "text", "text": "Please describe this image"},
|
|
],
|
|
}
|
|
history = [
|
|
_user_msg("msg_u1", "user with image"),
|
|
_assistant_msg("msg_a1"), # boundary (recent_window=1)
|
|
]
|
|
messages = [image_msg, _assistant_msg_dict()]
|
|
|
|
result = await compact(
|
|
messages,
|
|
history,
|
|
config=CompactionConfig(trigger_threshold=0.8, recent_window=1),
|
|
context_window=100000,
|
|
system_token_budget=0,
|
|
model="openai/gpt-4o",
|
|
task_id="task_001",
|
|
# _RaisesIfCalled: budget is met after Layer 1 so Layer 2 must not
|
|
# fire. Fails loudly if summarize_history() is unexpectedly reached.
|
|
llm_client=_RaisesIfCalled(),
|
|
)
|
|
|
|
image_block = result.messages[0]["content"][0]
|
|
text_block = result.messages[0]["content"][1]
|
|
|
|
# Image data must be cleared — the binary payload was replaced.
|
|
assert image_block["data"] == _BINARY_CONTENT_CLEARED, (
|
|
f"Expected image data to be cleared, got: {image_block['data']!r}"
|
|
)
|
|
# file_id must be preserved so the agent can re-fetch the image.
|
|
assert image_block["file_id"] == "file_abc123", (
|
|
f"Expected file_id 'file_abc123' preserved, got: {image_block['file_id']!r}"
|
|
)
|
|
# Text block in the same message must be untouched.
|
|
assert text_block["text"] == "Please describe this image"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_layer1_binary_content_inside_window_untouched(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Binary content inside the recent window must not be cleared by Layer 1."""
|
|
monkeypatch.setattr("omnigent.runtime.compaction.count_tokens", lambda msgs, model: 50)
|
|
|
|
image_msg_outside = {
|
|
"role": "user",
|
|
"content": [{"type": "image", "data": "OLD_DATA==", "file_id": "file_old"}],
|
|
}
|
|
image_msg_inside = {
|
|
"role": "user",
|
|
"content": [{"type": "image", "data": "NEW_DATA==", "file_id": "file_new"}],
|
|
}
|
|
# With recent_window=2 and history [u1, a1, u2, a2]:
|
|
# i=3: a2 → groups=1; i=1: a1 → groups=2 ≥ 2 → boundary=1.
|
|
# Items 0 outside window (image_msg_outside, messages[0]).
|
|
# Items 1..3 inside window (a1, image_msg_inside at msg index 2, a2).
|
|
history = [
|
|
_user_msg("msg_u1"),
|
|
_assistant_msg("msg_a1"),
|
|
_user_msg("msg_u2"),
|
|
_assistant_msg("msg_a2"),
|
|
]
|
|
messages = [image_msg_outside, _assistant_msg_dict(), image_msg_inside, _assistant_msg_dict()]
|
|
|
|
result = await compact(
|
|
messages,
|
|
history,
|
|
config=CompactionConfig(trigger_threshold=0.8, recent_window=2),
|
|
context_window=100000,
|
|
system_token_budget=0,
|
|
model="openai/gpt-4o",
|
|
task_id="task_001",
|
|
# _RaisesIfCalled: budget is met after Layer 1 so Layer 2 must not fire.
|
|
llm_client=_RaisesIfCalled(),
|
|
)
|
|
|
|
# The image OUTSIDE the window (index 0 < boundary=1) should be cleared.
|
|
assert result.messages[0]["content"][0]["data"] == _BINARY_CONTENT_CLEARED
|
|
# The image INSIDE the window (index 2 ≥ boundary=1) must be untouched.
|
|
assert result.messages[2]["content"][0]["data"] == "NEW_DATA==", (
|
|
"Image inside recent window must not be cleared by Layer 1."
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("recent_window", "outside_fco_idx", "protected_fco_idx"),
|
|
[
|
|
# recent_window=2: boundary at index 17 (fc17). fco18 protected; fco14 outside.
|
|
(2, 14, 18),
|
|
# recent_window=3: boundary at index 15 (a15). fco18 protected; fco10 outside.
|
|
(3, 10, 18),
|
|
# recent_window=4: boundary at index 13 (fc13). fco18 protected; fco6 outside.
|
|
(4, 6, 18),
|
|
],
|
|
ids=["window-2", "window-3", "window-4"],
|
|
)
|
|
@pytest.mark.asyncio
|
|
async def test_recent_window_boundary_parametrized(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
recent_window: int,
|
|
outside_fco_idx: int,
|
|
protected_fco_idx: int,
|
|
) -> None:
|
|
"""
|
|
Items inside the recent window must never be modified;
|
|
items outside must have their tool result bodies cleared.
|
|
"""
|
|
monkeypatch.setattr("omnigent.runtime.compaction.count_tokens", lambda msgs, model: 50)
|
|
|
|
history = []
|
|
messages = []
|
|
for i in range(5):
|
|
call_id = f"c{i}"
|
|
history.extend(
|
|
[
|
|
_user_msg(f"msg_u{i}"),
|
|
_fc_item(f"msg_fc{i}", call_id),
|
|
_fco_item(f"msg_fco{i}", call_id, f"output_iter_{i}"),
|
|
_assistant_msg(f"msg_a{i}"),
|
|
]
|
|
)
|
|
messages.extend(
|
|
[
|
|
_user_msg_dict(),
|
|
_fc_dict(call_id),
|
|
_fco_dict(call_id, f"output_iter_{i}"),
|
|
_assistant_msg_dict(),
|
|
]
|
|
)
|
|
|
|
result = await compact(
|
|
messages,
|
|
history,
|
|
config=CompactionConfig(trigger_threshold=0.8, recent_window=recent_window),
|
|
context_window=100000,
|
|
system_token_budget=0,
|
|
model="openai/gpt-4o",
|
|
task_id="task_001",
|
|
# _RaisesIfCalled: count_tokens is mocked below budget, so Layer 2
|
|
# must not fire. Fails loudly if summarize_history() is called.
|
|
llm_client=_RaisesIfCalled(),
|
|
)
|
|
|
|
outside_output = result.messages[outside_fco_idx]["output"]
|
|
protected_output = result.messages[protected_fco_idx]["output"]
|
|
|
|
# Tool result outside the recent window must be cleared.
|
|
assert outside_output == _TOOL_RESULT_CLEARED, (
|
|
f"fco at index {outside_fco_idx} should be cleared (outside window={recent_window}), "
|
|
f"got: {outside_output!r}"
|
|
)
|
|
# Tool result inside the recent window must be preserved.
|
|
assert protected_output != _TOOL_RESULT_CLEARED, (
|
|
f"fco at index {protected_fco_idx} should be preserved (inside window={recent_window}), "
|
|
f"but was cleared"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_layer2_triggers_when_layer1_insufficient(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""
|
|
When Layer 1 alone is insufficient (token count still above budget),
|
|
Layer 2 (LLM summarization) is triggered.
|
|
"""
|
|
call_counts = [0]
|
|
|
|
def mock_count_tokens(msgs: list[dict[str, Any]], model: str) -> int:
|
|
"""
|
|
Return above-budget on the first call to force Layer 2, then below-budget
|
|
for all subsequent calls so Layer 2 can succeed.
|
|
"""
|
|
call_counts[0] += 1
|
|
# First call (after Layer 1): above budget → trigger Layer 2
|
|
# Second call (inside _run_layer2): check if to_summarize too large
|
|
# Third call (after Layer 2): summary + recent fits budget
|
|
if call_counts[0] == 1:
|
|
return 10001 # above budget=10000 → Layer 2 needed
|
|
return 50 # all subsequent calls: below budget
|
|
|
|
monkeypatch.setattr("omnigent.runtime.compaction.count_tokens", mock_count_tokens)
|
|
|
|
async def _stub_summarize(
|
|
msgs: list[dict[str, Any]],
|
|
llm_client: Any,
|
|
model: str,
|
|
connection: dict[str, str] | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Return a fixed summary result."""
|
|
return {
|
|
"text": "Summary of earlier conversation",
|
|
"token_count": 50,
|
|
}
|
|
|
|
monkeypatch.setattr(
|
|
"omnigent.runtime.compaction.summarize_history",
|
|
_stub_summarize,
|
|
)
|
|
|
|
# 2 iterations; recent_window=1 → boundary at index 7 (last assistant)
|
|
history = [
|
|
_user_msg("msg_u1"),
|
|
_fc_item("msg_fc1", "c1"),
|
|
_fco_item("msg_fco1", "c1"),
|
|
_assistant_msg("msg_a1"),
|
|
_user_msg("msg_u2"),
|
|
_fc_item("msg_fc2", "c2"),
|
|
_fco_item("msg_fco2", "c2"),
|
|
_assistant_msg("msg_a2"),
|
|
]
|
|
messages = [
|
|
_user_msg_dict(),
|
|
_fc_dict("c1"),
|
|
_fco_dict("c1"),
|
|
_assistant_msg_dict(),
|
|
_user_msg_dict(),
|
|
_fc_dict("c2"),
|
|
_fco_dict("c2"),
|
|
_assistant_msg_dict(),
|
|
]
|
|
|
|
result = await compact(
|
|
messages,
|
|
history,
|
|
config=CompactionConfig(trigger_threshold=0.8, recent_window=1),
|
|
context_window=12500, # budget = int(12500*0.8) = 10000
|
|
system_token_budget=0,
|
|
model="openai/gpt-4o",
|
|
task_id="task_001",
|
|
# summarize_history is monkeypatched above so llm_client is never used.
|
|
# _RaisesIfCalled still catches any accidental bypass of the patch.
|
|
llm_client=_RaisesIfCalled(),
|
|
)
|
|
|
|
# summary_metadata being set proves Layer 2 fired successfully.
|
|
assert result.summary_metadata is not None, (
|
|
"Layer 2 should have triggered and set summary_metadata, "
|
|
"but it is None — check that mock count_tokens returns > budget on first call."
|
|
)
|
|
# The summary text must match what summarize_history returned.
|
|
assert result.summary_metadata.text == "Summary of earlier conversation"
|
|
# last_item_id must point to a real history item before the boundary.
|
|
# boundary=7 (last assistant) → last summarized item = msg_fco2 at index 6 (non-synthetic).
|
|
# Actually: _find_last_summarized_item_id looks for last non-synthetic item before boundary.
|
|
# With recent_window=1, boundary=7, last item before boundary is msg_fco2 at index 6.
|
|
# Wait - boundary=7 means items 7+ are protected. Items 0..6 are summarized.
|
|
# _find_last_summarized_item_id(history, boundary=7) → history[6] = msg_fco2.
|
|
assert result.summary_metadata.last_item_id == "msg_fco2", (
|
|
f"last_item_id should point to the last item before the boundary, "
|
|
f"got: {result.summary_metadata.last_item_id!r}"
|
|
)
|
|
# The compacted messages should start with the synthetic summary pair
|
|
# (user + assistant messages from _summary_to_messages).
|
|
assert result.messages[0]["role"] == "user"
|
|
assert "automatically generated summary" in result.messages[0]["content"]
|
|
assert result.messages[1]["role"] == "assistant"
|
|
assert result.messages[1]["content"] == "Summary of earlier conversation"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_layer2_failure_falls_back_to_layer3(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""
|
|
When Layer 2 summarization fails, compact() falls back to Layer 3
|
|
(truncation) without raising. summary_metadata is None.
|
|
"""
|
|
# First 2 calls above budget (trigger Layer 2); subsequent calls below budget
|
|
# so Layer 3 stops truncating after one pass (not emptying the list).
|
|
call_idx = [0]
|
|
|
|
def mock_count_tokens(msgs: list[dict[str, Any]], model: str) -> int:
|
|
"""
|
|
Return above-budget on the first 2 calls to trigger Layer 2,
|
|
then below-budget so Layer 3 terminates with remaining messages.
|
|
"""
|
|
call_idx[0] += 1
|
|
return 10001 if call_idx[0] <= 2 else 50
|
|
|
|
monkeypatch.setattr("omnigent.runtime.compaction.count_tokens", mock_count_tokens)
|
|
|
|
async def _raise_retryable(*args: Any, **kwargs: Any) -> dict[str, Any]:
|
|
"""Raise RetryableLLMError to simulate an unavailable LLM."""
|
|
raise RetryableLLMError("LLM unavailable", code="503")
|
|
|
|
monkeypatch.setattr(
|
|
"omnigent.runtime.compaction.summarize_history",
|
|
_raise_retryable,
|
|
)
|
|
|
|
history = [
|
|
_user_msg("msg_u1"),
|
|
_assistant_msg("msg_a1"),
|
|
_user_msg("msg_u2"),
|
|
_assistant_msg("msg_a2"),
|
|
]
|
|
messages = [
|
|
_user_msg_dict("first"),
|
|
_assistant_msg_dict(),
|
|
_user_msg_dict("second"),
|
|
_assistant_msg_dict(),
|
|
]
|
|
|
|
# Must not raise even though summarize_history fails.
|
|
result = await compact(
|
|
messages,
|
|
history,
|
|
config=CompactionConfig(trigger_threshold=0.8, recent_window=1),
|
|
context_window=12500,
|
|
system_token_budget=0,
|
|
model="openai/gpt-4o",
|
|
task_id="task_001",
|
|
# summarize_history is monkeypatched to raise before reaching llm_client.
|
|
# _RaisesIfCalled catches any accidental bypass of the monkeypatch.
|
|
llm_client=_RaisesIfCalled(),
|
|
)
|
|
|
|
# summary_metadata=None proves Layer 2 failed (not persisted).
|
|
assert result.summary_metadata is None, (
|
|
"summary_metadata must be None when Layer 2 summarization fails — "
|
|
"it is only set on successful summarization."
|
|
)
|
|
# Some messages must have been returned (Layer 3 truncated, not emptied).
|
|
assert len(result.messages) > 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_summarize_history_returns_text_and_token_count() -> None:
|
|
"""summarize_history calls the LLM and returns text + token_count > 0."""
|
|
summary_text = "Summary of earlier conversation context."
|
|
stub_llm = _ReturnsTextClient(text=summary_text, model="openai/gpt-4o")
|
|
|
|
messages = [{"role": "user", "content": "prior conversation"}]
|
|
result = await summarize_history(messages, stub_llm, "openai/gpt-4o")
|
|
|
|
# The "text" field must match what the LLM returned.
|
|
assert result["text"] == summary_text, (
|
|
f"Expected summary text from LLM response, got: {result['text']!r}"
|
|
)
|
|
# token_count must be positive — proves count_tokens ran on the text.
|
|
assert result["token_count"] > 0, (
|
|
"token_count must be > 0; failure means count_tokens wasn't called or returned 0."
|
|
)
|
|
# The LLM must have been called exactly once.
|
|
assert stub_llm.call_count == 1, (
|
|
f"Expected 1 LLM call, got {stub_llm.call_count}. "
|
|
"Failure means summarize_history called the LLM more than once or not at all."
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_summarize_history_recursive_prompt_includes_continuation_prefix() -> None:
|
|
"""
|
|
When history starts with a prior summary, the summarization prompt
|
|
includes a 'Incorporate it' continuation instruction.
|
|
"""
|
|
# Prior summary header that triggers recursive detection
|
|
prior_summary_header = (
|
|
"[This is an automatically generated summary of the prior conversation "
|
|
"context. The original messages are available but not included in this "
|
|
"prompt for brevity.]\n\n"
|
|
"Please provide a summary of our conversation so far."
|
|
)
|
|
messages = [
|
|
{"role": "user", "content": prior_summary_header},
|
|
{"role": "assistant", "content": "Earlier we discussed X and Y."},
|
|
{"role": "user", "content": "Now let's continue with Z."},
|
|
]
|
|
|
|
captured_instructions: list[str] = []
|
|
mock_resp = Response(
|
|
output=[MessageOutput(content=[OutputText(text="Combined summary.")])],
|
|
model="openai/gpt-4o",
|
|
)
|
|
|
|
class _CapturingClient:
|
|
"""LLM client stub that captures ``instructions`` from each call."""
|
|
|
|
class responses:
|
|
"""Namespace mirroring the real client's ``responses`` attribute."""
|
|
|
|
@staticmethod
|
|
async def create(**kwargs: Any) -> Response:
|
|
"""Capture the instructions kwarg and return the mock response."""
|
|
captured_instructions.append(kwargs.get("instructions", ""))
|
|
return mock_resp
|
|
|
|
result = await summarize_history(messages, _CapturingClient(), "openai/gpt-4o")
|
|
|
|
assert len(captured_instructions) == 1
|
|
# The continuation prefix must be present when history starts with a prior summary.
|
|
assert "Incorporate it into your new summary" in captured_instructions[0], (
|
|
"Recursive summarization prompt must include the 'Incorporate it' instruction; "
|
|
"failure means _build_summarization_prompt did not detect the prior summary header."
|
|
)
|
|
assert result["text"] == "Combined summary."
|
|
|
|
|
|
def test_compaction_to_history_items_produces_valid_pair() -> None:
|
|
"""
|
|
compaction_to_history_items() produces a user+assistant synthetic pair
|
|
for inclusion at the start of conversation history.
|
|
"""
|
|
compaction_item = ConversationItem(
|
|
id="cmp_abc123",
|
|
type="compaction",
|
|
status="completed",
|
|
response_id="task_001",
|
|
created_at=1000,
|
|
data=CompactionData(
|
|
summary="The user asked to analyze the dataset. The agent loaded data.csv.",
|
|
last_item_id="msg_xyz999",
|
|
model="openai/gpt-4o",
|
|
token_count=42,
|
|
),
|
|
)
|
|
|
|
result = compaction_to_history_items(compaction_item)
|
|
|
|
# Must return exactly 2 items: synthetic user + assistant.
|
|
assert len(result) == 2, (
|
|
f"Expected exactly 2 items (user + assistant), got {len(result)}. "
|
|
"Failure means compaction_to_history_items changed its output shape."
|
|
)
|
|
user_item = result[0]
|
|
assistant_item = result[1]
|
|
|
|
# Both items must be message type for history processing.
|
|
assert user_item.type == "message"
|
|
assert assistant_item.type == "message"
|
|
|
|
# User item must have role=user.
|
|
assert isinstance(user_item.data, MessageData)
|
|
assert user_item.data.role == "user"
|
|
|
|
# User content must contain the summary marker prefix so the LLM
|
|
# understands this is synthetic context, not a real prior message.
|
|
user_text = user_item.data.content[0]["text"]
|
|
assert "[This is an automatically generated summary" in user_text, (
|
|
"User content must contain the summary marker prefix — "
|
|
"failure means the synthetic header was changed or removed."
|
|
)
|
|
|
|
# Assistant item must have the summary text verbatim.
|
|
assert isinstance(assistant_item.data, MessageData)
|
|
assert assistant_item.data.role == "assistant"
|
|
assistant_text = assistant_item.data.content[0]["text"]
|
|
assert assistant_text == "The user asked to analyze the dataset. The agent loaded data.csv.", (
|
|
f"Assistant content must equal the CompactionData.summary, got: {assistant_text!r}"
|
|
)
|
|
|
|
# IDs must be derived from the compaction item ID.
|
|
assert user_item.id == "cmp_abc123_user"
|
|
assert assistant_item.id == "cmp_abc123_assistant"
|
|
|
|
|
|
def test_count_tokens_returns_positive_integer() -> None:
|
|
"""count_tokens returns a positive integer for non-empty messages."""
|
|
messages = [{"role": "user", "content": "Hello world, this is a test message."}]
|
|
result = count_tokens(messages, "openai/gpt-4o")
|
|
# Must be a positive integer; failure means tiktoken encoding failed.
|
|
assert isinstance(result, int)
|
|
assert result > 0
|
|
|
|
|
|
def test_count_tokens_unknown_model_falls_back() -> None:
|
|
"""Unknown model falls back to cl100k_base encoding without raising."""
|
|
messages = [{"role": "user", "content": "test"}]
|
|
# Should not raise even for completely unknown model names.
|
|
result = count_tokens(messages, "unknown/totally-fake-model-xyz")
|
|
assert result > 0
|
|
|
|
|
|
async def test_layer3_truncation_preserves_tool_call_pairs(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""
|
|
Layer 3 truncation drops tool call pairs together — never
|
|
orphans a ``function_call`` without its ``function_call_output``.
|
|
"""
|
|
|
|
def mock_count_tokens(
|
|
msgs: list[dict[str, Any]],
|
|
model: str,
|
|
) -> int:
|
|
"""
|
|
Simulates shrinking token count as messages are truncated.
|
|
|
|
:param msgs: Messages list (length used to simulate shrinking).
|
|
:param model: Model string (unused).
|
|
:returns: Token count proportional to message count.
|
|
"""
|
|
# Each message ~ 5000 tokens. Budget is 10000, so need <= 2 msgs.
|
|
return len(msgs) * 5000
|
|
|
|
monkeypatch.setattr(
|
|
"omnigent.runtime.compaction.count_tokens",
|
|
mock_count_tokens,
|
|
)
|
|
|
|
# Layer 2 fails so we fall through to Layer 3.
|
|
async def _raise_layer2(
|
|
msgs: list[dict[str, Any]],
|
|
llm_client: Any,
|
|
model: str,
|
|
connection: dict[str, str] | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Raise to simulate Layer 2 failure."""
|
|
raise RuntimeError("Simulated Layer 2 failure")
|
|
|
|
monkeypatch.setattr(
|
|
"omnigent.runtime.compaction.summarize_history",
|
|
_raise_layer2,
|
|
)
|
|
|
|
# Layout: user, fc+fco pair, assistant, user, assistant
|
|
# 6 messages x 5000 = 30000 > budget 10000
|
|
# Layer 3 must drop from front but keep fc+fco together.
|
|
history = [
|
|
_user_msg("msg_u1"),
|
|
_fc_item("msg_fc1", "c1"),
|
|
_fco_item("msg_fco1", "c1"),
|
|
_assistant_msg("msg_a1"),
|
|
_user_msg("msg_u2"),
|
|
_assistant_msg("msg_a2"),
|
|
]
|
|
messages = [
|
|
_user_msg_dict(),
|
|
_fc_dict("c1"),
|
|
_fco_dict("c1"),
|
|
_assistant_msg_dict(),
|
|
_user_msg_dict(),
|
|
_assistant_msg_dict(),
|
|
]
|
|
|
|
result = await compact(
|
|
messages,
|
|
history,
|
|
config=CompactionConfig(trigger_threshold=0.8, recent_window=1),
|
|
context_window=12500, # budget = 10000
|
|
system_token_budget=0,
|
|
model="openai/gpt-4o",
|
|
task_id="task_trunc",
|
|
llm_client=_RaisesIfCalled(),
|
|
)
|
|
|
|
# After truncation, remaining messages must not have orphaned pairs.
|
|
# An orphaned function_call_output without its function_call is a bug.
|
|
remaining_types = [m.get("type", m.get("role", "unknown")) for m in result.messages]
|
|
|
|
for i, msg in enumerate(result.messages):
|
|
if msg.get("type") == "function_call_output":
|
|
# The preceding message must be its matching function_call.
|
|
assert i > 0, "function_call_output at index 0 is orphaned without its function_call."
|
|
prev = result.messages[i - 1]
|
|
assert prev.get("type") == "function_call", (
|
|
f"function_call_output at index {i} is preceded by "
|
|
f"{prev.get('type', prev.get('role'))!r}, not "
|
|
f"'function_call'. The pair was broken by truncation. "
|
|
f"Remaining types: {remaining_types}"
|
|
)
|
|
assert prev.get("call_id") == msg.get("call_id"), (
|
|
f"function_call.call_id={prev.get('call_id')!r} doesn't "
|
|
f"match function_call_output.call_id="
|
|
f"{msg.get('call_id')!r} at index {i}."
|
|
)
|
|
|
|
# Must have truncated at least 2 messages (budget fits <= 2).
|
|
# Original had 6 messages x 5000 = 30000 > 10000 budget.
|
|
assert len(result.messages) <= 2, (
|
|
f"Expected <= 2 messages after truncation (budget=10000, "
|
|
f"5000 tokens/msg), got {len(result.messages)}. "
|
|
f"Layer 3 didn't truncate enough."
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_layer2_receives_cleared_content(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""
|
|
When Layer 2 fires, ``summarize_history`` receives messages
|
|
that already have tool result bodies cleared by Layer 1.
|
|
|
|
Captures the messages passed to ``summarize_history`` and
|
|
verifies the tool result body was replaced with the clearing
|
|
marker before summarization.
|
|
"""
|
|
call_counts = [0]
|
|
|
|
def mock_count_tokens(
|
|
msgs: list[dict[str, Any]],
|
|
model: str,
|
|
) -> int:
|
|
"""
|
|
First call above budget to trigger Layer 2, then below.
|
|
|
|
:param msgs: Messages (unused).
|
|
:param model: Model string (unused).
|
|
:returns: Token count.
|
|
"""
|
|
call_counts[0] += 1
|
|
if call_counts[0] == 1:
|
|
return 10001
|
|
return 50
|
|
|
|
monkeypatch.setattr(
|
|
"omnigent.runtime.compaction.count_tokens",
|
|
mock_count_tokens,
|
|
)
|
|
|
|
captured_inputs: list[list[dict[str, Any]]] = []
|
|
|
|
async def _capturing_summarize(
|
|
msgs: list[dict[str, Any]],
|
|
llm_client: Any,
|
|
model: str,
|
|
connection: dict[str, str] | None = None,
|
|
) -> dict[str, Any]:
|
|
"""
|
|
Capture the messages passed to summarize_history.
|
|
|
|
:param msgs: The messages to summarize — should have
|
|
cleared tool result bodies.
|
|
:param llm_client: LLM client (unused).
|
|
:param model: Model string (unused).
|
|
:param connection: Connection params (unused).
|
|
:returns: Fake summary result.
|
|
"""
|
|
captured_inputs.append(list(msgs))
|
|
return {"text": "Summary", "token_count": 10}
|
|
|
|
monkeypatch.setattr(
|
|
"omnigent.runtime.compaction.summarize_history",
|
|
_capturing_summarize,
|
|
)
|
|
|
|
# History with a tool call pair OUTSIDE the recent window.
|
|
history = [
|
|
_user_msg("msg_u1"),
|
|
_fc_item("msg_fc1", "c1"),
|
|
_fco_item("msg_fco1", "c1", output="verbose tool output"),
|
|
_assistant_msg("msg_a1"),
|
|
_user_msg("msg_u2"),
|
|
_assistant_msg("msg_a2"),
|
|
]
|
|
messages = [
|
|
_user_msg_dict(),
|
|
_fc_dict("c1"),
|
|
_fco_dict("c1", "verbose tool output"),
|
|
_assistant_msg_dict(),
|
|
_user_msg_dict(),
|
|
_assistant_msg_dict(),
|
|
]
|
|
|
|
await compact(
|
|
messages,
|
|
history,
|
|
config=CompactionConfig(trigger_threshold=0.8, recent_window=1),
|
|
context_window=12500,
|
|
system_token_budget=0,
|
|
model="openai/gpt-4o",
|
|
task_id="task_layer1_feeds_layer2",
|
|
llm_client=_RaisesIfCalled(),
|
|
)
|
|
|
|
# summarize_history must have been called exactly once.
|
|
assert len(captured_inputs) == 1, (
|
|
f"Expected 1 call to summarize_history, got {len(captured_inputs)}."
|
|
)
|
|
|
|
# The tool result body in the summarization input must be cleared.
|
|
# Layer 1 runs before Layer 2, so fco at index 2 (outside window)
|
|
# should have its output replaced with the clearing marker.
|
|
summarized = captured_inputs[0]
|
|
fco_in_summary = [m for m in summarized if m.get("type") == "function_call_output"]
|
|
assert len(fco_in_summary) >= 1, (
|
|
"Expected at least 1 function_call_output in summarization input."
|
|
)
|
|
assert fco_in_summary[0]["output"] == _TOOL_RESULT_CLEARED, (
|
|
f"Tool result body should be cleared before reaching "
|
|
f"summarize_history, got: {fco_in_summary[0]['output']!r}. "
|
|
f"If it contains 'verbose tool output', Layer 1 didn't "
|
|
f"clear before passing to Layer 2."
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_layer3_fires_when_summary_plus_recent_exceeds_budget(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""
|
|
Layer 3 fires as the primary path (not fallback) when Layer 2
|
|
succeeds but the summary + recent messages together still
|
|
exceed the budget.
|
|
|
|
This differs from ``test_layer2_failure_falls_back_to_layer3``
|
|
which tests Layer 3 after Layer 2 failure. Here Layer 2
|
|
succeeds but its output is still too large.
|
|
"""
|
|
call_counts = [0]
|
|
|
|
def mock_count_tokens(
|
|
msgs: list[dict[str, Any]],
|
|
model: str,
|
|
) -> int:
|
|
"""
|
|
Always above budget so Layer 3 must truncate.
|
|
|
|
:param msgs: Messages (length-based estimate).
|
|
:param model: Model string (unused).
|
|
:returns: Token count.
|
|
"""
|
|
call_counts[0] += 1
|
|
# First call (after Layer 1): above budget → Layer 2 fires
|
|
if call_counts[0] == 1:
|
|
return 10001
|
|
# Second call (inside _run_layer2 size check): below budget
|
|
# so summarization input fits the model
|
|
if call_counts[0] == 2:
|
|
return 50
|
|
# Third call (summary + recent budget check): ABOVE budget
|
|
# → Layer 2 output doesn't fit, fall through to Layer 3
|
|
if call_counts[0] == 3:
|
|
return 10001
|
|
# Layer 3 truncation calls: shrink per message
|
|
return len(msgs) * 3000
|
|
|
|
monkeypatch.setattr(
|
|
"omnigent.runtime.compaction.count_tokens",
|
|
mock_count_tokens,
|
|
)
|
|
|
|
async def _stub_summarize(
|
|
msgs: list[dict[str, Any]],
|
|
llm_client: Any,
|
|
model: str,
|
|
connection: dict[str, str] | None = None,
|
|
) -> dict[str, Any]:
|
|
"""Return a large summary that still exceeds budget."""
|
|
return {
|
|
"text": "A very long summary that still exceeds budget",
|
|
"token_count": 9000,
|
|
}
|
|
|
|
monkeypatch.setattr(
|
|
"omnigent.runtime.compaction.summarize_history",
|
|
_stub_summarize,
|
|
)
|
|
|
|
history = [
|
|
_user_msg("msg_u1"),
|
|
_assistant_msg("msg_a1"),
|
|
_user_msg("msg_u2"),
|
|
_assistant_msg("msg_a2"),
|
|
_user_msg("msg_u3"),
|
|
_assistant_msg("msg_a3"),
|
|
]
|
|
messages = [
|
|
_user_msg_dict("m1"),
|
|
_assistant_msg_dict("a1"),
|
|
_user_msg_dict("m2"),
|
|
_assistant_msg_dict("a2"),
|
|
_user_msg_dict("m3"),
|
|
_assistant_msg_dict("a3"),
|
|
]
|
|
|
|
result = await compact(
|
|
messages,
|
|
history,
|
|
config=CompactionConfig(trigger_threshold=0.8, recent_window=1),
|
|
context_window=12500,
|
|
system_token_budget=0,
|
|
model="openai/gpt-4o",
|
|
task_id="task_layer3_primary",
|
|
llm_client=_RaisesIfCalled(),
|
|
)
|
|
|
|
# Layer 2 succeeded (summary was produced) but the combined
|
|
# result was too large. Layer 3 must have truncated further.
|
|
# The result should have fewer messages than summary + recent.
|
|
assert len(result.messages) < 6, (
|
|
f"Expected Layer 3 truncation to reduce message count "
|
|
f"below 6, got {len(result.messages)}. "
|
|
f"If 6, Layer 3 didn't fire after Layer 2's output "
|
|
f"exceeded the budget."
|
|
)
|
|
# The first message should be the synthetic summary user
|
|
# message (from Layer 2's output), possibly truncated further.
|
|
assert result.messages[0]["role"] == "user", (
|
|
f"First message should be the summary user message, "
|
|
f"got role={result.messages[0].get('role')!r}."
|
|
)
|
|
|
|
|
|
def test_pair_aware_drop_count_drops_both_when_pair_at_front() -> None:
|
|
"""
|
|
When the first two messages are a function_call followed by its
|
|
matching function_call_output, both must be dropped together.
|
|
|
|
If only one were dropped, the LLM would see an orphaned
|
|
function_call_output without its parent call (or vice versa).
|
|
"""
|
|
messages = [
|
|
{"type": "function_call", "call_id": "c1", "name": "grep", "arguments": "{}"},
|
|
{"type": "function_call_output", "call_id": "c1", "output": "result"},
|
|
_user_msg_dict("after the pair"),
|
|
]
|
|
assert _pair_aware_drop_count(messages) == 2, (
|
|
"Expected 2 (drop both halves of the tool call pair). "
|
|
"If 1, the function_call_output would be orphaned."
|
|
)
|
|
|
|
|
|
def test_pair_aware_drop_count_drops_one_for_non_pair() -> None:
|
|
"""
|
|
When the front message is not part of a tool call pair, only
|
|
one item should be dropped.
|
|
"""
|
|
messages = [
|
|
_user_msg_dict("hello"),
|
|
_assistant_msg_dict("world"),
|
|
]
|
|
assert _pair_aware_drop_count(messages) == 1, (
|
|
"Expected 1 for a plain user message at the front."
|
|
)
|
|
|
|
|
|
def test_pair_aware_drop_count_drops_one_for_mismatched_call_ids() -> None:
|
|
"""
|
|
A function_call followed by a function_call_output with a
|
|
*different* call_id is not a pair — drop only the first item.
|
|
"""
|
|
messages = [
|
|
{"type": "function_call", "call_id": "c1", "name": "grep", "arguments": "{}"},
|
|
{"type": "function_call_output", "call_id": "c2", "output": "result"},
|
|
]
|
|
assert _pair_aware_drop_count(messages) == 1, (
|
|
"Expected 1 — mismatched call_ids means these are not a "
|
|
"pair, so only the first item should be dropped."
|
|
)
|
|
|
|
|
|
def test_pair_aware_drop_count_returns_zero_for_empty() -> None:
|
|
"""Empty list returns 0 — nothing to drop."""
|
|
assert _pair_aware_drop_count([]) == 0
|
|
|
|
|
|
def test_truncate_oldest_preserves_tool_call_pairs(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""
|
|
_truncate_oldest drops function_call + function_call_output
|
|
together, never leaving an orphaned half.
|
|
|
|
Uses a mock token counter that returns above-budget on the
|
|
first call (triggering one drop) then below-budget so
|
|
truncation stops.
|
|
"""
|
|
call_count = [0]
|
|
|
|
def mock_count_tokens(msgs: list[dict[str, Any]], model: str) -> int:
|
|
"""
|
|
Above budget on first call to trigger one truncation round,
|
|
then below budget so the loop exits.
|
|
|
|
:param msgs: Messages list.
|
|
:param model: Model string (unused).
|
|
:returns: Token count.
|
|
"""
|
|
call_count[0] += 1
|
|
# First call: above budget. After dropping the pair (2 items),
|
|
# second call: below budget.
|
|
return 10000 if call_count[0] == 1 else 50
|
|
|
|
monkeypatch.setattr(
|
|
"omnigent.runtime.compaction.count_tokens",
|
|
mock_count_tokens,
|
|
)
|
|
|
|
messages = [
|
|
{"type": "function_call", "call_id": "c1", "name": "grep", "arguments": "{}"},
|
|
{"type": "function_call_output", "call_id": "c1", "output": "grep result"},
|
|
_user_msg_dict("kept message"),
|
|
]
|
|
|
|
result = _truncate_oldest(messages, budget=100, model="test")
|
|
|
|
# The pair (indices 0-1) must be dropped together, leaving
|
|
# only the user message. If the pair were split, we'd see
|
|
# an orphaned function_call_output at index 0.
|
|
assert len(result) == 1, (
|
|
f"Expected 1 message after dropping the tool call pair, "
|
|
f"got {len(result)}. If 2, only one half of the pair was "
|
|
f"dropped (orphaned tool call)."
|
|
)
|
|
assert result[0]["role"] == "user", (
|
|
f"Expected the surviving message to be the user message, "
|
|
f"got type={result[0].get('type')!r} role={result[0].get('role')!r}."
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_compaction_strips_annotations_before_summarization(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""
|
|
Layer 1 strips ``annotations`` from ``output_text`` blocks
|
|
before they reach Layer 2 summarization.
|
|
|
|
Captures the messages passed to ``summarize_history`` and
|
|
verifies annotations are absent.
|
|
"""
|
|
call_counts = [0]
|
|
|
|
def mock_count_tokens(
|
|
msgs: list[dict[str, Any]],
|
|
model: str,
|
|
) -> int:
|
|
"""
|
|
First call above budget to trigger Layer 2, then below.
|
|
|
|
:param msgs: Messages (unused).
|
|
:param model: Model string (unused).
|
|
:returns: Token count.
|
|
"""
|
|
call_counts[0] += 1
|
|
if call_counts[0] == 1:
|
|
return 10001
|
|
return 50
|
|
|
|
monkeypatch.setattr(
|
|
"omnigent.runtime.compaction.count_tokens",
|
|
mock_count_tokens,
|
|
)
|
|
|
|
captured_inputs: list[list[dict[str, Any]]] = []
|
|
|
|
async def _capturing_summarize(
|
|
msgs: list[dict[str, Any]],
|
|
llm_client: Any,
|
|
model: str,
|
|
connection: dict[str, str] | None = None,
|
|
) -> dict[str, Any]:
|
|
"""
|
|
Capture summarization input for assertion.
|
|
|
|
:param msgs: Messages to summarize.
|
|
:param llm_client: LLM client (unused).
|
|
:param model: Model string (unused).
|
|
:param connection: Connection params (unused).
|
|
:returns: Fake summary.
|
|
"""
|
|
captured_inputs.append(list(msgs))
|
|
return {"text": "Summary", "token_count": 10}
|
|
|
|
monkeypatch.setattr(
|
|
"omnigent.runtime.compaction.summarize_history",
|
|
_capturing_summarize,
|
|
)
|
|
|
|
# Assistant message with file_citation annotation OUTSIDE recent window.
|
|
annotated_msg = {
|
|
"role": "assistant",
|
|
"content": [
|
|
{
|
|
"type": "output_text",
|
|
"text": "Here is the chart:",
|
|
"annotations": [
|
|
{
|
|
"type": "file_citation",
|
|
"file_id": "file_abc123",
|
|
"filename": "chart.png",
|
|
"content_type": "image/png",
|
|
}
|
|
],
|
|
}
|
|
],
|
|
}
|
|
history = [
|
|
_user_msg("msg_u1"),
|
|
_assistant_msg("msg_a1"),
|
|
_user_msg("msg_u2"),
|
|
_assistant_msg("msg_a2"),
|
|
]
|
|
messages = [
|
|
_user_msg_dict(),
|
|
annotated_msg,
|
|
_user_msg_dict(),
|
|
_assistant_msg_dict(),
|
|
]
|
|
|
|
await compact(
|
|
messages,
|
|
history,
|
|
config=CompactionConfig(trigger_threshold=0.8, recent_window=1),
|
|
context_window=12500,
|
|
system_token_budget=0,
|
|
model="openai/gpt-4o",
|
|
task_id="task_ann_strip",
|
|
llm_client=_RaisesIfCalled(),
|
|
)
|
|
|
|
# summarize_history must have been called.
|
|
assert len(captured_inputs) == 1, (
|
|
f"Expected 1 summarize_history call, got {len(captured_inputs)}."
|
|
)
|
|
|
|
# The annotated output_text block in the summarization input
|
|
# must NOT have annotations — they should be stripped by Layer 1.
|
|
summarized = captured_inputs[0]
|
|
for msg in summarized:
|
|
content = msg.get("content")
|
|
if not isinstance(content, list):
|
|
continue
|
|
for block in content:
|
|
if isinstance(block, dict) and block.get("type") == "output_text":
|
|
assert "annotations" not in block, (
|
|
f"output_text block still has annotations in "
|
|
f"summarization input: {block}. Layer 1 should "
|
|
f"have stripped them."
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Budget honors the declared/effective context window
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_declared_window_keeps_large_fill_under_budget(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""
|
|
A declared 1M window keeps a Polly-scale fill (~197K) under budget.
|
|
|
|
Regression for the runner over-compaction bug: budget is
|
|
``context_window * trigger_threshold``. With the declared 1M window
|
|
(resolved via resolve_effective_context_window), budget=800K and a 197K
|
|
fill does NOT trigger Layer 2. If the window were the 128K catalog default
|
|
(budget=102400), the same fill would compact — which is the bug.
|
|
"""
|
|
monkeypatch.setattr("omnigent.runtime.compaction.count_tokens", lambda msgs, model: 197_000)
|
|
messages = [_user_msg_dict("hi"), _assistant_msg_dict("hello")]
|
|
history = [_user_msg("msg_001", "hi"), _assistant_msg("msg_002", "hello")]
|
|
|
|
window = resolve_effective_context_window(1_000_000, "claude-opus-4-8")
|
|
assert window == 1_000_000
|
|
|
|
result = await compact(
|
|
messages,
|
|
history,
|
|
config=CompactionConfig(trigger_threshold=0.8, recent_window=1),
|
|
context_window=window,
|
|
system_token_budget=0,
|
|
model="claude-opus-4-8",
|
|
task_id="task_001",
|
|
# 197K <= 0.8 * 1M = 800K → under budget → Layer 2 must NOT fire.
|
|
llm_client=_RaisesIfCalled(),
|
|
)
|
|
assert result.summary_metadata is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_default_window_compacts_same_fill(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
"""
|
|
The same ~197K fill DOES compact against the 128K catalog default.
|
|
|
|
Contrast with the test above: this is the pre-fix behavior (budget=102400),
|
|
confirming the window value is what flips compaction on/off.
|
|
"""
|
|
monkeypatch.setattr("omnigent.runtime.compaction.count_tokens", lambda msgs, model: 197_000)
|
|
messages = [_user_msg_dict("hi"), _assistant_msg_dict("hello")]
|
|
history = [_user_msg("msg_001", "hi"), _assistant_msg("msg_002", "hello")]
|
|
|
|
result = await compact(
|
|
messages,
|
|
history,
|
|
config=CompactionConfig(trigger_threshold=0.8, recent_window=1),
|
|
context_window=128_000,
|
|
system_token_budget=0,
|
|
model="claude-opus-4-8",
|
|
task_id="task_001",
|
|
# 197K > 0.8 * 128K = 102400 → over budget → Layer 2 fires.
|
|
llm_client=_ReturnsTextClient("Summary of earlier context."),
|
|
)
|
|
assert result.summary_metadata is not None
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Layer-2 auth-failure detection (don't bury a 401)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_is_summary_auth_error_distinguishes_401_403() -> None:
|
|
"""401/403 (by status code or message) are auth errors; others are not."""
|
|
|
|
class _Resp:
|
|
def __init__(self, code: int) -> None:
|
|
self.status_code = code
|
|
|
|
class _HTTPStatusError(Exception):
|
|
def __init__(self, code: int) -> None:
|
|
super().__init__(f"status {code}")
|
|
self.response = _Resp(code)
|
|
|
|
# By response.status_code (as httpx.HTTPStatusError carries it).
|
|
assert _is_summary_auth_error(_HTTPStatusError(401)) is True
|
|
assert _is_summary_auth_error(_HTTPStatusError(403)) is True
|
|
assert _is_summary_auth_error(_HTTPStatusError(500)) is False
|
|
# By message text (when no structured response is attached).
|
|
assert _is_summary_auth_error(Exception("Client error '401 Unauthorized'")) is True
|
|
assert _is_summary_auth_error(Exception("Forbidden")) is True
|
|
# Non-auth failures fall through to the generic warning path.
|
|
assert _is_summary_auth_error(Exception("connection reset by peer")) is False
|
|
assert _is_summary_auth_error(TimeoutError("read timed out")) is False
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Compaction model routing (issue #1950)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def test_route_bare_model_prefixes_anthropic_claude() -> None:
|
|
"""A bare ``claude-*`` id must route to Anthropic, not the OpenAI default.
|
|
|
|
Regression for #1950: explicit ``/compact`` on a ``claude-sdk`` agent with a
|
|
pinned bare Anthropic model (e.g. ``claude-haiku-4-5-20251001``) sent the id
|
|
to ``api.openai.com`` (routing.py defaults prefix-less ids to OpenAI), so the
|
|
summarization LLM call 500'd. It must be nudged to ``anthropic/…``.
|
|
"""
|
|
from omnigent.llms.routing import parse_model_string
|
|
|
|
out = _route_bare_model_for_compaction(LLMConfig(model="claude-haiku-4-5-20251001"))
|
|
assert out.model == "anthropic/claude-haiku-4-5-20251001"
|
|
# And the generic client now routes it to Anthropic rather than OpenAI.
|
|
assert parse_model_string(out.model).provider == "anthropic"
|
|
|
|
|
|
def test_route_bare_model_preserves_databricks_and_prefixed_ids() -> None:
|
|
"""databricks-* still gets the databricks/ prefix; already-routed ids pass through."""
|
|
assert (
|
|
_route_bare_model_for_compaction(LLMConfig(model="databricks-claude-sonnet-4")).model
|
|
== "databricks/databricks-claude-sonnet-4"
|
|
)
|
|
# Already provider-prefixed — left untouched (no double prefixing).
|
|
for prefixed in ("anthropic/claude-haiku-4-5-20251001", "openai/gpt-4o", "databricks/foo"):
|
|
assert _route_bare_model_for_compaction(LLMConfig(model=prefixed)).model == prefixed
|
|
# Bare gpt-* is correctly OpenAI already — unchanged.
|
|
assert _route_bare_model_for_compaction(LLMConfig(model="gpt-4o")).model == "gpt-4o"
|