test/middleware/test_conditional.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
from collections.abc import Sequence
from unittest.mock import MagicMock
import pytest
from ag2 import Agent, Context
from ag2.events import (
BaseEvent,
HumanInputRequest,
HumanMessage,
ModelResponse,
ToolCallEvent,
ToolResultEvent,
)
from ag2.middleware import (
AgentTurn,
BaseMiddleware,
ConditionalMiddleware,
HumanInputHook,
LLMCall,
Middleware,
ToolExecution,
)
from ag2.testing import TestConfig
class TrackingMiddleware(BaseMiddleware):
def __init__(
self,
event: BaseEvent,
ctx: Context,
mock: MagicMock,
) -> None:
super().__init__(event, ctx)
self.mock = mock
async def on_turn(
self,
call_next: AgentTurn,
event: BaseEvent,
ctx: Context,
) -> ModelResponse:
self.mock.on_turn(event.__class__.__name__)
return await call_next(event, ctx)
async def on_tool_execution(
self,
call_next: ToolExecution,
event: ToolCallEvent,
ctx: Context,
) -> ToolResultEvent:
self.mock.on_tool(event.name)
return await call_next(event, ctx)
async def on_llm_call(
self,
call_next: LLMCall,
events: Sequence[BaseEvent],
ctx: Context,
) -> ModelResponse:
self.mock.on_llm(len(events))
return await call_next(events, ctx)
async def on_human_input(
self,
call_next: HumanInputHook,
event: HumanInputRequest,
ctx: Context,
) -> HumanMessage:
self.mock.on_human(event.content)
return await call_next(event, ctx)
@pytest.mark.asyncio()
class TestConditionalMiddleware:
async def test_tool_condition_activates_on_matching_tool(self, mock: MagicMock) -> None:
"""The motivating example from #2781: condition on ToolCallEvent.name."""
def my_tool() -> str:
return "done"
agent = Agent(
"",
config=TestConfig(
ToolCallEvent(name="my_tool"),
"result",
),
tools=[my_tool],
middleware=[
ConditionalMiddleware(
Middleware(TrackingMiddleware, mock=mock),
condition=ToolCallEvent.name == "my_tool",
),
],
)
await agent.ask("Hi!")
mock.on_tool.assert_called_once_with("my_tool")
async def test_tool_condition_skips_non_matching_tool(self, mock: MagicMock) -> None:
"""Condition targets a different tool name — middleware should not fire."""
def my_tool() -> str:
return "done"
agent = Agent(
"",
config=TestConfig(
ToolCallEvent(name="my_tool"),
"result",
),
tools=[my_tool],
middleware=[
ConditionalMiddleware(
Middleware(TrackingMiddleware, mock=mock),
condition=ToolCallEvent.name == "other_tool",
),
],
)
await agent.ask("Hi!")
mock.on_tool.assert_not_called()
async def test_type_condition_activates_per_hook(self, mock: MagicMock) -> None:
"""Bare event type as condition: activates on_tool_execution but not on_turn."""
def my_tool() -> str:
return "done"
agent = Agent(
"",
config=TestConfig(
ToolCallEvent(name="my_tool"),
"result",
),
tools=[my_tool],
middleware=[
ConditionalMiddleware(
Middleware(TrackingMiddleware, mock=mock),
condition=ToolCallEvent,
),
],
)
await agent.ask("Hi!")
mock.on_tool.assert_called_once_with("my_tool")
mock.on_turn.assert_not_called()
async def test_llm_call_skipped_when_condition_does_not_match(self, mock: MagicMock) -> None:
"""on_llm_call checks condition against initial event — skipped when it doesn't match."""
agent = Agent(
"",
config=TestConfig("result"),
middleware=[
ConditionalMiddleware(
Middleware(TrackingMiddleware, mock=mock),
condition=ToolCallEvent,
),
],
)
await agent.ask("Hi!")
mock.on_turn.assert_not_called()
mock.on_llm.assert_not_called()
async def test_llm_call_delegates_when_condition_matches(self, mock: MagicMock) -> None:
"""on_llm_call delegates when condition matches the initial event."""
agent = Agent(
"",
config=TestConfig("result"),
middleware=[
ConditionalMiddleware(
Middleware(TrackingMiddleware, mock=mock),
condition=BaseEvent,
),
],
)
await agent.ask("Hi!")
mock.on_turn.assert_called()
mock.on_llm.assert_called()
async def test_inverted_condition(self, mock: MagicMock) -> None:
"""~ToolCallEvent activates on_turn but not on_tool_execution."""
def my_tool() -> str:
return "done"
agent = Agent(
"",
config=TestConfig(
ToolCallEvent(name="my_tool"),
"result",
),
tools=[my_tool],
middleware=[
ConditionalMiddleware(
Middleware(TrackingMiddleware, mock=mock),
condition=~(ToolCallEvent.name == "my_tool"),
),
],
)
await agent.ask("Hi!")
mock.on_turn.assert_called()
mock.on_tool.assert_not_called()
async def test_composed_or_condition(self, mock: MagicMock) -> None:
"""Composed condition from issue: (name == 'a') | (name == 'b')."""
condition = (ToolCallEvent.name == "tool_a") | (ToolCallEvent.name == "tool_b")
def tool_a() -> str:
return "a"
def tool_b() -> str:
return "b"
def tool_c() -> str:
return "c"
agent = Agent(
"",
config=TestConfig(
[
ToolCallEvent(name="tool_a"),
ToolCallEvent(name="tool_b"),
ToolCallEvent(name="tool_c"),
],
"result",
),
tools=[tool_a, tool_b, tool_c],
middleware=[
ConditionalMiddleware(
Middleware(TrackingMiddleware, mock=mock),
condition=condition,
),
],
)
await agent.ask("Hi!")
assert mock.on_tool.call_count == 2
mock.on_tool.assert_any_call("tool_a")
mock.on_tool.assert_any_call("tool_b")
async def test_human_input_activates_when_condition_matches(self, mock: MagicMock) -> None:
"""on_human_input fires when condition matches HumanInputRequest."""
async def my_tool(ctx: Context) -> str:
await ctx.input("Approve?", timeout=1.0)
return "done"
def hitl_hook(event: HumanInputRequest) -> HumanMessage:
return HumanMessage("yes")
agent = Agent(
"",
config=TestConfig(
ToolCallEvent(name="my_tool"),
"result",
),
tools=[my_tool],
hitl_hook=hitl_hook,
middleware=[
ConditionalMiddleware(
Middleware(TrackingMiddleware, mock=mock),
condition=HumanInputRequest,
),
],
)
await agent.ask("Hi!")
mock.on_human.assert_called_once_with("Approve?")
mock.on_turn.assert_not_called()
mock.on_tool.assert_not_called()
async def test_human_input_skipped_when_condition_does_not_match(self, mock: MagicMock) -> None:
"""on_human_input skipped when condition targets a different event type."""
async def my_tool(ctx: Context) -> str:
await ctx.input("Approve?", timeout=1.0)
return "done"
def hitl_hook(event: HumanInputRequest) -> HumanMessage:
return HumanMessage("yes")
agent = Agent(
"",
config=TestConfig(
ToolCallEvent(name="my_tool"),
"result",
),
tools=[my_tool],
hitl_hook=hitl_hook,
middleware=[
ConditionalMiddleware(
Middleware(TrackingMiddleware, mock=mock),
condition=ToolCallEvent.name == "other_tool",
),
],
)
await agent.ask("Hi!")
mock.on_human.assert_not_called()