test/policies/test_reasoning_pairing.py
Rama main · réplica de lectura del repositorio original.
# Copyright (c) 2026, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
#
# SPDX-License-Identifier: Apache-2.0
"""Trimming policies must not split a retained event from the provider-native
item it needs — a builtin tool call from its reasoning item, or a response
carrying tool calls from the turn object that is the only record of them."""
import pytest
from ag2 import Context, ToolResult
from ag2.events import (
BaseEvent,
BuiltinToolCallEvent,
BuiltinToolResultEvent,
ModelMessage,
ModelReasoning,
ModelRequest,
ModelResponse,
ProviderReplay,
TextInput,
ToolCallEvent,
ToolCallsEvent,
ToolResultEvent,
ToolResultsEvent,
UsageEvent,
)
from ag2.policies.sliding_window import SlidingWindowPolicy
from ag2.policies.token_budget import TokenBudgetPolicy
from test._helpers import DurableReasoning, ProviderTurnState
def _call(call_id: str) -> BuiltinToolCallEvent:
return BuiltinToolCallEvent(id=call_id, name="web_search", arguments="{}")
def _result(parent_id: str) -> BuiltinToolResultEvent:
return BuiltinToolResultEvent(parent_id=parent_id, name="web_search", result=ToolResult("ok"))
def _chars(events: list[BaseEvent]) -> int:
"""Size of the given events in the characters the policy counts."""
return sum(len(str(e)) for e in events)
def _budget_for(events: list[BaseEvent]) -> int:
"""Token budget that fits exactly the given events."""
return _chars(events) // 4 + 1
@pytest.mark.asyncio
class TestSlidingWindow:
async def test_orphaned_builtin_call_is_dropped(self, context: Context) -> None:
events = [
DurableReasoning("plan"),
_call("ws_1"),
_result("ws_1"),
ModelRequest([TextInput("next")]),
]
policy = SlidingWindowPolicy(max_events=3)
_, result = await policy.apply([], events, context)
assert result == [events[-1]]
async def test_intact_group_is_kept(self, context: Context) -> None:
events = [
ModelRequest([TextInput("old")]),
DurableReasoning("plan"),
_call("ws_1"),
_result("ws_1"),
]
policy = SlidingWindowPolicy(max_events=3)
_, result = await policy.apply([], events, context)
assert result == events[1:]
async def test_builtin_calls_without_reasoning_are_kept(self, context: Context) -> None:
# Non-reasoning models emit no reasoning item, so nothing can be orphaned.
events = [
ModelRequest([TextInput("old")]),
_call("ws_1"),
_result("ws_1"),
ModelRequest([TextInput("next")]),
]
policy = SlidingWindowPolicy(max_events=3)
_, result = await policy.apply([], events, context)
assert result == events[1:]
async def test_later_group_keeps_its_own_anchor(self, context: Context) -> None:
# The cut splits group one; group two carries its own reasoning item.
events = [
DurableReasoning("plan one"),
_call("ws_1"),
DurableReasoning("plan two"),
_call("ws_2"),
]
policy = SlidingWindowPolicy(max_events=3)
_, result = await policy.apply([], events, context)
assert result == events[2:]
async def test_anchor_does_not_leak_past_a_response_boundary(self, context: Context) -> None:
# The second group is self-contained: it has no reasoning item of its own,
# so nothing about it was orphaned and the first group's anchor must not
# condemn it. Live, the same model emits a reasoning item only when asked
# for a summary, so one history really can mix anchored and unanchored
# builtin calls.
events = [
DurableReasoning("plan one"),
_call("ws_1"),
_result("ws_1"),
ModelRequest([TextInput("next")]),
_call("ws_2"),
_result("ws_2"),
]
policy = SlidingWindowPolicy(max_events=2)
_, result = await policy.apply([], events, context)
assert result == events[4:]
async def test_anchor_does_not_leak_past_a_model_response(self, context: Context) -> None:
# A ModelResponse closes the response that emitted the anchor, so a builtin
# call after it belongs to a later response and needs its own anchor.
events = [
DurableReasoning("plan one"),
_call("ws_1"),
ModelResponse(tool_calls=ToolCallsEvent(calls=[])),
_call("ws_2"),
_result("ws_2"),
]
policy = SlidingWindowPolicy(max_events=2)
_, result = await policy.apply([], events, context)
assert result == events[3:]
async def test_transparent_count_reflects_dropped_group(self, context: Context) -> None:
events = [
DurableReasoning("plan"),
_call("ws_1"),
_result("ws_1"),
ModelRequest([TextInput("next")]),
]
policy = SlidingWindowPolicy(max_events=3, transparent=True)
prompts, result = await policy.apply([], events, context)
assert len(result) == 1
assert "last 1 of 4" in prompts[-1]
async def test_group_is_kept_whole_when_no_smaller_window_is_legal(self, context: Context) -> None:
# Nothing follows the group, so advancing the cut would leave an empty
# request. Overshooting the window beats sending nothing.
events = [
ModelRequest([TextInput("q")]),
DurableReasoning("plan"),
_call("ws_1"),
_result("ws_1"),
]
policy = SlidingWindowPolicy(max_events=2)
_, result = await policy.apply([], events, context)
assert result == events[1:]
async def test_anchor_is_found_across_interleaved_events(self, context: Context) -> None:
# UsageEvent is persisted but not conversation, so it can sit between a
# reasoning item and the call it anchors. The anchor is still its anchor.
events = [
DurableReasoning("plan"),
UsageEvent(),
_call("ws_1"),
_result("ws_1"),
ModelRequest([TextInput("next")]),
]
policy = SlidingWindowPolicy(max_events=3)
_, result = await policy.apply([], events, context)
assert result == [events[-1]]
async def test_window_widens_past_stray_local_call_events(self, context: Context) -> None:
# A local call is announced by the ModelResponse that requested it; the
# standalone call events map to no provider item. A window holding only
# those maps to an empty request, so it must widen to reach the response.
call = ToolCallEvent(id="c_1", name="convert", arguments="{}")
events = [
ModelRequest([TextInput("q")]),
ModelResponse(tool_calls=ToolCallsEvent(calls=[call])),
ToolCallsEvent(calls=[call]),
call,
ToolResultsEvent(results=[ToolResultEvent(parent_id="c_1", name="convert", result=ToolResult("ok"))]),
]
policy = SlidingWindowPolicy(max_events=3)
_, result = await policy.apply([], events, context)
assert result == events[1:]
async def test_orphaned_builtin_result_is_dropped(self, context: Context) -> None:
# A builtin call never appears in ModelResponse.tool_calls, so its result
# needs the call event itself to survive the cut.
events = [
_call("ws_1"),
_result("ws_1"),
ModelRequest([TextInput("next")]),
]
policy = SlidingWindowPolicy(max_events=2)
_, result = await policy.apply([], events, context)
assert result == [events[-1]]
@pytest.mark.asyncio
class TestProviderTurnItem:
"""A response's tool calls may live only in a provider-native turn object."""
async def test_response_orphaned_from_its_turn_item_is_dropped(self, context: Context) -> None:
# Keeping the response without the turn object rebuilds the turn text-only:
# the tool calls vanish and the results below reference calls the model was
# never told it made. Dropping the response takes the results with it.
events = [
ProviderTurnState(),
ModelResponse(tool_calls=ToolCallsEvent(calls=[ToolCallEvent(id="tc_1", name="multiply")])),
ToolResultsEvent(results=[ToolResultEvent(parent_id="tc_1", name="multiply", result=ToolResult("ok"))]),
ModelRequest([TextInput("next")]),
]
policy = SlidingWindowPolicy(max_events=3)
_, result = await policy.apply([], events, context)
assert result == [events[-1]]
async def test_intact_turn_is_kept(self, context: Context) -> None:
events = [
ModelRequest([TextInput("old")]),
ProviderTurnState(),
ModelResponse(tool_calls=ToolCallsEvent(calls=[ToolCallEvent(id="tc_1", name="multiply")])),
ToolResultsEvent(results=[ToolResultEvent(parent_id="tc_1", name="multiply", result=ToolResult("ok"))]),
]
policy = SlidingWindowPolicy(max_events=3)
_, result = await policy.apply([], events, context)
assert result == events[1:]
async def test_plain_response_survives_without_its_turn_item(self, context: Context) -> None:
# Only tool calls live exclusively in the turn object; the text is on the
# response, so an answer with no tool calls loses nothing.
events = [
ProviderTurnState(),
ModelResponse(message=ModelMessage("hi")),
ModelRequest([TextInput("next")]),
]
policy = SlidingWindowPolicy(max_events=2)
_, result = await policy.apply([], events, context)
assert result == events[1:]
async def test_turn_item_does_not_leak_to_a_later_response(self, context: Context) -> None:
# A turn object belongs to the response that follows it. A later response
# that never had one must not be condemned by the earlier one being cut.
events = [
ProviderTurnState(),
ModelResponse(tool_calls=ToolCallsEvent(calls=[ToolCallEvent(id="tc_1", name="multiply")])),
ModelRequest([TextInput("next")]),
ModelResponse(tool_calls=ToolCallsEvent(calls=[ToolCallEvent(id="tc_2", name="multiply")])),
ToolResultsEvent(results=[ToolResultEvent(parent_id="tc_2", name="multiply", result=ToolResult("ok"))]),
]
policy = SlidingWindowPolicy(max_events=3)
_, result = await policy.apply([], events, context)
assert result == events[2:]
class ReasoningShapedTurnState(ModelReasoning, ProviderReplay):
"""A turn object that happens to subclass ``ModelReasoning``.
Nothing stops a provider from shaping its turn carrier this way, and the two
roles have opposite remedies — so the role has to be read off the event's own
declaration rather than inferred from what it inherits.
"""
__transient__ = False
__replay_role__ = "turn"
def test_marker_requires_a_declared_role() -> None:
with pytest.raises(TypeError, match="__replay_role__"):
class Undeclared(BaseEvent, ProviderReplay):
pass
@pytest.mark.asyncio
async def test_declared_role_decides_the_remedy_not_the_base_class(context: Context) -> None:
# Read as an anchor, this event would condemn builtin calls and leave the
# response standing — rebuilding the turn text-only and losing its tool calls.
events = [
ReasoningShapedTurnState("state"),
ModelResponse(tool_calls=ToolCallsEvent(calls=[ToolCallEvent(id="tc_1", name="multiply")])),
ToolResultsEvent(results=[ToolResultEvent(parent_id="tc_1", name="multiply", result=ToolResult("ok"))]),
ModelRequest([TextInput("next")]),
]
policy = SlidingWindowPolicy(max_events=3)
_, result = await policy.apply([], events, context)
assert result == [events[-1]]
@pytest.mark.asyncio
class TestTokenBudget:
async def test_orphaned_builtin_call_is_dropped(self, context: Context) -> None:
events = [
DurableReasoning("plan"),
_call("ws_1"),
_result("ws_1"),
ModelRequest([TextInput("next")]),
]
policy = TokenBudgetPolicy(max_tokens=_budget_for(events[1:]))
_, result = await policy.apply([], events, context)
assert result == [events[-1]]
async def test_stays_within_budget_after_dropping_the_group(self, context: Context) -> None:
events = [
DurableReasoning("plan"),
_call("ws_1"),
_result("ws_1"),
ModelRequest([TextInput("next")]),
]
budget = _budget_for(events[1:])
policy = TokenBudgetPolicy(max_tokens=budget)
_, result = await policy.apply([], events, context)
assert _chars(result) <= budget * 4
async def test_intact_group_is_kept(self, context: Context) -> None:
events = [
ModelRequest([TextInput("a" * 5000)]),
DurableReasoning("plan"),
_call("ws_1"),
_result("ws_1"),
]
policy = TokenBudgetPolicy(max_tokens=_budget_for(events[1:]))
_, result = await policy.apply([], events, context)
assert result == events[1:]
async def test_budget_is_overshot_when_no_smaller_span_is_legal(self, context: Context) -> None:
# The budget fits only the result, whose call is outside it. Dropping the
# orphan would leave nothing, so the span widens past the budget instead —
# an oversized request is recoverable, an empty one is rejected outright.
events = [
DurableReasoning("plan"),
_call("ws_1"),
_result("ws_1"),
]
budget = _budget_for(events[-1:])
policy = TokenBudgetPolicy(max_tokens=budget)
_, result = await policy.apply([], events, context)
assert result == events
assert _chars(result) > budget * 4