mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: Agent and Function middleware (#770)
* Initial middleware implementation
* Small fixes
* Small updates
* Small updates in samples
* Moved middleware functionality to decorator
* Removed obsolete file
* Renamed AgentInvocationContext to AzureRunContext
* Added unit tests
* Small settings update for test discovery in VS Code
* Added unit tests
* Reverted changes in environment settings
* Added context result override
* Renaming and updates to logic
* Added more samples
* Updated DEV_SETUP.md
* Addressed PR feedback
* Addressed PR feedback
* Removed unused parameter
* Small fix
* Small fix in telemetry logic
* Revert "Small fix in telemetry logic"
This reverts commit 6f82660d2d.
* Small fix
---------
Co-authored-by: Chris <66376200+crickman@users.noreply.github.com>
This commit is contained in:
Dmytro Struk
committed by
GitHub
Unverified
GitHub
parent
538be4c149
commit
99860a5d07
@@ -105,6 +105,9 @@ class MockChatClient:
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.additional_properties: dict[str, Any] = {}
|
||||
self.call_count: int = 0
|
||||
self.responses: list[ChatResponse] = []
|
||||
self.streaming_responses: list[list[ChatResponseUpdate]] = []
|
||||
|
||||
async def get_response(
|
||||
self,
|
||||
@@ -112,6 +115,9 @@ class MockChatClient:
|
||||
**kwargs: Any,
|
||||
) -> ChatResponse:
|
||||
logger.debug(f"Running custom chat client, with: {messages=}, {kwargs=}")
|
||||
self.call_count += 1
|
||||
if self.responses:
|
||||
return self.responses.pop(0)
|
||||
return ChatResponse(messages=ChatMessage(role="assistant", text="test response"))
|
||||
|
||||
async def get_streaming_response(
|
||||
@@ -120,8 +126,13 @@ class MockChatClient:
|
||||
**kwargs: Any,
|
||||
) -> AsyncIterable[ChatResponseUpdate]:
|
||||
logger.debug(f"Running custom chat client stream, with: {messages=}, {kwargs=}")
|
||||
yield ChatResponseUpdate(text=TextContent(text="test streaming response "), role="assistant")
|
||||
yield ChatResponseUpdate(contents=[TextContent(text="another update")], role="assistant")
|
||||
self.call_count += 1
|
||||
if self.streaming_responses:
|
||||
for update in self.streaming_responses.pop(0):
|
||||
yield update
|
||||
else:
|
||||
yield ChatResponseUpdate(text=TextContent(text="test streaming response "), role="assistant")
|
||||
yield ChatResponseUpdate(contents=[TextContent(text="another update")], role="assistant")
|
||||
|
||||
|
||||
class MockBaseChatClient(BaseChatClient):
|
||||
|
||||
@@ -0,0 +1,939 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from collections.abc import AsyncIterable, Awaitable, Callable
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from agent_framework import (
|
||||
AgentProtocol,
|
||||
AgentRunResponse,
|
||||
AgentRunResponseUpdate,
|
||||
ChatMessage,
|
||||
Role,
|
||||
TextContent,
|
||||
)
|
||||
from agent_framework._middleware import (
|
||||
AgentMiddleware,
|
||||
AgentMiddlewarePipeline,
|
||||
AgentRunContext,
|
||||
FunctionInvocationContext,
|
||||
FunctionMiddleware,
|
||||
FunctionMiddlewarePipeline,
|
||||
)
|
||||
from agent_framework._tools import AIFunction
|
||||
|
||||
|
||||
class TestAgentRunContext:
|
||||
"""Test cases for AgentRunContext."""
|
||||
|
||||
def test_init_with_defaults(self, mock_agent: AgentProtocol) -> None:
|
||||
"""Test AgentRunContext initialization with default values."""
|
||||
messages = [ChatMessage(role=Role.USER, text="test")]
|
||||
context = AgentRunContext(agent=mock_agent, messages=messages)
|
||||
|
||||
assert context.agent is mock_agent
|
||||
assert context.messages == messages
|
||||
assert context.is_streaming is False
|
||||
assert context.metadata == {}
|
||||
|
||||
def test_init_with_custom_values(self, mock_agent: AgentProtocol) -> None:
|
||||
"""Test AgentRunContext initialization with custom values."""
|
||||
messages = [ChatMessage(role=Role.USER, text="test")]
|
||||
metadata = {"key": "value"}
|
||||
context = AgentRunContext(agent=mock_agent, messages=messages, is_streaming=True, metadata=metadata)
|
||||
|
||||
assert context.agent is mock_agent
|
||||
assert context.messages == messages
|
||||
assert context.is_streaming is True
|
||||
assert context.metadata == metadata
|
||||
|
||||
|
||||
class TestFunctionInvocationContext:
|
||||
"""Test cases for FunctionInvocationContext."""
|
||||
|
||||
def test_init_with_defaults(self, mock_function: AIFunction[Any, Any]) -> None:
|
||||
"""Test FunctionInvocationContext initialization with default values."""
|
||||
arguments = FunctionTestArgs(name="test")
|
||||
context = FunctionInvocationContext(function=mock_function, arguments=arguments)
|
||||
|
||||
assert context.function is mock_function
|
||||
assert context.arguments == arguments
|
||||
assert context.metadata == {}
|
||||
|
||||
def test_init_with_custom_metadata(self, mock_function: AIFunction[Any, Any]) -> None:
|
||||
"""Test FunctionInvocationContext initialization with custom metadata."""
|
||||
arguments = FunctionTestArgs(name="test")
|
||||
metadata = {"key": "value"}
|
||||
context = FunctionInvocationContext(function=mock_function, arguments=arguments, metadata=metadata)
|
||||
|
||||
assert context.function is mock_function
|
||||
assert context.arguments == arguments
|
||||
assert context.metadata == metadata
|
||||
|
||||
|
||||
class TestAgentMiddlewarePipeline:
|
||||
"""Test cases for AgentMiddlewarePipeline."""
|
||||
|
||||
def test_init_empty(self) -> None:
|
||||
"""Test AgentMiddlewarePipeline initialization with no middlewares."""
|
||||
pipeline = AgentMiddlewarePipeline()
|
||||
assert not pipeline.has_middlewares
|
||||
|
||||
def test_init_with_class_middleware(self) -> None:
|
||||
"""Test AgentMiddlewarePipeline initialization with class-based middleware."""
|
||||
middleware = TestAgentMiddleware()
|
||||
pipeline = AgentMiddlewarePipeline([middleware])
|
||||
assert pipeline.has_middlewares
|
||||
|
||||
def test_init_with_function_middleware(self) -> None:
|
||||
"""Test AgentMiddlewarePipeline initialization with function-based middleware."""
|
||||
|
||||
async def test_middleware(context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]) -> None:
|
||||
await next(context)
|
||||
|
||||
pipeline = AgentMiddlewarePipeline([test_middleware])
|
||||
assert pipeline.has_middlewares
|
||||
|
||||
async def test_execute_no_middleware(self, mock_agent: AgentProtocol) -> None:
|
||||
"""Test pipeline execution with no middleware."""
|
||||
pipeline = AgentMiddlewarePipeline()
|
||||
messages = [ChatMessage(role=Role.USER, text="test")]
|
||||
context = AgentRunContext(agent=mock_agent, messages=messages)
|
||||
|
||||
expected_response = AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
|
||||
|
||||
async def final_handler(ctx: AgentRunContext) -> AgentRunResponse:
|
||||
return expected_response
|
||||
|
||||
result = await pipeline.execute(mock_agent, messages, context, final_handler)
|
||||
assert result == expected_response
|
||||
|
||||
async def test_execute_with_middleware(self, mock_agent: AgentProtocol) -> None:
|
||||
"""Test pipeline execution with middleware."""
|
||||
execution_order: list[str] = []
|
||||
|
||||
class OrderTrackingMiddleware(AgentMiddleware):
|
||||
def __init__(self, name: str):
|
||||
self.name = name
|
||||
|
||||
async def process(
|
||||
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
execution_order.append(f"{self.name}_before")
|
||||
await next(context)
|
||||
execution_order.append(f"{self.name}_after")
|
||||
|
||||
middleware = OrderTrackingMiddleware("test")
|
||||
pipeline = AgentMiddlewarePipeline([middleware])
|
||||
messages = [ChatMessage(role=Role.USER, text="test")]
|
||||
context = AgentRunContext(agent=mock_agent, messages=messages)
|
||||
|
||||
expected_response = AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
|
||||
|
||||
async def final_handler(ctx: AgentRunContext) -> AgentRunResponse:
|
||||
execution_order.append("handler")
|
||||
return expected_response
|
||||
|
||||
result = await pipeline.execute(mock_agent, messages, context, final_handler)
|
||||
assert result == expected_response
|
||||
assert execution_order == ["test_before", "handler", "test_after"]
|
||||
|
||||
async def test_execute_stream_no_middleware(self, mock_agent: AgentProtocol) -> None:
|
||||
"""Test pipeline streaming execution with no middleware."""
|
||||
pipeline = AgentMiddlewarePipeline()
|
||||
messages = [ChatMessage(role=Role.USER, text="test")]
|
||||
context = AgentRunContext(agent=mock_agent, messages=messages)
|
||||
|
||||
async def final_handler(ctx: AgentRunContext) -> AsyncIterable[AgentRunResponseUpdate]:
|
||||
yield AgentRunResponseUpdate(contents=[TextContent(text="chunk1")])
|
||||
yield AgentRunResponseUpdate(contents=[TextContent(text="chunk2")])
|
||||
|
||||
updates: list[AgentRunResponseUpdate] = []
|
||||
async for update in pipeline.execute_stream(mock_agent, messages, context, final_handler):
|
||||
updates.append(update)
|
||||
|
||||
assert len(updates) == 2
|
||||
assert updates[0].text == "chunk1"
|
||||
assert updates[1].text == "chunk2"
|
||||
|
||||
async def test_execute_stream_with_middleware(self, mock_agent: AgentProtocol) -> None:
|
||||
"""Test pipeline streaming execution with middleware."""
|
||||
execution_order: list[str] = []
|
||||
|
||||
class StreamOrderTrackingMiddleware(AgentMiddleware):
|
||||
def __init__(self, name: str):
|
||||
self.name = name
|
||||
|
||||
async def process(
|
||||
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
execution_order.append(f"{self.name}_before")
|
||||
await next(context)
|
||||
execution_order.append(f"{self.name}_after")
|
||||
|
||||
middleware = StreamOrderTrackingMiddleware("test")
|
||||
pipeline = AgentMiddlewarePipeline([middleware])
|
||||
messages = [ChatMessage(role=Role.USER, text="test")]
|
||||
context = AgentRunContext(agent=mock_agent, messages=messages)
|
||||
|
||||
async def final_handler(ctx: AgentRunContext) -> AsyncIterable[AgentRunResponseUpdate]:
|
||||
execution_order.append("handler_start")
|
||||
yield AgentRunResponseUpdate(contents=[TextContent(text="chunk1")])
|
||||
yield AgentRunResponseUpdate(contents=[TextContent(text="chunk2")])
|
||||
execution_order.append("handler_end")
|
||||
|
||||
updates: list[AgentRunResponseUpdate] = []
|
||||
async for update in pipeline.execute_stream(mock_agent, messages, context, final_handler):
|
||||
updates.append(update)
|
||||
|
||||
assert len(updates) == 2
|
||||
assert updates[0].text == "chunk1"
|
||||
assert updates[1].text == "chunk2"
|
||||
assert execution_order == ["test_before", "test_after", "handler_start", "handler_end"]
|
||||
|
||||
|
||||
class TestFunctionMiddlewarePipeline:
|
||||
"""Test cases for FunctionMiddlewarePipeline."""
|
||||
|
||||
def test_init_empty(self) -> None:
|
||||
"""Test FunctionMiddlewarePipeline initialization with no middlewares."""
|
||||
pipeline = FunctionMiddlewarePipeline()
|
||||
assert not pipeline.has_middlewares
|
||||
|
||||
def test_init_with_class_middleware(self) -> None:
|
||||
"""Test FunctionMiddlewarePipeline initialization with class-based middleware."""
|
||||
middleware = TestFunctionMiddleware()
|
||||
pipeline = FunctionMiddlewarePipeline([middleware])
|
||||
assert pipeline.has_middlewares
|
||||
|
||||
def test_init_with_function_middleware(self) -> None:
|
||||
"""Test FunctionMiddlewarePipeline initialization with function-based middleware."""
|
||||
|
||||
async def test_middleware(
|
||||
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
|
||||
) -> None:
|
||||
await next(context)
|
||||
|
||||
pipeline = FunctionMiddlewarePipeline([test_middleware])
|
||||
assert pipeline.has_middlewares
|
||||
|
||||
async def test_execute_no_middleware(self, mock_function: AIFunction[Any, Any]) -> None:
|
||||
"""Test pipeline execution with no middleware."""
|
||||
pipeline = FunctionMiddlewarePipeline()
|
||||
arguments = FunctionTestArgs(name="test")
|
||||
context = FunctionInvocationContext(function=mock_function, arguments=arguments)
|
||||
|
||||
expected_result = "function_result"
|
||||
|
||||
async def final_handler(ctx: FunctionInvocationContext) -> str:
|
||||
return expected_result
|
||||
|
||||
result = await pipeline.execute(mock_function, arguments, context, final_handler)
|
||||
assert result == expected_result
|
||||
|
||||
async def test_execute_with_middleware(self, mock_function: AIFunction[Any, Any]) -> None:
|
||||
"""Test pipeline execution with middleware."""
|
||||
execution_order: list[str] = []
|
||||
|
||||
class OrderTrackingFunctionMiddleware(FunctionMiddleware):
|
||||
def __init__(self, name: str):
|
||||
self.name = name
|
||||
|
||||
async def process(
|
||||
self,
|
||||
context: FunctionInvocationContext,
|
||||
next: Callable[[FunctionInvocationContext], Awaitable[None]],
|
||||
) -> None:
|
||||
execution_order.append(f"{self.name}_before")
|
||||
await next(context)
|
||||
execution_order.append(f"{self.name}_after")
|
||||
|
||||
middleware = OrderTrackingFunctionMiddleware("test")
|
||||
pipeline = FunctionMiddlewarePipeline([middleware])
|
||||
arguments = FunctionTestArgs(name="test")
|
||||
context = FunctionInvocationContext(function=mock_function, arguments=arguments)
|
||||
|
||||
expected_result = "function_result"
|
||||
|
||||
async def final_handler(ctx: FunctionInvocationContext) -> str:
|
||||
execution_order.append("handler")
|
||||
return expected_result
|
||||
|
||||
result = await pipeline.execute(mock_function, arguments, context, final_handler)
|
||||
assert result == expected_result
|
||||
assert execution_order == ["test_before", "handler", "test_after"]
|
||||
|
||||
|
||||
class TestClassBasedMiddleware:
|
||||
"""Test cases for class-based middleware implementations."""
|
||||
|
||||
async def test_agent_middleware_execution(self, mock_agent: AgentProtocol) -> None:
|
||||
"""Test class-based agent middleware execution."""
|
||||
metadata_updates: list[str] = []
|
||||
|
||||
class MetadataAgentMiddleware(AgentMiddleware):
|
||||
async def process(
|
||||
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
context.metadata["before"] = True
|
||||
metadata_updates.append("before")
|
||||
await next(context)
|
||||
context.metadata["after"] = True
|
||||
metadata_updates.append("after")
|
||||
|
||||
middleware = MetadataAgentMiddleware()
|
||||
pipeline = AgentMiddlewarePipeline([middleware])
|
||||
messages = [ChatMessage(role=Role.USER, text="test")]
|
||||
context = AgentRunContext(agent=mock_agent, messages=messages)
|
||||
|
||||
async def final_handler(ctx: AgentRunContext) -> AgentRunResponse:
|
||||
metadata_updates.append("handler")
|
||||
return AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
|
||||
|
||||
result = await pipeline.execute(mock_agent, messages, context, final_handler)
|
||||
|
||||
assert result is not None
|
||||
assert context.metadata["before"] is True
|
||||
assert context.metadata["after"] is True
|
||||
assert metadata_updates == ["before", "handler", "after"]
|
||||
|
||||
async def test_function_middleware_execution(self, mock_function: AIFunction[Any, Any]) -> None:
|
||||
"""Test class-based function middleware execution."""
|
||||
metadata_updates: list[str] = []
|
||||
|
||||
class MetadataFunctionMiddleware(FunctionMiddleware):
|
||||
async def process(
|
||||
self,
|
||||
context: FunctionInvocationContext,
|
||||
next: Callable[[FunctionInvocationContext], Awaitable[None]],
|
||||
) -> None:
|
||||
context.metadata["before"] = True
|
||||
metadata_updates.append("before")
|
||||
await next(context)
|
||||
context.metadata["after"] = True
|
||||
metadata_updates.append("after")
|
||||
|
||||
middleware = MetadataFunctionMiddleware()
|
||||
pipeline = FunctionMiddlewarePipeline([middleware])
|
||||
arguments = FunctionTestArgs(name="test")
|
||||
context = FunctionInvocationContext(function=mock_function, arguments=arguments)
|
||||
|
||||
async def final_handler(ctx: FunctionInvocationContext) -> str:
|
||||
metadata_updates.append("handler")
|
||||
return "result"
|
||||
|
||||
result = await pipeline.execute(mock_function, arguments, context, final_handler)
|
||||
|
||||
assert result == "result"
|
||||
assert context.metadata["before"] is True
|
||||
assert context.metadata["after"] is True
|
||||
assert metadata_updates == ["before", "handler", "after"]
|
||||
|
||||
|
||||
class TestFunctionBasedMiddleware:
|
||||
"""Test cases for function-based middleware implementations."""
|
||||
|
||||
async def test_agent_function_middleware(self, mock_agent: AgentProtocol) -> None:
|
||||
"""Test function-based agent middleware."""
|
||||
execution_order: list[str] = []
|
||||
|
||||
async def test_agent_middleware(
|
||||
context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
execution_order.append("function_before")
|
||||
context.metadata["function_middleware"] = True
|
||||
await next(context)
|
||||
execution_order.append("function_after")
|
||||
|
||||
pipeline = AgentMiddlewarePipeline([test_agent_middleware])
|
||||
messages = [ChatMessage(role=Role.USER, text="test")]
|
||||
context = AgentRunContext(agent=mock_agent, messages=messages)
|
||||
|
||||
async def final_handler(ctx: AgentRunContext) -> AgentRunResponse:
|
||||
execution_order.append("handler")
|
||||
return AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
|
||||
|
||||
result = await pipeline.execute(mock_agent, messages, context, final_handler)
|
||||
|
||||
assert result is not None
|
||||
assert context.metadata["function_middleware"] is True
|
||||
assert execution_order == ["function_before", "handler", "function_after"]
|
||||
|
||||
async def test_function_function_middleware(self, mock_function: AIFunction[Any, Any]) -> None:
|
||||
"""Test function-based function middleware."""
|
||||
execution_order: list[str] = []
|
||||
|
||||
async def test_function_middleware(
|
||||
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
|
||||
) -> None:
|
||||
execution_order.append("function_before")
|
||||
context.metadata["function_middleware"] = True
|
||||
await next(context)
|
||||
execution_order.append("function_after")
|
||||
|
||||
pipeline = FunctionMiddlewarePipeline([test_function_middleware])
|
||||
arguments = FunctionTestArgs(name="test")
|
||||
context = FunctionInvocationContext(function=mock_function, arguments=arguments)
|
||||
|
||||
async def final_handler(ctx: FunctionInvocationContext) -> str:
|
||||
execution_order.append("handler")
|
||||
return "result"
|
||||
|
||||
result = await pipeline.execute(mock_function, arguments, context, final_handler)
|
||||
|
||||
assert result == "result"
|
||||
assert context.metadata["function_middleware"] is True
|
||||
assert execution_order == ["function_before", "handler", "function_after"]
|
||||
|
||||
|
||||
class TestMixedMiddleware:
|
||||
"""Test cases for mixed class and function-based middleware."""
|
||||
|
||||
async def test_mixed_agent_middleware(self, mock_agent: AgentProtocol) -> None:
|
||||
"""Test mixed class and function-based agent middleware."""
|
||||
execution_order: list[str] = []
|
||||
|
||||
class ClassMiddleware(AgentMiddleware):
|
||||
async def process(
|
||||
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
execution_order.append("class_before")
|
||||
await next(context)
|
||||
execution_order.append("class_after")
|
||||
|
||||
async def function_middleware(
|
||||
context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
execution_order.append("function_before")
|
||||
await next(context)
|
||||
execution_order.append("function_after")
|
||||
|
||||
pipeline = AgentMiddlewarePipeline([ClassMiddleware(), function_middleware])
|
||||
messages = [ChatMessage(role=Role.USER, text="test")]
|
||||
context = AgentRunContext(agent=mock_agent, messages=messages)
|
||||
|
||||
async def final_handler(ctx: AgentRunContext) -> AgentRunResponse:
|
||||
execution_order.append("handler")
|
||||
return AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
|
||||
|
||||
result = await pipeline.execute(mock_agent, messages, context, final_handler)
|
||||
|
||||
assert result is not None
|
||||
assert execution_order == ["class_before", "function_before", "handler", "function_after", "class_after"]
|
||||
|
||||
async def test_mixed_function_middleware(self, mock_function: AIFunction[Any, Any]) -> None:
|
||||
"""Test mixed class and function-based function middleware."""
|
||||
execution_order: list[str] = []
|
||||
|
||||
class ClassMiddleware(FunctionMiddleware):
|
||||
async def process(
|
||||
self,
|
||||
context: FunctionInvocationContext,
|
||||
next: Callable[[FunctionInvocationContext], Awaitable[None]],
|
||||
) -> None:
|
||||
execution_order.append("class_before")
|
||||
await next(context)
|
||||
execution_order.append("class_after")
|
||||
|
||||
async def function_middleware(
|
||||
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
|
||||
) -> None:
|
||||
execution_order.append("function_before")
|
||||
await next(context)
|
||||
execution_order.append("function_after")
|
||||
|
||||
pipeline = FunctionMiddlewarePipeline([ClassMiddleware(), function_middleware])
|
||||
arguments = FunctionTestArgs(name="test")
|
||||
context = FunctionInvocationContext(function=mock_function, arguments=arguments)
|
||||
|
||||
async def final_handler(ctx: FunctionInvocationContext) -> str:
|
||||
execution_order.append("handler")
|
||||
return "result"
|
||||
|
||||
result = await pipeline.execute(mock_function, arguments, context, final_handler)
|
||||
|
||||
assert result == "result"
|
||||
assert execution_order == ["class_before", "function_before", "handler", "function_after", "class_after"]
|
||||
|
||||
|
||||
class TestMultipleMiddlewareOrdering:
|
||||
"""Test cases for multiple middleware execution order."""
|
||||
|
||||
async def test_agent_middleware_execution_order(self, mock_agent: AgentProtocol) -> None:
|
||||
"""Test that multiple agent middlewares execute in registration order."""
|
||||
execution_order: list[str] = []
|
||||
|
||||
class FirstMiddleware(AgentMiddleware):
|
||||
async def process(
|
||||
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
execution_order.append("first_before")
|
||||
await next(context)
|
||||
execution_order.append("first_after")
|
||||
|
||||
class SecondMiddleware(AgentMiddleware):
|
||||
async def process(
|
||||
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
execution_order.append("second_before")
|
||||
await next(context)
|
||||
execution_order.append("second_after")
|
||||
|
||||
class ThirdMiddleware(AgentMiddleware):
|
||||
async def process(
|
||||
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
execution_order.append("third_before")
|
||||
await next(context)
|
||||
execution_order.append("third_after")
|
||||
|
||||
middlewares = [FirstMiddleware(), SecondMiddleware(), ThirdMiddleware()]
|
||||
pipeline = AgentMiddlewarePipeline(middlewares) # type: ignore
|
||||
messages = [ChatMessage(role=Role.USER, text="test")]
|
||||
context = AgentRunContext(agent=mock_agent, messages=messages)
|
||||
|
||||
async def final_handler(ctx: AgentRunContext) -> AgentRunResponse:
|
||||
execution_order.append("handler")
|
||||
return AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
|
||||
|
||||
result = await pipeline.execute(mock_agent, messages, context, final_handler)
|
||||
|
||||
assert result is not None
|
||||
expected_order = [
|
||||
"first_before",
|
||||
"second_before",
|
||||
"third_before",
|
||||
"handler",
|
||||
"third_after",
|
||||
"second_after",
|
||||
"first_after",
|
||||
]
|
||||
assert execution_order == expected_order
|
||||
|
||||
async def test_function_middleware_execution_order(self, mock_function: AIFunction[Any, Any]) -> None:
|
||||
"""Test that multiple function middlewares execute in registration order."""
|
||||
execution_order: list[str] = []
|
||||
|
||||
class FirstMiddleware(FunctionMiddleware):
|
||||
async def process(
|
||||
self,
|
||||
context: FunctionInvocationContext,
|
||||
next: Callable[[FunctionInvocationContext], Awaitable[None]],
|
||||
) -> None:
|
||||
execution_order.append("first_before")
|
||||
await next(context)
|
||||
execution_order.append("first_after")
|
||||
|
||||
class SecondMiddleware(FunctionMiddleware):
|
||||
async def process(
|
||||
self,
|
||||
context: FunctionInvocationContext,
|
||||
next: Callable[[FunctionInvocationContext], Awaitable[None]],
|
||||
) -> None:
|
||||
execution_order.append("second_before")
|
||||
await next(context)
|
||||
execution_order.append("second_after")
|
||||
|
||||
middlewares = [FirstMiddleware(), SecondMiddleware()]
|
||||
pipeline = FunctionMiddlewarePipeline(middlewares) # type: ignore
|
||||
arguments = FunctionTestArgs(name="test")
|
||||
context = FunctionInvocationContext(function=mock_function, arguments=arguments)
|
||||
|
||||
async def final_handler(ctx: FunctionInvocationContext) -> str:
|
||||
execution_order.append("handler")
|
||||
return "result"
|
||||
|
||||
result = await pipeline.execute(mock_function, arguments, context, final_handler)
|
||||
|
||||
assert result == "result"
|
||||
expected_order = ["first_before", "second_before", "handler", "second_after", "first_after"]
|
||||
assert execution_order == expected_order
|
||||
|
||||
|
||||
class TestContextContentValidation:
|
||||
"""Test cases for validating middleware context content."""
|
||||
|
||||
async def test_agent_context_validation(self, mock_agent: AgentProtocol) -> None:
|
||||
"""Test that agent context contains expected data."""
|
||||
|
||||
class ContextValidationMiddleware(AgentMiddleware):
|
||||
async def process(
|
||||
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
# Verify context has all expected attributes
|
||||
assert hasattr(context, "agent")
|
||||
assert hasattr(context, "messages")
|
||||
assert hasattr(context, "is_streaming")
|
||||
assert hasattr(context, "metadata")
|
||||
|
||||
# Verify context content
|
||||
assert context.agent is mock_agent
|
||||
assert len(context.messages) == 1
|
||||
assert context.messages[0].role == Role.USER
|
||||
assert context.messages[0].text == "test"
|
||||
assert context.is_streaming is False
|
||||
assert isinstance(context.metadata, dict)
|
||||
|
||||
# Add custom metadata
|
||||
context.metadata["validated"] = True
|
||||
|
||||
await next(context)
|
||||
|
||||
middleware = ContextValidationMiddleware()
|
||||
pipeline = AgentMiddlewarePipeline([middleware])
|
||||
messages = [ChatMessage(role=Role.USER, text="test")]
|
||||
context = AgentRunContext(agent=mock_agent, messages=messages)
|
||||
|
||||
async def final_handler(ctx: AgentRunContext) -> AgentRunResponse:
|
||||
# Verify metadata was set by middleware
|
||||
assert ctx.metadata.get("validated") is True
|
||||
return AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
|
||||
|
||||
result = await pipeline.execute(mock_agent, messages, context, final_handler)
|
||||
assert result is not None
|
||||
|
||||
async def test_function_context_validation(self, mock_function: AIFunction[Any, Any]) -> None:
|
||||
"""Test that function context contains expected data."""
|
||||
|
||||
class ContextValidationMiddleware(FunctionMiddleware):
|
||||
async def process(
|
||||
self,
|
||||
context: FunctionInvocationContext,
|
||||
next: Callable[[FunctionInvocationContext], Awaitable[None]],
|
||||
) -> None:
|
||||
# Verify context has all expected attributes
|
||||
assert hasattr(context, "function")
|
||||
assert hasattr(context, "arguments")
|
||||
assert hasattr(context, "metadata")
|
||||
|
||||
# Verify context content
|
||||
assert context.function is mock_function
|
||||
assert isinstance(context.arguments, FunctionTestArgs)
|
||||
assert context.arguments.name == "test"
|
||||
assert isinstance(context.metadata, dict)
|
||||
|
||||
# Add custom metadata
|
||||
context.metadata["validated"] = True
|
||||
|
||||
await next(context)
|
||||
|
||||
middleware = ContextValidationMiddleware()
|
||||
pipeline = FunctionMiddlewarePipeline([middleware])
|
||||
arguments = FunctionTestArgs(name="test")
|
||||
context = FunctionInvocationContext(function=mock_function, arguments=arguments)
|
||||
|
||||
async def final_handler(ctx: FunctionInvocationContext) -> str:
|
||||
# Verify metadata was set by middleware
|
||||
assert ctx.metadata.get("validated") is True
|
||||
return "result"
|
||||
|
||||
result = await pipeline.execute(mock_function, arguments, context, final_handler)
|
||||
assert result == "result"
|
||||
|
||||
|
||||
class TestStreamingScenarios:
|
||||
"""Test cases for streaming and non-streaming scenarios."""
|
||||
|
||||
async def test_streaming_flag_validation(self, mock_agent: AgentProtocol) -> None:
|
||||
"""Test that is_streaming flag is correctly set for streaming calls."""
|
||||
streaming_flags: list[bool] = []
|
||||
|
||||
class StreamingFlagMiddleware(AgentMiddleware):
|
||||
async def process(
|
||||
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
streaming_flags.append(context.is_streaming)
|
||||
await next(context)
|
||||
|
||||
middleware = StreamingFlagMiddleware()
|
||||
pipeline = AgentMiddlewarePipeline([middleware])
|
||||
messages = [ChatMessage(role=Role.USER, text="test")]
|
||||
|
||||
# Test non-streaming
|
||||
context = AgentRunContext(agent=mock_agent, messages=messages)
|
||||
|
||||
async def final_handler(ctx: AgentRunContext) -> AgentRunResponse:
|
||||
streaming_flags.append(ctx.is_streaming)
|
||||
return AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response")])
|
||||
|
||||
await pipeline.execute(mock_agent, messages, context, final_handler)
|
||||
|
||||
# Test streaming
|
||||
context_stream = AgentRunContext(agent=mock_agent, messages=messages)
|
||||
|
||||
async def final_stream_handler(ctx: AgentRunContext) -> AsyncIterable[AgentRunResponseUpdate]:
|
||||
streaming_flags.append(ctx.is_streaming)
|
||||
yield AgentRunResponseUpdate(contents=[TextContent(text="chunk")])
|
||||
|
||||
updates: list[AgentRunResponseUpdate] = []
|
||||
async for update in pipeline.execute_stream(mock_agent, messages, context_stream, final_stream_handler):
|
||||
updates.append(update)
|
||||
|
||||
# Verify flags: [non-streaming middleware, non-streaming handler, streaming middleware, streaming handler]
|
||||
assert streaming_flags == [False, False, True, True]
|
||||
|
||||
async def test_streaming_middleware_behavior(self, mock_agent: AgentProtocol) -> None:
|
||||
"""Test middleware behavior with streaming responses."""
|
||||
chunks_processed: list[str] = []
|
||||
|
||||
class StreamProcessingMiddleware(AgentMiddleware):
|
||||
async def process(
|
||||
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
chunks_processed.append("before_stream")
|
||||
await next(context)
|
||||
chunks_processed.append("after_stream")
|
||||
|
||||
middleware = StreamProcessingMiddleware()
|
||||
pipeline = AgentMiddlewarePipeline([middleware])
|
||||
messages = [ChatMessage(role=Role.USER, text="test")]
|
||||
context = AgentRunContext(agent=mock_agent, messages=messages)
|
||||
|
||||
async def final_stream_handler(ctx: AgentRunContext) -> AsyncIterable[AgentRunResponseUpdate]:
|
||||
chunks_processed.append("stream_start")
|
||||
yield AgentRunResponseUpdate(contents=[TextContent(text="chunk1")])
|
||||
chunks_processed.append("chunk1_yielded")
|
||||
yield AgentRunResponseUpdate(contents=[TextContent(text="chunk2")])
|
||||
chunks_processed.append("chunk2_yielded")
|
||||
chunks_processed.append("stream_end")
|
||||
|
||||
updates: list[str] = []
|
||||
async for update in pipeline.execute_stream(mock_agent, messages, context, final_stream_handler):
|
||||
updates.append(update.text)
|
||||
|
||||
assert updates == ["chunk1", "chunk2"]
|
||||
assert chunks_processed == [
|
||||
"before_stream",
|
||||
"after_stream",
|
||||
"stream_start",
|
||||
"chunk1_yielded",
|
||||
"chunk2_yielded",
|
||||
"stream_end",
|
||||
]
|
||||
|
||||
|
||||
# region Helper classes and fixtures
|
||||
|
||||
|
||||
class FunctionTestArgs(BaseModel):
|
||||
"""Test arguments for function middleware tests."""
|
||||
|
||||
name: str = Field(description="Test name parameter")
|
||||
|
||||
|
||||
class TestAgentMiddleware(AgentMiddleware):
|
||||
"""Test implementation of AgentMiddleware."""
|
||||
|
||||
async def process(self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]) -> None:
|
||||
await next(context)
|
||||
|
||||
|
||||
class TestFunctionMiddleware(FunctionMiddleware):
|
||||
"""Test implementation of FunctionMiddleware."""
|
||||
|
||||
async def process(
|
||||
self, context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
|
||||
) -> None:
|
||||
await next(context)
|
||||
|
||||
|
||||
class MockFunctionArgs(BaseModel):
|
||||
"""Test arguments for function middleware tests."""
|
||||
|
||||
name: str = Field(description="Test name parameter")
|
||||
|
||||
|
||||
class TestMiddlewareExecutionControl:
|
||||
"""Test cases for middleware execution control (when next() is called vs not called)."""
|
||||
|
||||
async def test_agent_middleware_no_next_no_execution(self, mock_agent: AgentProtocol) -> None:
|
||||
"""Test that when agent middleware doesn't call next(), no execution happens."""
|
||||
|
||||
class NoNextMiddleware(AgentMiddleware):
|
||||
async def process(
|
||||
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
# Don't call next() - this should prevent any execution
|
||||
pass
|
||||
|
||||
middleware = NoNextMiddleware()
|
||||
pipeline = AgentMiddlewarePipeline([middleware])
|
||||
messages = [ChatMessage(role=Role.USER, text="test")]
|
||||
context = AgentRunContext(agent=mock_agent, messages=messages)
|
||||
|
||||
handler_called = False
|
||||
|
||||
async def final_handler(ctx: AgentRunContext) -> AgentRunResponse:
|
||||
nonlocal handler_called
|
||||
handler_called = True
|
||||
return AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="should not execute")])
|
||||
|
||||
result = await pipeline.execute(mock_agent, messages, context, final_handler)
|
||||
|
||||
# Verify no execution happened - should return empty AgentRunResponse
|
||||
assert result is not None
|
||||
assert isinstance(result, AgentRunResponse)
|
||||
assert result.messages == [] # Empty response
|
||||
assert not handler_called
|
||||
assert context.result is None
|
||||
|
||||
async def test_agent_middleware_no_next_no_streaming_execution(self, mock_agent: AgentProtocol) -> None:
|
||||
"""Test that when agent middleware doesn't call next(), no streaming execution happens."""
|
||||
|
||||
class NoNextStreamingMiddleware(AgentMiddleware):
|
||||
async def process(
|
||||
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
# Don't call next() - this should prevent any execution
|
||||
pass
|
||||
|
||||
middleware = NoNextStreamingMiddleware()
|
||||
pipeline = AgentMiddlewarePipeline([middleware])
|
||||
messages = [ChatMessage(role=Role.USER, text="test")]
|
||||
context = AgentRunContext(agent=mock_agent, messages=messages)
|
||||
|
||||
handler_called = False
|
||||
|
||||
async def final_handler(ctx: AgentRunContext) -> AsyncIterable[AgentRunResponseUpdate]:
|
||||
nonlocal handler_called
|
||||
handler_called = True
|
||||
yield AgentRunResponseUpdate(contents=[TextContent(text="should not execute")])
|
||||
|
||||
# When middleware doesn't call next(), streaming should yield no updates
|
||||
updates: list[AgentRunResponseUpdate] = []
|
||||
async for update in pipeline.execute_stream(mock_agent, messages, context, final_handler):
|
||||
updates.append(update)
|
||||
|
||||
# Verify no execution happened and no updates were yielded
|
||||
assert len(updates) == 0
|
||||
assert not handler_called
|
||||
assert context.result is None
|
||||
|
||||
async def test_function_middleware_no_next_no_execution(self, mock_function: AIFunction[Any, Any]) -> None:
|
||||
"""Test that when function middleware doesn't call next(), no execution happens."""
|
||||
|
||||
class FunctionTestArgs(BaseModel):
|
||||
name: str = Field(description="Test name parameter")
|
||||
|
||||
class NoNextFunctionMiddleware(FunctionMiddleware):
|
||||
async def process(
|
||||
self,
|
||||
context: FunctionInvocationContext,
|
||||
next: Callable[[FunctionInvocationContext], Awaitable[None]],
|
||||
) -> None:
|
||||
# Don't call next() - this should prevent any execution
|
||||
pass
|
||||
|
||||
middleware = NoNextFunctionMiddleware()
|
||||
pipeline = FunctionMiddlewarePipeline([middleware])
|
||||
arguments = FunctionTestArgs(name="test")
|
||||
context = FunctionInvocationContext(function=mock_function, arguments=arguments)
|
||||
|
||||
handler_called = False
|
||||
|
||||
async def final_handler(ctx: FunctionInvocationContext) -> str:
|
||||
nonlocal handler_called
|
||||
handler_called = True
|
||||
return "should not execute"
|
||||
|
||||
result = await pipeline.execute(mock_function, arguments, context, final_handler)
|
||||
|
||||
# Verify no execution happened
|
||||
assert result is None
|
||||
assert not handler_called
|
||||
assert context.result is None
|
||||
|
||||
async def test_multiple_middlewares_early_stop(self, mock_agent: AgentProtocol) -> None:
|
||||
"""Test that when first middleware doesn't call next(), subsequent middlewares are not called."""
|
||||
execution_order: list[str] = []
|
||||
|
||||
class FirstMiddleware(AgentMiddleware):
|
||||
async def process(
|
||||
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
execution_order.append("first")
|
||||
# Don't call next() - this should stop the pipeline
|
||||
|
||||
class SecondMiddleware(AgentMiddleware):
|
||||
async def process(
|
||||
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
execution_order.append("second")
|
||||
await next(context)
|
||||
|
||||
pipeline = AgentMiddlewarePipeline([FirstMiddleware(), SecondMiddleware()])
|
||||
messages = [ChatMessage(role=Role.USER, text="test")]
|
||||
context = AgentRunContext(agent=mock_agent, messages=messages)
|
||||
|
||||
handler_called = False
|
||||
|
||||
async def final_handler(ctx: AgentRunContext) -> AgentRunResponse:
|
||||
nonlocal handler_called
|
||||
handler_called = True
|
||||
return AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="should not execute")])
|
||||
|
||||
result = await pipeline.execute(mock_agent, messages, context, final_handler)
|
||||
|
||||
# Verify only first middleware was called and empty response returned
|
||||
assert execution_order == ["first"]
|
||||
assert result is not None
|
||||
assert isinstance(result, AgentRunResponse)
|
||||
assert result.messages == [] # Empty response
|
||||
assert not handler_called
|
||||
|
||||
async def test_function_middleware_pre_execution_override_with_next(
|
||||
self, mock_function: AIFunction[Any, Any]
|
||||
) -> None:
|
||||
"""Test that function middleware can override result before calling next() - this skips handler execution."""
|
||||
|
||||
class FunctionTestArgs(BaseModel):
|
||||
name: str = Field(description="Test name parameter")
|
||||
|
||||
class PreOverrideFunctionMiddleware(FunctionMiddleware):
|
||||
async def process(
|
||||
self,
|
||||
context: FunctionInvocationContext,
|
||||
next: Callable[[FunctionInvocationContext], Awaitable[None]],
|
||||
) -> None:
|
||||
# Set override first
|
||||
context.result = "pre-override result"
|
||||
# Then call next() to continue middleware pipeline
|
||||
await next(context)
|
||||
|
||||
middleware = PreOverrideFunctionMiddleware()
|
||||
pipeline = FunctionMiddlewarePipeline([middleware])
|
||||
arguments = FunctionTestArgs(name="test")
|
||||
context = FunctionInvocationContext(function=mock_function, arguments=arguments)
|
||||
|
||||
handler_called = False
|
||||
|
||||
async def final_handler(ctx: FunctionInvocationContext) -> str:
|
||||
nonlocal handler_called
|
||||
handler_called = True
|
||||
# This should not be called when result is pre-set
|
||||
return "original result"
|
||||
|
||||
result = await pipeline.execute(mock_function, arguments, context, final_handler)
|
||||
|
||||
# Verify pre-override worked and handler was NOT called (because result was already set)
|
||||
assert result == "pre-override result"
|
||||
assert not handler_called
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_agent() -> AgentProtocol:
|
||||
"""Mock agent for testing."""
|
||||
agent = MagicMock(spec=AgentProtocol)
|
||||
agent.name = "test_agent"
|
||||
return agent
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_function() -> AIFunction[Any, Any]:
|
||||
"""Mock function for testing."""
|
||||
function = MagicMock(spec=AIFunction[Any, Any])
|
||||
function.name = "test_function"
|
||||
return function
|
||||
@@ -0,0 +1,463 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from collections.abc import AsyncIterable, Awaitable, Callable
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from agent_framework import (
|
||||
AgentProtocol,
|
||||
AgentRunResponse,
|
||||
AgentRunResponseUpdate,
|
||||
ChatAgent,
|
||||
ChatMessage,
|
||||
Role,
|
||||
TextContent,
|
||||
)
|
||||
from agent_framework._middleware import (
|
||||
AgentMiddleware,
|
||||
AgentMiddlewarePipeline,
|
||||
AgentRunContext,
|
||||
FunctionInvocationContext,
|
||||
FunctionMiddleware,
|
||||
FunctionMiddlewarePipeline,
|
||||
)
|
||||
from agent_framework._tools import AIFunction
|
||||
|
||||
from .conftest import MockChatClient
|
||||
|
||||
|
||||
class FunctionTestArgs(BaseModel):
|
||||
"""Test arguments for function middleware tests."""
|
||||
|
||||
name: str = Field(description="Test name parameter")
|
||||
|
||||
|
||||
class TestResultOverrideMiddleware:
|
||||
"""Test cases for middleware result override functionality."""
|
||||
|
||||
async def test_agent_middleware_response_override_non_streaming(self, mock_agent: AgentProtocol) -> None:
|
||||
"""Test that agent middleware can override response for non-streaming execution."""
|
||||
override_response = AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="overridden response")])
|
||||
|
||||
class ResponseOverrideMiddleware(AgentMiddleware):
|
||||
async def process(
|
||||
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
# Execute the pipeline first, then override the response
|
||||
await next(context)
|
||||
context.result = override_response
|
||||
|
||||
middleware = ResponseOverrideMiddleware()
|
||||
pipeline = AgentMiddlewarePipeline([middleware])
|
||||
messages = [ChatMessage(role=Role.USER, text="test")]
|
||||
context = AgentRunContext(agent=mock_agent, messages=messages)
|
||||
|
||||
handler_called = False
|
||||
|
||||
async def final_handler(ctx: AgentRunContext) -> AgentRunResponse:
|
||||
nonlocal handler_called
|
||||
handler_called = True
|
||||
return AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="original response")])
|
||||
|
||||
result = await pipeline.execute(mock_agent, messages, context, final_handler)
|
||||
|
||||
# Verify the overridden response is returned
|
||||
assert result is not None
|
||||
assert result == override_response
|
||||
assert result.messages[0].text == "overridden response"
|
||||
# Verify original handler was called since middleware called next()
|
||||
assert handler_called
|
||||
|
||||
async def test_agent_middleware_response_override_streaming(self, mock_agent: AgentProtocol) -> None:
|
||||
"""Test that agent middleware can override response for streaming execution."""
|
||||
|
||||
async def override_stream() -> AsyncIterable[AgentRunResponseUpdate]:
|
||||
yield AgentRunResponseUpdate(contents=[TextContent(text="overridden")])
|
||||
yield AgentRunResponseUpdate(contents=[TextContent(text=" stream")])
|
||||
|
||||
class StreamResponseOverrideMiddleware(AgentMiddleware):
|
||||
async def process(
|
||||
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
# Execute the pipeline first, then override the response stream
|
||||
await next(context)
|
||||
context.result = override_stream()
|
||||
|
||||
middleware = StreamResponseOverrideMiddleware()
|
||||
pipeline = AgentMiddlewarePipeline([middleware])
|
||||
messages = [ChatMessage(role=Role.USER, text="test")]
|
||||
context = AgentRunContext(agent=mock_agent, messages=messages)
|
||||
|
||||
async def final_handler(ctx: AgentRunContext) -> AsyncIterable[AgentRunResponseUpdate]:
|
||||
yield AgentRunResponseUpdate(contents=[TextContent(text="original")])
|
||||
|
||||
updates: list[AgentRunResponseUpdate] = []
|
||||
async for update in pipeline.execute_stream(mock_agent, messages, context, final_handler):
|
||||
updates.append(update)
|
||||
|
||||
# Verify the overridden response stream is returned
|
||||
assert len(updates) == 2
|
||||
assert updates[0].text == "overridden"
|
||||
assert updates[1].text == " stream"
|
||||
|
||||
async def test_function_middleware_result_override(self, mock_function: AIFunction[Any, Any]) -> None:
|
||||
"""Test that function middleware can override result."""
|
||||
override_result = "overridden function result"
|
||||
|
||||
class ResultOverrideMiddleware(FunctionMiddleware):
|
||||
async def process(
|
||||
self,
|
||||
context: FunctionInvocationContext,
|
||||
next: Callable[[FunctionInvocationContext], Awaitable[None]],
|
||||
) -> None:
|
||||
# Execute the pipeline first, then override the result
|
||||
await next(context)
|
||||
context.result = override_result
|
||||
|
||||
middleware = ResultOverrideMiddleware()
|
||||
pipeline = FunctionMiddlewarePipeline([middleware])
|
||||
arguments = FunctionTestArgs(name="test")
|
||||
context = FunctionInvocationContext(function=mock_function, arguments=arguments)
|
||||
|
||||
handler_called = False
|
||||
|
||||
async def final_handler(ctx: FunctionInvocationContext) -> str:
|
||||
nonlocal handler_called
|
||||
handler_called = True
|
||||
return "original function result"
|
||||
|
||||
result = await pipeline.execute(mock_function, arguments, context, final_handler)
|
||||
|
||||
# Verify the overridden result is returned
|
||||
assert result == override_result
|
||||
# Verify original handler was called since middleware called next()
|
||||
assert handler_called
|
||||
|
||||
async def test_chat_agent_middleware_response_override(self) -> None:
|
||||
"""Test result override functionality with ChatAgent integration."""
|
||||
mock_chat_client = MockChatClient()
|
||||
|
||||
class ChatAgentResponseOverrideMiddleware(AgentMiddleware):
|
||||
async def process(
|
||||
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
# Always call next() first to allow execution
|
||||
await next(context)
|
||||
# Then conditionally override based on content
|
||||
if any("special" in msg.text for msg in context.messages if msg.text):
|
||||
context.result = AgentRunResponse(
|
||||
messages=[ChatMessage(role=Role.ASSISTANT, text="Special response from middleware!")]
|
||||
)
|
||||
|
||||
# Create ChatAgent with override middleware
|
||||
middleware = ChatAgentResponseOverrideMiddleware()
|
||||
agent = ChatAgent(chat_client=mock_chat_client, middleware=[middleware])
|
||||
|
||||
# Test override case
|
||||
override_messages = [ChatMessage(role=Role.USER, text="Give me a special response")]
|
||||
override_response = await agent.run(override_messages)
|
||||
assert override_response.messages[0].text == "Special response from middleware!"
|
||||
# Verify chat client was called since middleware called next()
|
||||
assert mock_chat_client.call_count == 1
|
||||
|
||||
# Test normal case
|
||||
normal_messages = [ChatMessage(role=Role.USER, text="Normal request")]
|
||||
normal_response = await agent.run(normal_messages)
|
||||
assert normal_response.messages[0].text == "test response"
|
||||
# Verify chat client was called for normal case
|
||||
assert mock_chat_client.call_count == 2
|
||||
|
||||
async def test_chat_agent_middleware_streaming_override(self) -> None:
|
||||
"""Test streaming result override functionality with ChatAgent integration."""
|
||||
mock_chat_client = MockChatClient()
|
||||
|
||||
async def custom_stream() -> AsyncIterable[AgentRunResponseUpdate]:
|
||||
yield AgentRunResponseUpdate(contents=[TextContent(text="Custom")])
|
||||
yield AgentRunResponseUpdate(contents=[TextContent(text=" streaming")])
|
||||
yield AgentRunResponseUpdate(contents=[TextContent(text=" response!")])
|
||||
|
||||
class ChatAgentStreamOverrideMiddleware(AgentMiddleware):
|
||||
async def process(
|
||||
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
# Always call next() first to allow execution
|
||||
await next(context)
|
||||
# Then conditionally override based on content
|
||||
if any("custom stream" in msg.text for msg in context.messages if msg.text):
|
||||
context.result = custom_stream()
|
||||
|
||||
# Create ChatAgent with override middleware
|
||||
middleware = ChatAgentStreamOverrideMiddleware()
|
||||
agent = ChatAgent(chat_client=mock_chat_client, middleware=[middleware])
|
||||
|
||||
# Test streaming override case
|
||||
override_messages = [ChatMessage(role=Role.USER, text="Give me a custom stream")]
|
||||
override_updates: list[AgentRunResponseUpdate] = []
|
||||
async for update in agent.run_stream(override_messages):
|
||||
override_updates.append(update)
|
||||
|
||||
assert len(override_updates) == 3
|
||||
assert override_updates[0].text == "Custom"
|
||||
assert override_updates[1].text == " streaming"
|
||||
assert override_updates[2].text == " response!"
|
||||
|
||||
# Test normal streaming case
|
||||
normal_messages = [ChatMessage(role=Role.USER, text="Normal streaming request")]
|
||||
normal_updates: list[AgentRunResponseUpdate] = []
|
||||
async for update in agent.run_stream(normal_messages):
|
||||
normal_updates.append(update)
|
||||
|
||||
assert len(normal_updates) == 2
|
||||
assert normal_updates[0].text == "test streaming response "
|
||||
assert normal_updates[1].text == "another update"
|
||||
|
||||
async def test_agent_middleware_conditional_no_next(self, mock_agent: AgentProtocol) -> None:
|
||||
"""Test that when agent middleware conditionally doesn't call next(), no execution happens."""
|
||||
|
||||
class ConditionalNoNextMiddleware(AgentMiddleware):
|
||||
async def process(
|
||||
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
# Only call next() if message contains "execute"
|
||||
if any("execute" in msg.text for msg in context.messages if msg.text):
|
||||
await next(context)
|
||||
# Otherwise, don't call next() - no execution should happen
|
||||
|
||||
middleware = ConditionalNoNextMiddleware()
|
||||
pipeline = AgentMiddlewarePipeline([middleware])
|
||||
|
||||
handler_called = False
|
||||
|
||||
async def final_handler(ctx: AgentRunContext) -> AgentRunResponse:
|
||||
nonlocal handler_called
|
||||
handler_called = True
|
||||
return AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="executed response")])
|
||||
|
||||
# Test case where next() is NOT called
|
||||
no_execute_messages = [ChatMessage(role=Role.USER, text="Don't run this")]
|
||||
no_execute_context = AgentRunContext(agent=mock_agent, messages=no_execute_messages)
|
||||
no_execute_result = await pipeline.execute(mock_agent, no_execute_messages, no_execute_context, final_handler)
|
||||
|
||||
# When middleware doesn't call next(), result should be empty AgentRunResponse
|
||||
assert no_execute_result is not None
|
||||
assert isinstance(no_execute_result, AgentRunResponse)
|
||||
assert no_execute_result.messages == [] # Empty response
|
||||
assert not handler_called
|
||||
assert no_execute_context.result is None
|
||||
|
||||
# Reset for next test
|
||||
handler_called = False
|
||||
|
||||
# Test case where next() IS called
|
||||
execute_messages = [ChatMessage(role=Role.USER, text="Please execute this")]
|
||||
execute_context = AgentRunContext(agent=mock_agent, messages=execute_messages)
|
||||
execute_result = await pipeline.execute(mock_agent, execute_messages, execute_context, final_handler)
|
||||
|
||||
assert execute_result is not None
|
||||
assert execute_result.messages[0].text == "executed response"
|
||||
assert handler_called
|
||||
|
||||
async def test_function_middleware_conditional_no_next(self, mock_function: AIFunction[Any, Any]) -> None:
|
||||
"""Test that when function middleware conditionally doesn't call next(), no execution happens."""
|
||||
|
||||
class ConditionalNoNextFunctionMiddleware(FunctionMiddleware):
|
||||
async def process(
|
||||
self,
|
||||
context: FunctionInvocationContext,
|
||||
next: Callable[[FunctionInvocationContext], Awaitable[None]],
|
||||
) -> None:
|
||||
# Only call next() if argument name contains "execute"
|
||||
args = context.arguments
|
||||
assert isinstance(args, FunctionTestArgs)
|
||||
if "execute" in args.name:
|
||||
await next(context)
|
||||
# Otherwise, don't call next() - no execution should happen
|
||||
|
||||
middleware = ConditionalNoNextFunctionMiddleware()
|
||||
pipeline = FunctionMiddlewarePipeline([middleware])
|
||||
|
||||
handler_called = False
|
||||
|
||||
async def final_handler(ctx: FunctionInvocationContext) -> str:
|
||||
nonlocal handler_called
|
||||
handler_called = True
|
||||
return "executed function result"
|
||||
|
||||
# Test case where next() is NOT called
|
||||
no_execute_args = FunctionTestArgs(name="test_no_action")
|
||||
no_execute_context = FunctionInvocationContext(function=mock_function, arguments=no_execute_args)
|
||||
no_execute_result = await pipeline.execute(mock_function, no_execute_args, no_execute_context, final_handler)
|
||||
|
||||
# When middleware doesn't call next(), function result should be None (functions can return None)
|
||||
assert no_execute_result is None
|
||||
assert not handler_called
|
||||
assert no_execute_context.result is None
|
||||
|
||||
# Reset for next test
|
||||
handler_called = False
|
||||
|
||||
# Test case where next() IS called
|
||||
execute_args = FunctionTestArgs(name="test_execute")
|
||||
execute_context = FunctionInvocationContext(function=mock_function, arguments=execute_args)
|
||||
execute_result = await pipeline.execute(mock_function, execute_args, execute_context, final_handler)
|
||||
|
||||
assert execute_result == "executed function result"
|
||||
assert handler_called
|
||||
|
||||
|
||||
class TestResultObservability:
|
||||
"""Test cases for middleware result observability functionality."""
|
||||
|
||||
async def test_agent_middleware_response_observability(self, mock_agent: AgentProtocol) -> None:
|
||||
"""Test that middleware can observe response after execution."""
|
||||
observed_responses: list[AgentRunResponse] = []
|
||||
|
||||
class ObservabilityMiddleware(AgentMiddleware):
|
||||
async def process(
|
||||
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
# Context should be empty before next()
|
||||
assert context.result is None
|
||||
|
||||
# Call next to execute
|
||||
await next(context)
|
||||
|
||||
# Context should now contain the response for observability
|
||||
assert context.result is not None
|
||||
assert isinstance(context.result, AgentRunResponse)
|
||||
observed_responses.append(context.result)
|
||||
|
||||
middleware = ObservabilityMiddleware()
|
||||
pipeline = AgentMiddlewarePipeline([middleware])
|
||||
messages = [ChatMessage(role=Role.USER, text="test")]
|
||||
context = AgentRunContext(agent=mock_agent, messages=messages)
|
||||
|
||||
async def final_handler(ctx: AgentRunContext) -> AgentRunResponse:
|
||||
return AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="executed response")])
|
||||
|
||||
result = await pipeline.execute(mock_agent, messages, context, final_handler)
|
||||
|
||||
# Verify response was observed
|
||||
assert len(observed_responses) == 1
|
||||
assert observed_responses[0].messages[0].text == "executed response"
|
||||
assert result == observed_responses[0]
|
||||
|
||||
async def test_function_middleware_result_observability(self, mock_function: AIFunction[Any, Any]) -> None:
|
||||
"""Test that middleware can observe function result after execution."""
|
||||
observed_results: list[str] = []
|
||||
|
||||
class ObservabilityMiddleware(FunctionMiddleware):
|
||||
async def process(
|
||||
self,
|
||||
context: FunctionInvocationContext,
|
||||
next: Callable[[FunctionInvocationContext], Awaitable[None]],
|
||||
) -> None:
|
||||
# Context should be empty before next()
|
||||
assert context.result is None
|
||||
|
||||
# Call next to execute
|
||||
await next(context)
|
||||
|
||||
# Context should now contain the result for observability
|
||||
assert context.result is not None
|
||||
observed_results.append(context.result)
|
||||
|
||||
middleware = ObservabilityMiddleware()
|
||||
pipeline = FunctionMiddlewarePipeline([middleware])
|
||||
arguments = FunctionTestArgs(name="test")
|
||||
context = FunctionInvocationContext(function=mock_function, arguments=arguments)
|
||||
|
||||
async def final_handler(ctx: FunctionInvocationContext) -> str:
|
||||
return "executed function result"
|
||||
|
||||
result = await pipeline.execute(mock_function, arguments, context, final_handler)
|
||||
|
||||
# Verify result was observed
|
||||
assert len(observed_results) == 1
|
||||
assert observed_results[0] == "executed function result"
|
||||
assert result == observed_results[0]
|
||||
|
||||
async def test_agent_middleware_post_execution_override(self, mock_agent: AgentProtocol) -> None:
|
||||
"""Test that middleware can override response after observing execution."""
|
||||
|
||||
class PostExecutionOverrideMiddleware(AgentMiddleware):
|
||||
async def process(
|
||||
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
# Call next to execute first
|
||||
await next(context)
|
||||
|
||||
# Now observe and conditionally override
|
||||
assert context.result is not None
|
||||
assert isinstance(context.result, AgentRunResponse)
|
||||
|
||||
if "modify" in context.result.messages[0].text:
|
||||
# Override after observing
|
||||
context.result = AgentRunResponse(
|
||||
messages=[ChatMessage(role=Role.ASSISTANT, text="modified after execution")]
|
||||
)
|
||||
|
||||
middleware = PostExecutionOverrideMiddleware()
|
||||
pipeline = AgentMiddlewarePipeline([middleware])
|
||||
messages = [ChatMessage(role=Role.USER, text="test")]
|
||||
context = AgentRunContext(agent=mock_agent, messages=messages)
|
||||
|
||||
async def final_handler(ctx: AgentRunContext) -> AgentRunResponse:
|
||||
return AgentRunResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="response to modify")])
|
||||
|
||||
result = await pipeline.execute(mock_agent, messages, context, final_handler)
|
||||
|
||||
# Verify response was modified after execution
|
||||
assert result is not None
|
||||
assert result.messages[0].text == "modified after execution"
|
||||
|
||||
async def test_function_middleware_post_execution_override(self, mock_function: AIFunction[Any, Any]) -> None:
|
||||
"""Test that middleware can override function result after observing execution."""
|
||||
|
||||
class PostExecutionOverrideMiddleware(FunctionMiddleware):
|
||||
async def process(
|
||||
self,
|
||||
context: FunctionInvocationContext,
|
||||
next: Callable[[FunctionInvocationContext], Awaitable[None]],
|
||||
) -> None:
|
||||
# Call next to execute first
|
||||
await next(context)
|
||||
|
||||
# Now observe and conditionally override
|
||||
assert context.result is not None
|
||||
|
||||
if "modify" in context.result:
|
||||
# Override after observing
|
||||
context.result = "modified after execution"
|
||||
|
||||
middleware = PostExecutionOverrideMiddleware()
|
||||
pipeline = FunctionMiddlewarePipeline([middleware])
|
||||
arguments = FunctionTestArgs(name="test")
|
||||
context = FunctionInvocationContext(function=mock_function, arguments=arguments)
|
||||
|
||||
async def final_handler(ctx: FunctionInvocationContext) -> str:
|
||||
return "result to modify"
|
||||
|
||||
result = await pipeline.execute(mock_function, arguments, context, final_handler)
|
||||
|
||||
# Verify result was modified after execution
|
||||
assert result == "modified after execution"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_agent() -> AgentProtocol:
|
||||
"""Mock agent for testing."""
|
||||
agent = MagicMock(spec=AgentProtocol)
|
||||
agent.name = "test_agent"
|
||||
return agent
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_function() -> AIFunction[Any, Any]:
|
||||
"""Mock function for testing."""
|
||||
function = MagicMock(spec=AIFunction[Any, Any])
|
||||
function.name = "test_function"
|
||||
return function
|
||||
@@ -0,0 +1,544 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from collections.abc import Awaitable, Callable
|
||||
|
||||
from agent_framework import (
|
||||
AgentRunResponseUpdate,
|
||||
ChatAgent,
|
||||
ChatMessage,
|
||||
ChatResponse,
|
||||
ChatResponseUpdate,
|
||||
FunctionCallContent,
|
||||
FunctionResultContent,
|
||||
Role,
|
||||
TextContent,
|
||||
)
|
||||
from agent_framework._middleware import (
|
||||
AgentMiddleware,
|
||||
AgentRunContext,
|
||||
FunctionInvocationContext,
|
||||
FunctionMiddleware,
|
||||
)
|
||||
|
||||
from .conftest import MockChatClient
|
||||
|
||||
# region ChatAgent Tests
|
||||
|
||||
|
||||
class TestChatAgentClassBasedMiddleware:
|
||||
"""Test cases for class-based middleware integration with ChatAgent."""
|
||||
|
||||
async def test_class_based_agent_middleware_with_chat_agent(self, chat_client: "MockChatClient") -> None:
|
||||
"""Test class-based agent middleware with ChatAgent."""
|
||||
execution_order: list[str] = []
|
||||
|
||||
class TrackingAgentMiddleware(AgentMiddleware):
|
||||
def __init__(self, name: str):
|
||||
self.name = name
|
||||
|
||||
async def process(
|
||||
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
execution_order.append(f"{self.name}_before")
|
||||
await next(context)
|
||||
execution_order.append(f"{self.name}_after")
|
||||
|
||||
# Create ChatAgent with middleware
|
||||
middleware = TrackingAgentMiddleware("agent_middleware")
|
||||
agent = ChatAgent(chat_client=chat_client, middleware=[middleware])
|
||||
|
||||
# Execute the agent
|
||||
messages = [ChatMessage(role=Role.USER, text="test message")]
|
||||
response = await agent.run(messages)
|
||||
|
||||
# Verify response
|
||||
assert response is not None
|
||||
assert len(response.messages) > 0
|
||||
assert response.messages[0].role == Role.ASSISTANT
|
||||
# Note: conftest "MockChatClient" returns different text format
|
||||
assert "test response" in response.messages[0].text
|
||||
|
||||
# Verify middleware execution order
|
||||
assert execution_order == ["agent_middleware_before", "agent_middleware_after"]
|
||||
|
||||
async def test_class_based_function_middleware_with_chat_agent(self, chat_client: "MockChatClient") -> None:
|
||||
"""Test class-based function middleware with ChatAgent."""
|
||||
execution_order: list[str] = []
|
||||
|
||||
class TrackingFunctionMiddleware(FunctionMiddleware):
|
||||
def __init__(self, name: str):
|
||||
self.name = name
|
||||
|
||||
async def process(
|
||||
self,
|
||||
context: FunctionInvocationContext,
|
||||
next: Callable[[FunctionInvocationContext], Awaitable[None]],
|
||||
) -> None:
|
||||
execution_order.append(f"{self.name}_before")
|
||||
await next(context)
|
||||
execution_order.append(f"{self.name}_after")
|
||||
|
||||
# Create ChatAgent with function middleware (no tools, so function middleware won't be triggered)
|
||||
middleware = TrackingFunctionMiddleware("function_middleware")
|
||||
agent = ChatAgent(chat_client=chat_client, middleware=[middleware])
|
||||
|
||||
# Execute the agent
|
||||
messages = [ChatMessage(role=Role.USER, text="test message")]
|
||||
response = await agent.run(messages)
|
||||
|
||||
# Verify response
|
||||
assert response is not None
|
||||
assert len(response.messages) > 0
|
||||
assert chat_client.call_count == 1
|
||||
|
||||
# Note: Function middleware won't execute since no function calls are made
|
||||
assert execution_order == []
|
||||
|
||||
|
||||
class TestChatAgentFunctionBasedMiddleware:
|
||||
"""Test cases for function-based middleware integration with ChatAgent."""
|
||||
|
||||
async def test_function_based_agent_middleware_with_chat_agent(self, chat_client: "MockChatClient") -> None:
|
||||
"""Test function-based agent middleware with ChatAgent."""
|
||||
execution_order: list[str] = []
|
||||
|
||||
async def tracking_agent_middleware(
|
||||
context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
execution_order.append("agent_function_before")
|
||||
await next(context)
|
||||
execution_order.append("agent_function_after")
|
||||
|
||||
# Create ChatAgent with function middleware
|
||||
agent = ChatAgent(chat_client=chat_client, middleware=[tracking_agent_middleware])
|
||||
|
||||
# Execute the agent
|
||||
messages = [ChatMessage(role=Role.USER, text="test message")]
|
||||
response = await agent.run(messages)
|
||||
|
||||
# Verify response
|
||||
assert response is not None
|
||||
assert len(response.messages) > 0
|
||||
assert response.messages[0].role == Role.ASSISTANT
|
||||
assert response.messages[0].text == "test response"
|
||||
assert chat_client.call_count == 1
|
||||
|
||||
# Verify middleware execution order
|
||||
assert execution_order == ["agent_function_before", "agent_function_after"]
|
||||
|
||||
async def test_function_based_function_middleware_with_chat_agent(self, chat_client: "MockChatClient") -> None:
|
||||
"""Test function-based function middleware with ChatAgent."""
|
||||
execution_order: list[str] = []
|
||||
|
||||
async def tracking_function_middleware(
|
||||
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
|
||||
) -> None:
|
||||
execution_order.append("function_function_before")
|
||||
await next(context)
|
||||
execution_order.append("function_function_after")
|
||||
|
||||
# Create ChatAgent with function middleware (no tools, so function middleware won't be triggered)
|
||||
agent = ChatAgent(chat_client=chat_client, middleware=[tracking_function_middleware])
|
||||
|
||||
# Execute the agent
|
||||
messages = [ChatMessage(role=Role.USER, text="test message")]
|
||||
response = await agent.run(messages)
|
||||
|
||||
# Verify response
|
||||
assert response is not None
|
||||
assert len(response.messages) > 0
|
||||
assert chat_client.call_count == 1
|
||||
|
||||
# Note: Function middleware won't execute since no function calls are made
|
||||
assert execution_order == []
|
||||
|
||||
|
||||
class TestChatAgentStreamingMiddleware:
|
||||
"""Test cases for streaming middleware integration with ChatAgent."""
|
||||
|
||||
async def test_agent_middleware_with_streaming(self, chat_client: "MockChatClient") -> None:
|
||||
"""Test agent middleware with streaming ChatAgent responses."""
|
||||
execution_order: list[str] = []
|
||||
streaming_flags: list[bool] = []
|
||||
|
||||
class StreamingTrackingMiddleware(AgentMiddleware):
|
||||
async def process(
|
||||
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
execution_order.append("middleware_before")
|
||||
streaming_flags.append(context.is_streaming)
|
||||
await next(context)
|
||||
execution_order.append("middleware_after")
|
||||
|
||||
# Create ChatAgent with middleware
|
||||
middleware = StreamingTrackingMiddleware()
|
||||
agent = ChatAgent(chat_client=chat_client, middleware=[middleware])
|
||||
|
||||
# Set up mock streaming responses
|
||||
chat_client.streaming_responses = [
|
||||
[
|
||||
ChatResponseUpdate(contents=[TextContent(text="Streaming")], role=Role.ASSISTANT),
|
||||
ChatResponseUpdate(contents=[TextContent(text=" response")], role=Role.ASSISTANT),
|
||||
]
|
||||
]
|
||||
|
||||
# Execute streaming
|
||||
messages = [ChatMessage(role=Role.USER, text="test message")]
|
||||
updates: list[AgentRunResponseUpdate] = []
|
||||
async for update in agent.run_stream(messages):
|
||||
updates.append(update)
|
||||
|
||||
# Verify streaming response
|
||||
assert len(updates) == 2
|
||||
assert updates[0].text == "Streaming"
|
||||
assert updates[1].text == " response"
|
||||
assert chat_client.call_count == 1
|
||||
|
||||
# Verify middleware was called and streaming flag was set correctly
|
||||
assert execution_order == ["middleware_before", "middleware_after"]
|
||||
assert streaming_flags == [True] # Context should indicate streaming
|
||||
|
||||
async def test_non_streaming_vs_streaming_flag_validation(self, chat_client: "MockChatClient") -> None:
|
||||
"""Test that is_streaming flag is correctly set for different execution modes."""
|
||||
streaming_flags: list[bool] = []
|
||||
|
||||
class FlagTrackingMiddleware(AgentMiddleware):
|
||||
async def process(
|
||||
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
streaming_flags.append(context.is_streaming)
|
||||
await next(context)
|
||||
|
||||
# Create ChatAgent with middleware
|
||||
middleware = FlagTrackingMiddleware()
|
||||
agent = ChatAgent(chat_client=chat_client, middleware=[middleware])
|
||||
messages = [ChatMessage(role=Role.USER, text="test message")]
|
||||
|
||||
# Test non-streaming execution
|
||||
response = await agent.run(messages)
|
||||
assert response is not None
|
||||
|
||||
# Test streaming execution
|
||||
async for _ in agent.run_stream(messages):
|
||||
pass
|
||||
|
||||
# Verify flags: [non-streaming, streaming]
|
||||
assert streaming_flags == [False, True]
|
||||
|
||||
|
||||
class TestChatAgentMultipleMiddlewareOrdering:
|
||||
"""Test cases for multiple middleware execution order with ChatAgent."""
|
||||
|
||||
async def test_multiple_agent_middleware_execution_order(self, chat_client: "MockChatClient") -> None:
|
||||
"""Test that multiple agent middlewares execute in correct order with ChatAgent."""
|
||||
execution_order: list[str] = []
|
||||
|
||||
class OrderedMiddleware(AgentMiddleware):
|
||||
def __init__(self, name: str):
|
||||
self.name = name
|
||||
|
||||
async def process(
|
||||
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
execution_order.append(f"{self.name}_before")
|
||||
await next(context)
|
||||
execution_order.append(f"{self.name}_after")
|
||||
|
||||
# Create multiple middlewares
|
||||
middleware1 = OrderedMiddleware("first")
|
||||
middleware2 = OrderedMiddleware("second")
|
||||
middleware3 = OrderedMiddleware("third")
|
||||
|
||||
# Create ChatAgent with multiple middlewares
|
||||
agent = ChatAgent(chat_client=chat_client, middleware=[middleware1, middleware2, middleware3])
|
||||
|
||||
# Execute the agent
|
||||
messages = [ChatMessage(role=Role.USER, text="test message")]
|
||||
response = await agent.run(messages)
|
||||
|
||||
# Verify response
|
||||
assert response is not None
|
||||
assert chat_client.call_count == 1
|
||||
|
||||
# Verify execution order (should be nested: first wraps second wraps third)
|
||||
expected_order = ["first_before", "second_before", "third_before", "third_after", "second_after", "first_after"]
|
||||
assert execution_order == expected_order
|
||||
|
||||
async def test_mixed_middleware_types_with_chat_agent(self, chat_client: "MockChatClient") -> None:
|
||||
"""Test mixed class and function-based middlewares with ChatAgent."""
|
||||
execution_order: list[str] = []
|
||||
|
||||
class ClassAgentMiddleware(AgentMiddleware):
|
||||
async def process(
|
||||
self, context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
execution_order.append("class_agent_before")
|
||||
await next(context)
|
||||
execution_order.append("class_agent_after")
|
||||
|
||||
async def function_agent_middleware(
|
||||
context: AgentRunContext, next: Callable[[AgentRunContext], Awaitable[None]]
|
||||
) -> None:
|
||||
execution_order.append("function_agent_before")
|
||||
await next(context)
|
||||
execution_order.append("function_agent_after")
|
||||
|
||||
class ClassFunctionMiddleware(FunctionMiddleware):
|
||||
async def process(
|
||||
self,
|
||||
context: FunctionInvocationContext,
|
||||
next: Callable[[FunctionInvocationContext], Awaitable[None]],
|
||||
) -> None:
|
||||
execution_order.append("class_function_before")
|
||||
await next(context)
|
||||
execution_order.append("class_function_after")
|
||||
|
||||
async def function_function_middleware(
|
||||
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
|
||||
) -> None:
|
||||
execution_order.append("function_function_before")
|
||||
await next(context)
|
||||
execution_order.append("function_function_after")
|
||||
|
||||
# Create ChatAgent with mixed middleware types (no tools, focusing on agent middleware)
|
||||
agent = ChatAgent(
|
||||
chat_client=chat_client,
|
||||
middleware=[
|
||||
ClassAgentMiddleware(),
|
||||
function_agent_middleware,
|
||||
ClassFunctionMiddleware(), # Won't execute without function calls
|
||||
function_function_middleware, # Won't execute without function calls
|
||||
],
|
||||
)
|
||||
|
||||
# Execute the agent
|
||||
messages = [ChatMessage(role=Role.USER, text="test message")]
|
||||
response = await agent.run(messages)
|
||||
|
||||
# Verify response
|
||||
assert response is not None
|
||||
assert chat_client.call_count == 1
|
||||
|
||||
# Verify that agent middlewares were executed in correct order
|
||||
# (Function middlewares won't execute since no functions are called)
|
||||
expected_order = ["class_agent_before", "function_agent_before", "function_agent_after", "class_agent_after"]
|
||||
assert execution_order == expected_order
|
||||
|
||||
|
||||
# region Tool Functions for Testing
|
||||
|
||||
|
||||
def sample_tool_function(location: str) -> str:
|
||||
"""A simple tool function for middleware testing."""
|
||||
return f"Weather in {location}: sunny"
|
||||
|
||||
|
||||
# region ChatAgent Function Middleware Tests with Tools
|
||||
|
||||
|
||||
class TestChatAgentFunctionMiddlewareWithTools:
|
||||
"""Test cases for function middleware integration with ChatAgent when tools are used."""
|
||||
|
||||
async def test_class_based_function_middleware_with_tool_calls(self, chat_client: "MockChatClient") -> None:
|
||||
"""Test class-based function middleware with ChatAgent when function calls are made."""
|
||||
execution_order: list[str] = []
|
||||
|
||||
class TrackingFunctionMiddleware(FunctionMiddleware):
|
||||
def __init__(self, name: str):
|
||||
self.name = name
|
||||
|
||||
async def process(
|
||||
self,
|
||||
context: FunctionInvocationContext,
|
||||
next: Callable[[FunctionInvocationContext], Awaitable[None]],
|
||||
) -> None:
|
||||
execution_order.append(f"{self.name}_before")
|
||||
await next(context)
|
||||
execution_order.append(f"{self.name}_after")
|
||||
|
||||
# Set up mock to return a function call first, then a regular response
|
||||
function_call_response = ChatResponse(
|
||||
messages=[
|
||||
ChatMessage(
|
||||
role=Role.ASSISTANT,
|
||||
contents=[
|
||||
FunctionCallContent(
|
||||
call_id="call_123",
|
||||
name="sample_tool_function",
|
||||
arguments='{"location": "Seattle"}',
|
||||
)
|
||||
],
|
||||
)
|
||||
]
|
||||
)
|
||||
final_response = ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="Final response")])
|
||||
|
||||
chat_client.responses = [function_call_response, final_response]
|
||||
|
||||
# Create ChatAgent with function middleware and tools
|
||||
middleware = TrackingFunctionMiddleware("function_middleware")
|
||||
agent = ChatAgent(
|
||||
chat_client=chat_client,
|
||||
middleware=[middleware],
|
||||
tools=[sample_tool_function],
|
||||
)
|
||||
|
||||
# Execute the agent
|
||||
messages = [ChatMessage(role=Role.USER, text="Get weather for Seattle")]
|
||||
response = await agent.run(messages)
|
||||
|
||||
# Verify response
|
||||
assert response is not None
|
||||
assert len(response.messages) > 0
|
||||
assert chat_client.call_count == 2 # Two calls: one for function call, one for final response
|
||||
|
||||
# Verify function middleware was executed
|
||||
assert execution_order == ["function_middleware_before", "function_middleware_after"]
|
||||
|
||||
# Verify function call and result are in the response
|
||||
all_contents = [content for message in response.messages for content in message.contents]
|
||||
function_calls = [c for c in all_contents if isinstance(c, FunctionCallContent)]
|
||||
function_results = [c for c in all_contents if isinstance(c, FunctionResultContent)]
|
||||
|
||||
assert len(function_calls) == 1
|
||||
assert len(function_results) == 1
|
||||
assert function_calls[0].name == "sample_tool_function"
|
||||
assert function_results[0].call_id == function_calls[0].call_id
|
||||
|
||||
async def test_function_based_function_middleware_with_tool_calls(self, chat_client: "MockChatClient") -> None:
|
||||
"""Test function-based function middleware with ChatAgent when function calls are made."""
|
||||
execution_order: list[str] = []
|
||||
|
||||
async def tracking_function_middleware(
|
||||
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
|
||||
) -> None:
|
||||
execution_order.append("function_middleware_before")
|
||||
await next(context)
|
||||
execution_order.append("function_middleware_after")
|
||||
|
||||
# Set up mock to return a function call first, then a regular response
|
||||
function_call_response = ChatResponse(
|
||||
messages=[
|
||||
ChatMessage(
|
||||
role=Role.ASSISTANT,
|
||||
contents=[
|
||||
FunctionCallContent(
|
||||
call_id="call_456",
|
||||
name="sample_tool_function",
|
||||
arguments='{"location": "San Francisco"}',
|
||||
)
|
||||
],
|
||||
)
|
||||
]
|
||||
)
|
||||
final_response = ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="Final response")])
|
||||
|
||||
chat_client.responses = [function_call_response, final_response]
|
||||
|
||||
# Create ChatAgent with function middleware and tools
|
||||
agent = ChatAgent(
|
||||
chat_client=chat_client,
|
||||
middleware=[tracking_function_middleware],
|
||||
tools=[sample_tool_function],
|
||||
)
|
||||
|
||||
# Execute the agent
|
||||
messages = [ChatMessage(role=Role.USER, text="Get weather for San Francisco")]
|
||||
response = await agent.run(messages)
|
||||
|
||||
# Verify response
|
||||
assert response is not None
|
||||
assert len(response.messages) > 0
|
||||
assert chat_client.call_count == 2 # Two calls: one for function call, one for final response
|
||||
|
||||
# Verify function middleware was executed
|
||||
assert execution_order == ["function_middleware_before", "function_middleware_after"]
|
||||
|
||||
# Verify function call and result are in the response
|
||||
all_contents = [content for message in response.messages for content in message.contents]
|
||||
function_calls = [c for c in all_contents if isinstance(c, FunctionCallContent)]
|
||||
function_results = [c for c in all_contents if isinstance(c, FunctionResultContent)]
|
||||
|
||||
assert len(function_calls) == 1
|
||||
assert len(function_results) == 1
|
||||
assert function_calls[0].name == "sample_tool_function"
|
||||
assert function_results[0].call_id == function_calls[0].call_id
|
||||
|
||||
async def test_mixed_agent_and_function_middleware_with_tool_calls(self, chat_client: "MockChatClient") -> None:
|
||||
"""Test both agent and function middleware with ChatAgent when function calls are made."""
|
||||
execution_order: list[str] = []
|
||||
|
||||
class TrackingAgentMiddleware(AgentMiddleware):
|
||||
async def process(
|
||||
self,
|
||||
context: AgentRunContext,
|
||||
next: Callable[[AgentRunContext], Awaitable[None]],
|
||||
) -> None:
|
||||
execution_order.append("agent_middleware_before")
|
||||
await next(context)
|
||||
execution_order.append("agent_middleware_after")
|
||||
|
||||
class TrackingFunctionMiddleware(FunctionMiddleware):
|
||||
async def process(
|
||||
self,
|
||||
context: FunctionInvocationContext,
|
||||
next: Callable[[FunctionInvocationContext], Awaitable[None]],
|
||||
) -> None:
|
||||
execution_order.append("function_middleware_before")
|
||||
await next(context)
|
||||
execution_order.append("function_middleware_after")
|
||||
|
||||
# Set up mock to return a function call first, then a regular response
|
||||
function_call_response = ChatResponse(
|
||||
messages=[
|
||||
ChatMessage(
|
||||
role=Role.ASSISTANT,
|
||||
contents=[
|
||||
FunctionCallContent(
|
||||
call_id="call_789",
|
||||
name="sample_tool_function",
|
||||
arguments='{"location": "New York"}',
|
||||
)
|
||||
],
|
||||
)
|
||||
]
|
||||
)
|
||||
final_response = ChatResponse(messages=[ChatMessage(role=Role.ASSISTANT, text="Final response")])
|
||||
|
||||
chat_client.responses = [function_call_response, final_response]
|
||||
|
||||
# Create ChatAgent with both agent and function middleware and tools
|
||||
agent = ChatAgent(
|
||||
chat_client=chat_client,
|
||||
middleware=[TrackingAgentMiddleware(), TrackingFunctionMiddleware()],
|
||||
tools=[sample_tool_function],
|
||||
)
|
||||
|
||||
# Execute the agent
|
||||
messages = [ChatMessage(role=Role.USER, text="Get weather for New York")]
|
||||
response = await agent.run(messages)
|
||||
|
||||
# Verify response
|
||||
assert response is not None
|
||||
assert len(response.messages) > 0
|
||||
assert chat_client.call_count == 2 # Two calls: one for function call, one for final response
|
||||
|
||||
# Verify middleware execution order: agent middleware wraps everything,
|
||||
# function middleware only for function calls
|
||||
expected_order = [
|
||||
"agent_middleware_before",
|
||||
"function_middleware_before",
|
||||
"function_middleware_after",
|
||||
"agent_middleware_after",
|
||||
]
|
||||
assert execution_order == expected_order
|
||||
|
||||
# Verify function call and result are in the response
|
||||
all_contents = [content for message in response.messages for content in message.contents]
|
||||
function_calls = [c for c in all_contents if isinstance(c, FunctionCallContent)]
|
||||
function_results = [c for c in all_contents if isinstance(c, FunctionResultContent)]
|
||||
|
||||
assert len(function_calls) == 1
|
||||
assert len(function_results) == 1
|
||||
assert function_calls[0].name == "sample_tool_function"
|
||||
assert function_results[0].call_id == function_calls[0].call_id
|
||||
Reference in New Issue
Block a user