test/test_aggregate.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
"""Tests for AggregateStrategy, AggregateTrigger, and built-in strategies."""
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from dirty_equals import IsPartialDict
from ag2 import Agent, KnowledgeConfig
from ag2 import Context as Context
from ag2.aggregate import (
AggregateTrigger,
ConversationSummaryAggregate,
WorkingMemoryAggregate,
)
from ag2.events import (
AggregationCompleted,
AggregationFailed,
AggregationStarted,
ModelMessage,
ModelRequest,
ModelResponse,
TextInput,
Usage,
)
from ag2.knowledge import KnowledgePlugin, MemoryKnowledgeStore
from ag2.stream import MemoryStream
from ag2.testing import TestConfig
class TestAggregateTrigger:
def test_defaults(self) -> None:
trigger = AggregateTrigger()
assert trigger.every_n_turns == 0
assert trigger.every_n_events == 0
assert trigger.on_end is False
def test_custom_values(self) -> None:
trigger = AggregateTrigger(every_n_turns=5, every_n_events=50, on_end=True)
assert trigger.every_n_turns == 5
assert trigger.every_n_events == 50
assert trigger.on_end is True
class TestConversationSummaryAggregate:
@pytest.mark.asyncio
async def test_writes_timestamped_summary(self) -> None:
"""Summary files should be prefixed with a timestamp for chronological sorting."""
strategy = ConversationSummaryAggregate(
config=TestConfig(
ModelResponse(
ModelMessage("This conversation covered X and Y."),
usage=Usage(prompt_tokens=100, completion_tokens=50),
)
),
)
store = MemoryKnowledgeStore()
stream = MemoryStream()
ctx = Context(stream=stream)
events = [ModelRequest([TextInput("hello")]), ModelRequest([TextInput("world")])]
await strategy.aggregate(events, ctx, store)
entries = await store.list("/memory/conversations/")
assert len(entries) == 1
filename = entries[0]
# Filename should start with a timestamp: YYYYMMDDTHHmmss_
assert filename[8] == "T" # ISO date separator
assert filename[15] == "_" # separator before stream ID
assert filename.endswith(".md")
@pytest.mark.asyncio
async def test_skips_empty_events(self) -> None:
mock_config = MagicMock()
strategy = ConversationSummaryAggregate(config=mock_config)
store = MemoryKnowledgeStore()
ctx = Context(stream=MemoryStream())
await strategy.aggregate([], ctx, store)
entries = await store.list("/memory/conversations/")
assert len(entries) == 0
@pytest.mark.asyncio
async def test_stores_usage(self) -> None:
strategy = ConversationSummaryAggregate(
config=TestConfig(ModelResponse(ModelMessage("Summary"), usage=Usage(prompt_tokens=200))),
)
store = MemoryKnowledgeStore()
ctx = Context(stream=MemoryStream())
await strategy.aggregate([ModelRequest([TextInput("hi")])], ctx, store)
assert strategy.last_usage == IsPartialDict({"prompt_tokens": 200})
@pytest.mark.asyncio
async def test_chronological_ordering_of_summaries(self) -> None:
"""Multiple summaries should sort chronologically by filename."""
mock_response = MagicMock()
mock_response.content = "Summary"
mock_response.usage = {}
mock_client = AsyncMock(return_value=mock_response)
mock_config = MagicMock()
mock_config.create.return_value = mock_client
strategy = ConversationSummaryAggregate(config=mock_config)
store = MemoryKnowledgeStore()
# Write two summaries with different streams
stream1 = MemoryStream()
ctx1 = Context(stream=stream1)
with patch("ag2.aggregate.datetime") as mock_dt:
mock_dt.now.return_value.strftime.return_value = "20260101T120000"
mock_dt.side_effect = None
await strategy.aggregate([ModelRequest([TextInput("first")])], ctx1, store)
stream2 = MemoryStream()
ctx2 = Context(stream=stream2)
with patch("ag2.aggregate.datetime") as mock_dt:
mock_dt.now.return_value.strftime.return_value = "20260201T120000"
mock_dt.side_effect = None
await strategy.aggregate([ModelRequest([TextInput("second")])], ctx2, store)
entries = await store.list("/memory/conversations/")
assert len(entries) == 2
# Alphabetical sort == chronological sort because of timestamp prefix
assert entries[0] < entries[1]
assert entries[0].startswith("20260101")
assert entries[1].startswith("20260201")
class TestWorkingMemoryAggregate:
@pytest.mark.asyncio
async def test_writes_working_memory(self) -> None:
strategy = WorkingMemoryAggregate(
config=TestConfig(
ModelResponse(ModelMessage("Updated working memory content."), usage=Usage(prompt_tokens=100))
),
)
store = MemoryKnowledgeStore()
ctx = Context(stream=MemoryStream())
await strategy.aggregate([ModelRequest([TextInput("hi")])], ctx, store)
content = await store.read("/memory/working.md")
assert content == "Updated working memory content."
@pytest.mark.asyncio
async def test_merges_with_existing(self) -> None:
"""Should pass existing working memory to LLM for merging."""
mock_response = MagicMock()
mock_response.content = "Merged memory."
mock_response.usage = {}
mock_client = AsyncMock(return_value=mock_response)
mock_config = MagicMock()
mock_config.create.return_value = mock_client
strategy = WorkingMemoryAggregate(config=mock_config)
store = MemoryKnowledgeStore()
await store.write("/memory/working.md", "Existing context about project X.")
ctx = Context(stream=MemoryStream())
await strategy.aggregate([ModelRequest([TextInput("update")])], ctx, store)
# Verify the LLM was called with existing memory in the prompt
call_args = mock_client.call_args
prompt_event = call_args[0][0][0]
assert any(
isinstance(inp, TextInput) and "Existing context about project X." in inp.content
for inp in prompt_event.parts
)
content = await store.read("/memory/working.md")
assert content == "Merged memory."
@pytest.mark.asyncio
async def test_skips_empty_events(self) -> None:
mock_config = MagicMock()
strategy = WorkingMemoryAggregate(config=mock_config)
store = MemoryKnowledgeStore()
ctx = Context(stream=MemoryStream())
await strategy.aggregate([], ctx, store)
assert await store.read("/memory/working.md") is None
@pytest.mark.asyncio
async def test_custom_prompt_template_is_used(self) -> None:
"""A user-supplied prompt= template replaces the default verbatim."""
mock_response = MagicMock()
mock_response.content = "lessons"
mock_response.usage = {}
mock_client = AsyncMock(return_value=mock_response)
mock_config = MagicMock()
mock_config.create.return_value = mock_client
strategy = WorkingMemoryAggregate(
config=mock_config,
prompt=("Track tactics, not facts. Existing notes:\n{existing}\n---\nRound:\n{events}"),
)
store = MemoryKnowledgeStore()
await store.write("/memory/working.md", "prior tactics")
ctx = Context(stream=MemoryStream())
await strategy.aggregate([ModelRequest([TextInput("hi")])], ctx, store)
prompt_event = mock_client.call_args[0][0][0]
rendered = next(inp.content for inp in prompt_event.parts if isinstance(inp, TextInput))
assert "Track tactics, not facts." in rendered
assert "prior tactics" in rendered
# The default journal-style phrasing must not leak in.
assert "Preserve important existing context" not in rendered
@pytest.mark.asyncio
async def test_falls_back_to_existing_on_empty_response(self) -> None:
mock_response = MagicMock()
mock_response.content = "" # LLM returns empty
mock_response.usage = {}
mock_client = AsyncMock(return_value=mock_response)
mock_config = MagicMock()
mock_config.create.return_value = mock_client
strategy = WorkingMemoryAggregate(config=mock_config)
store = MemoryKnowledgeStore()
await store.write("/memory/working.md", "Existing content.")
ctx = Context(stream=MemoryStream())
await strategy.aggregate([ModelRequest([TextInput("update")])], ctx, store)
# Should fall back to existing content when LLM returns empty
content = await store.read("/memory/working.md")
assert content == "Existing content."
class TestConversationSummaryStreamId:
@pytest.mark.asyncio
async def test_filename_uses_full_stream_id(self) -> None:
"""The summary filename must contain the full stream UUID, not a prefix."""
store = MemoryKnowledgeStore()
stream = MemoryStream()
ctx = Context(stream=stream)
full_id = str(stream.id)
mock_response = MagicMock()
mock_response.content = "summary"
mock_response.usage = {}
mock_client = AsyncMock(return_value=mock_response)
mock_config = MagicMock()
mock_config.create.return_value = mock_client
agg = ConversationSummaryAggregate(config=mock_config)
await agg.aggregate([ModelRequest([TextInput("hello")])], ctx, store)
entries = await store.list("/memory/conversations/")
assert len(entries) == 1
assert full_id in entries[0]
class _RecordingAggregate:
"""In-process AggregateStrategy that records calls without LLM traffic."""
def __init__(self) -> None:
self.calls = 0
self.last_usage: dict = {}
async def aggregate(self, events, context, store) -> None:
self.calls += 1
await store.write(f"/memory/runs/{self.calls}.md", "rolled up")
class TestAggregationWiredOnAgent:
"""End-to-end behaviour of the aggregation middleware on an Agent."""
@pytest.mark.asyncio
async def test_on_end_fires_once_per_ask(self) -> None:
store = MemoryKnowledgeStore()
strategy = _RecordingAggregate()
stream = MemoryStream()
completions: list[AggregationCompleted] = []
stream.where(AggregationCompleted).subscribe(lambda e: completions.append(e))
agent = Agent(
"roller",
config=TestConfig(ModelResponse(ModelMessage("done"))),
plugins=[
KnowledgePlugin(
store=store,
aggregate=strategy,
aggregate_trigger=AggregateTrigger(on_end=True),
)
],
)
await agent.ask("go", stream=stream)
assert strategy.calls == 1
assert len(completions) == 1
assert completions[0].agent == "roller"
@pytest.mark.asyncio
async def test_completion_reports_the_summarization_usage_as_a_dict(self) -> None:
stream = MemoryStream()
completions: list[AggregationCompleted] = []
stream.where(AggregationCompleted).subscribe(lambda e: completions.append(e))
agent = Agent(
"roller",
config=TestConfig(ModelResponse(ModelMessage("done"))),
plugins=[
KnowledgePlugin(
store=MemoryKnowledgeStore(),
aggregate=ConversationSummaryAggregate(
config=TestConfig(ModelResponse(ModelMessage("Summary"), usage=Usage(total_tokens=30))),
),
aggregate_trigger=AggregateTrigger(on_end=True),
)
],
)
await agent.ask("go", stream=stream)
assert [(c.llm_calls, c.usage) for c in completions] == [(1, IsPartialDict({"total_tokens": 30}))]
@pytest.mark.asyncio
@pytest.mark.filterwarnings(r"ignore:Agent\(knowledge=:DeprecationWarning")
async def test_legacy_knowledge_config_aggregates_once_per_ask(self) -> None:
strategy = _RecordingAggregate()
stream = MemoryStream()
completions: list[AggregationCompleted] = []
stream.where(AggregationCompleted).subscribe(lambda e: completions.append(e))
agent = Agent(
"roller",
config=TestConfig(ModelResponse(ModelMessage("done"))),
knowledge=KnowledgeConfig(
store=MemoryKnowledgeStore(),
aggregate=strategy,
aggregate_trigger=AggregateTrigger(on_end=True),
),
)
await agent.ask("go", stream=stream)
assert strategy.calls == 1
assert len(completions) == 1
assert completions[0].agent == "roller"
@pytest.mark.asyncio
async def test_on_end_does_not_double_fire_with_other_triggers(self) -> None:
"""``on_end=True`` plus ``every_n_turns=1`` still aggregates once per ask."""
store = MemoryKnowledgeStore()
strategy = _RecordingAggregate()
agent = Agent(
"once",
config=TestConfig(ModelResponse(ModelMessage("ok"))),
plugins=[
KnowledgePlugin(
store=store,
aggregate=strategy,
aggregate_trigger=AggregateTrigger(every_n_turns=1, on_end=True),
)
],
)
await agent.ask("go")
assert strategy.calls == 1
@pytest.mark.asyncio
async def test_every_n_events_counts_work_events_not_telemetry(self) -> None:
"""every_n_events must ignore telemetry (UsageEvent) on the stream.
A usage-carrying turn persists ModelRequest + UsageEvent + ModelResponse
(3 events) but only 2 are conversational. With every_n_events=3 the
crossings are at 3 and 6 -> asks 2 and 3 fire, not every ask.
"""
store = MemoryKnowledgeStore()
strategy = _RecordingAggregate()
stream = MemoryStream()
completions: list[AggregationCompleted] = []
stream.where(AggregationCompleted).subscribe(lambda e: completions.append(e))
usage = Usage(prompt_tokens=5, completion_tokens=3, total_tokens=8)
agent = Agent(
"counter",
config=TestConfig(*(ModelResponse(ModelMessage("ok"), usage=usage) for _ in range(4))),
plugins=[
KnowledgePlugin(
store=store,
aggregate=strategy,
aggregate_trigger=AggregateTrigger(every_n_events=3, on_end=False),
)
],
)
r = await agent.ask("a", stream=stream)
r = await r.ask("b")
r = await r.ask("c")
await r.ask("d")
assert strategy.calls == 2
assert len(completions) == 2
@pytest.mark.asyncio
async def test_on_end_runs_even_when_turn_raises(self) -> None:
"""If the LLM call errors, ``on_end`` aggregation still fires.
TestConfig with no canned responses raises on the first call,
simulating a failed turn.
"""
store = MemoryKnowledgeStore()
strategy = _RecordingAggregate()
agent = Agent(
"boom",
config=TestConfig(),
plugins=[
KnowledgePlugin(
store=store,
aggregate=strategy,
aggregate_trigger=AggregateTrigger(on_end=True),
)
],
)
with pytest.raises(Exception):
await agent.ask("go")
assert strategy.calls == 1
class _RaisingAggregate:
"""AggregateStrategy that always raises — for failure-path tests."""
last_usage: dict = {}
async def aggregate(self, events, context, store) -> None:
raise RuntimeError("aggregate boom")
@pytest.mark.asyncio
class TestAggregationLifecycleEvents:
"""Started + Failed events must reach the stream so failures are observable
without configuring Python logging."""
async def test_started_event_fires_before_strategy_runs(self) -> None:
store = MemoryKnowledgeStore()
stream = MemoryStream()
started: list[AggregationStarted] = []
stream.where(AggregationStarted).subscribe(lambda e: started.append(e))
agent = Agent(
"lifecycle",
config=TestConfig(ModelResponse(ModelMessage("ok"))),
plugins=[
KnowledgePlugin(
store=store,
aggregate=_RecordingAggregate(),
aggregate_trigger=AggregateTrigger(on_end=True),
)
],
)
await agent.ask("go", stream=stream)
assert len(started) == 1
assert started[0].agent == "lifecycle"
assert started[0].strategy == "_RecordingAggregate"
async def test_failed_event_fires_when_strategy_raises(self) -> None:
store = MemoryKnowledgeStore()
stream = MemoryStream()
failures: list[AggregationFailed] = []
completions: list[AggregationCompleted] = []
stream.where(AggregationFailed).subscribe(lambda e: failures.append(e))
stream.where(AggregationCompleted).subscribe(lambda e: completions.append(e))
agent = Agent(
"broken-aggregator",
config=TestConfig(ModelResponse(ModelMessage("ok"))),
plugins=[
KnowledgePlugin(
store=store,
aggregate=_RaisingAggregate(),
aggregate_trigger=AggregateTrigger(on_end=True),
)
],
)
# The turn itself must succeed — only the aggregation failed.
await agent.ask("go", stream=stream)
assert len(failures) == 1
assert failures[0].agent == "broken-aggregator"
assert failures[0].strategy == "_RaisingAggregate"
assert failures[0].error_type == "RuntimeError"
assert "aggregate boom" in failures[0].error
assert completions == []