agent/agent-framework
2
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Python: [BREAKING] Renamed next middleware parameter to call_next (#3735)

Browse Source 
* Renamed next middleware parameter to call_next

* Resolved comments
This commit is contained in:
Dmytro Struk
2026-02-09 13:21:27 -08:00
committed by GitHub
Unverified
parent 977c3adfb2
commit e4ca3e60f8
26 changed files with 529 additions and 430 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/core/tests/core/test_as_tool_kwargs_propagation.py
+28 -14
View File
@@ -19,10 +19,12 @@ class TestAsToolKwargsPropagation:
captured_kwargs: dict[str, Any] = {}
@agent_middleware
async def capture_middleware(context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def capture_middleware(
context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
# Capture kwargs passed to the sub-agent
captured_kwargs.update(context.kwargs)
await next(context)
await call_next(context)
# Setup mock response
chat_client.responses = [
@@ -60,9 +62,11 @@ class TestAsToolKwargsPropagation:
captured_kwargs: dict[str, Any] = {}
@agent_middleware
async def capture_middleware(context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def capture_middleware(
context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
captured_kwargs.update(context.kwargs)
await next(context)
await call_next(context)
# Setup mock response
chat_client.responses = [
@@ -95,10 +99,12 @@ class TestAsToolKwargsPropagation:
captured_kwargs_list: list[dict[str, Any]] = []
@agent_middleware
async def capture_middleware(context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def capture_middleware(
context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
# Capture kwargs at each level
captured_kwargs_list.append(dict(context.kwargs))
await next(context)
await call_next(context)
# Setup mock responses to trigger nested tool invocation: B calls tool C, then completes.
chat_client.responses = [
@@ -156,9 +162,11 @@ class TestAsToolKwargsPropagation:
captured_kwargs: dict[str, Any] = {}
@agent_middleware
async def capture_middleware(context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def capture_middleware(
context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
captured_kwargs.update(context.kwargs)
await next(context)
await call_next(context)
# Setup mock streaming responses
from agent_framework import ChatResponseUpdate
@@ -216,9 +224,11 @@ class TestAsToolKwargsPropagation:
captured_kwargs: dict[str, Any] = {}
@agent_middleware
async def capture_middleware(context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def capture_middleware(
context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
captured_kwargs.update(context.kwargs)
await next(context)
await call_next(context)
# Setup mock response
chat_client.responses = [
@@ -256,14 +266,16 @@ class TestAsToolKwargsPropagation:
call_count = 0
@agent_middleware
async def capture_middleware(context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def capture_middleware(
context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
nonlocal call_count
call_count += 1
if call_count == 1:
first_call_kwargs.update(context.kwargs)
elif call_count == 2:
second_call_kwargs.update(context.kwargs)
await next(context)
await call_next(context)
# Setup mock responses for both calls
chat_client.responses = [
@@ -306,9 +318,11 @@ class TestAsToolKwargsPropagation:
captured_kwargs: dict[str, Any] = {}
@agent_middleware
async def capture_middleware(context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def capture_middleware(
context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
captured_kwargs.update(context.kwargs)
await next(context)
await call_next(context)
# Setup mock response
chat_client.responses = [
python/packages/core/tests/core/test_middleware.py
+132 -96
View File
@@ -135,12 +135,12 @@ class TestAgentMiddlewarePipeline:
"""Test cases for AgentMiddlewarePipeline."""
class PreNextTerminateMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]) -> None:
raise MiddlewareTermination
class PostNextTerminateMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Any) -> None:
await next(context)
async def process(self, context: AgentContext, call_next: Any) -> None:
await call_next(context)
raise MiddlewareTermination
def test_init_empty(self) -> None:
@@ -157,8 +157,8 @@ class TestAgentMiddlewarePipeline:
def test_init_with_function_middleware(self) -> None:
"""Test AgentMiddlewarePipeline initialization with function-based middleware."""
async def test_middleware(context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
await next(context)
async def test_middleware(context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]) -> None:
await call_next(context)
pipeline = AgentMiddlewarePipeline(test_middleware)
assert pipeline.has_middlewares
@@ -185,9 +185,11 @@ class TestAgentMiddlewarePipeline:
def __init__(self, name: str):
self.name = name
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
execution_order.append(f"{self.name}_before")
await next(context)
await call_next(context)
execution_order.append(f"{self.name}_after")
middleware = OrderTrackingMiddleware("test")
@@ -236,9 +238,11 @@ class TestAgentMiddlewarePipeline:
def __init__(self, name: str):
self.name = name
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
execution_order.append(f"{self.name}_before")
await next(context)
await call_next(context)
execution_order.append(f"{self.name}_after")
middleware = StreamOrderTrackingMiddleware("test")
@@ -363,10 +367,12 @@ class TestAgentMiddlewarePipeline:
captured_thread = None
class ThreadCapturingMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
nonlocal captured_thread
captured_thread = context.thread
await next(context)
await call_next(context)
middleware = ThreadCapturingMiddleware()
pipeline = AgentMiddlewarePipeline(middleware)
@@ -388,10 +394,12 @@ class TestAgentMiddlewarePipeline:
captured_thread = "not_none" # Use string to distinguish from None
class ThreadCapturingMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
nonlocal captured_thread
captured_thread = context.thread
await next(context)
await call_next(context)
middleware = ThreadCapturingMiddleware()
pipeline = AgentMiddlewarePipeline(middleware)
@@ -412,12 +420,12 @@ class TestFunctionMiddlewarePipeline:
"""Test cases for FunctionMiddlewarePipeline."""
class PreNextTerminateFunctionMiddleware(FunctionMiddleware):
async def process(self, context: FunctionInvocationContext, next: Any) -> None:
async def process(self, context: FunctionInvocationContext, call_next: Any) -> None:
raise MiddlewareTermination
class PostNextTerminateFunctionMiddleware(FunctionMiddleware):
async def process(self, context: FunctionInvocationContext, next: Any) -> None:
await next(context)
async def process(self, context: FunctionInvocationContext, call_next: Any) -> None:
await call_next(context)
raise MiddlewareTermination
async def test_execute_with_pre_next_termination(self, mock_function: FunctionTool[Any, Any]) -> None:
@@ -475,9 +483,9 @@ class TestFunctionMiddlewarePipeline:
"""Test FunctionMiddlewarePipeline initialization with function-based middleware."""
async def test_middleware(
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
context: FunctionInvocationContext, call_next: Callable[[FunctionInvocationContext], Awaitable[None]]
) -> None:
await next(context)
await call_next(context)
pipeline = FunctionMiddlewarePipeline(test_middleware)
assert pipeline.has_middlewares
@@ -507,10 +515,10 @@ class TestFunctionMiddlewarePipeline:
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
execution_order.append(f"{self.name}_before")
await next(context)
await call_next(context)
execution_order.append(f"{self.name}_after")
middleware = OrderTrackingFunctionMiddleware("test")
@@ -533,12 +541,12 @@ class TestChatMiddlewarePipeline:
"""Test cases for ChatMiddlewarePipeline."""
class PreNextTerminateChatMiddleware(ChatMiddleware):
async def process(self, context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def process(self, context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
raise MiddlewareTermination
class PostNextTerminateChatMiddleware(ChatMiddleware):
async def process(self, context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
await next(context)
async def process(self, context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
await call_next(context)
raise MiddlewareTermination
def test_init_empty(self) -> None:
@@ -555,8 +563,8 @@ class TestChatMiddlewarePipeline:
def test_init_with_function_middleware(self) -> None:
"""Test ChatMiddlewarePipeline initialization with function-based middleware."""
async def test_middleware(context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
await next(context)
async def test_middleware(context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
await call_next(context)
pipeline = ChatMiddlewarePipeline(test_middleware)
assert pipeline.has_middlewares
@@ -584,9 +592,9 @@ class TestChatMiddlewarePipeline:
def __init__(self, name: str):
self.name = name
async def process(self, context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def process(self, context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
execution_order.append(f"{self.name}_before")
await next(context)
await call_next(context)
execution_order.append(f"{self.name}_after")
middleware = OrderTrackingChatMiddleware("test")
@@ -636,9 +644,9 @@ class TestChatMiddlewarePipeline:
def __init__(self, name: str):
self.name = name
async def process(self, context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def process(self, context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
execution_order.append(f"{self.name}_before")
await next(context)
await call_next(context)
execution_order.append(f"{self.name}_after")
middleware = StreamOrderTrackingChatMiddleware("test")
@@ -766,10 +774,12 @@ class TestClassBasedMiddleware:
metadata_updates: list[str] = []
class MetadataAgentMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
context.metadata["before"] = True
metadata_updates.append("before")
await next(context)
await call_next(context)
context.metadata["after"] = True
metadata_updates.append("after")
@@ -797,11 +807,11 @@ class TestClassBasedMiddleware:
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
context.metadata["before"] = True
metadata_updates.append("before")
await next(context)
await call_next(context)
context.metadata["after"] = True
metadata_updates.append("after")
@@ -829,10 +839,12 @@ class TestFunctionBasedMiddleware:
"""Test function-based agent middleware."""
execution_order: list[str] = []
async def test_agent_middleware(context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def test_agent_middleware(
context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
execution_order.append("function_before")
context.metadata["function_middleware"] = True
await next(context)
await call_next(context)
execution_order.append("function_after")
pipeline = AgentMiddlewarePipeline(test_agent_middleware)
@@ -854,11 +866,11 @@ class TestFunctionBasedMiddleware:
execution_order: list[str] = []
async def test_function_middleware(
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
context: FunctionInvocationContext, call_next: Callable[[FunctionInvocationContext], Awaitable[None]]
) -> None:
execution_order.append("function_before")
context.metadata["function_middleware"] = True
await next(context)
await call_next(context)
execution_order.append("function_after")
pipeline = FunctionMiddlewarePipeline(test_function_middleware)
@@ -884,14 +896,18 @@ class TestMixedMiddleware:
execution_order: list[str] = []
class ClassMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
execution_order.append("class_before")
await next(context)
await call_next(context)
execution_order.append("class_after")
async def function_middleware(context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def function_middleware(
context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
execution_order.append("function_before")
await next(context)
await call_next(context)
execution_order.append("function_after")
pipeline = AgentMiddlewarePipeline(ClassMiddleware(), function_middleware)
@@ -915,17 +931,17 @@ class TestMixedMiddleware:
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
execution_order.append("class_before")
await next(context)
await call_next(context)
execution_order.append("class_after")
async def function_middleware(
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
context: FunctionInvocationContext, call_next: Callable[[FunctionInvocationContext], Awaitable[None]]
) -> None:
execution_order.append("function_before")
await next(context)
await call_next(context)
execution_order.append("function_after")
pipeline = FunctionMiddlewarePipeline(ClassMiddleware(), function_middleware)
@@ -946,16 +962,16 @@ class TestMixedMiddleware:
execution_order: list[str] = []
class ClassChatMiddleware(ChatMiddleware):
async def process(self, context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def process(self, context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
execution_order.append("class_before")
await next(context)
await call_next(context)
execution_order.append("class_after")
async def function_chat_middleware(
context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]
context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]
) -> None:
execution_order.append("function_before")
await next(context)
await call_next(context)
execution_order.append("function_after")
pipeline = ChatMiddlewarePipeline(ClassChatMiddleware(), function_chat_middleware)
@@ -981,21 +997,27 @@ class TestMultipleMiddlewareOrdering:
execution_order: list[str] = []
class FirstMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
execution_order.append("first_before")
await next(context)
await call_next(context)
execution_order.append("first_after")
class SecondMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
execution_order.append("second_before")
await next(context)
await call_next(context)
execution_order.append("second_after")
class ThirdMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
execution_order.append("third_before")
await next(context)
await call_next(context)
execution_order.append("third_after")
middleware = [FirstMiddleware(), SecondMiddleware(), ThirdMiddleware()]
@@ -1029,20 +1051,20 @@ class TestMultipleMiddlewareOrdering:
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
execution_order.append("first_before")
await next(context)
await call_next(context)
execution_order.append("first_after")
class SecondMiddleware(FunctionMiddleware):
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
execution_order.append("second_before")
await next(context)
await call_next(context)
execution_order.append("second_after")
middleware = [FirstMiddleware(), SecondMiddleware()]
@@ -1065,21 +1087,21 @@ class TestMultipleMiddlewareOrdering:
execution_order: list[str] = []
class FirstChatMiddleware(ChatMiddleware):
async def process(self, context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def process(self, context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
execution_order.append("first_before")
await next(context)
await call_next(context)
execution_order.append("first_after")
class SecondChatMiddleware(ChatMiddleware):
async def process(self, context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def process(self, context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
execution_order.append("second_before")
await next(context)
await call_next(context)
execution_order.append("second_after")
class ThirdChatMiddleware(ChatMiddleware):
async def process(self, context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def process(self, context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
execution_order.append("third_before")
await next(context)
await call_next(context)
execution_order.append("third_after")
middleware = [FirstChatMiddleware(), SecondChatMiddleware(), ThirdChatMiddleware()]
@@ -1114,7 +1136,9 @@ class TestContextContentValidation:
"""Test that agent context contains expected data."""
class ContextValidationMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
# Verify context has all expected attributes
assert hasattr(context, "agent")
assert hasattr(context, "messages")
@@ -1132,7 +1156,7 @@ class TestContextContentValidation:
# Add custom metadata
context.metadata["validated"] = True
await next(context)
await call_next(context)
middleware = ContextValidationMiddleware()
pipeline = AgentMiddlewarePipeline(middleware)
@@ -1154,7 +1178,7 @@ class TestContextContentValidation:
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
# Verify context has all expected attributes
assert hasattr(context, "function")
@@ -1170,7 +1194,7 @@ class TestContextContentValidation:
# Add custom metadata
context.metadata["validated"] = True
await next(context)
await call_next(context)
middleware = ContextValidationMiddleware()
pipeline = FunctionMiddlewarePipeline(middleware)
@@ -1189,7 +1213,7 @@ class TestContextContentValidation:
"""Test that chat context contains expected data."""
class ChatContextValidationMiddleware(ChatMiddleware):
async def process(self, context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def process(self, context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
# Verify context has all expected attributes
assert hasattr(context, "chat_client")
assert hasattr(context, "messages")
@@ -1211,7 +1235,7 @@ class TestContextContentValidation:
# Add custom metadata
context.metadata["validated"] = True
await next(context)
await call_next(context)
middleware = ChatContextValidationMiddleware()
pipeline = ChatMiddlewarePipeline(middleware)
@@ -1236,9 +1260,11 @@ class TestStreamingScenarios:
streaming_flags: list[bool] = []
class StreamingFlagMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
streaming_flags.append(context.stream)
await next(context)
await call_next(context)
middleware = StreamingFlagMiddleware()
pipeline = AgentMiddlewarePipeline(middleware)
@@ -1276,9 +1302,11 @@ class TestStreamingScenarios:
chunks_processed: list[str] = []
class StreamProcessingMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
chunks_processed.append("before_stream")
await next(context)
await call_next(context)
chunks_processed.append("after_stream")
middleware = StreamProcessingMiddleware()
@@ -1317,9 +1345,9 @@ class TestStreamingScenarios:
streaming_flags: list[bool] = []
class ChatStreamingFlagMiddleware(ChatMiddleware):
async def process(self, context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def process(self, context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
streaming_flags.append(context.stream)
await next(context)
await call_next(context)
middleware = ChatStreamingFlagMiddleware()
pipeline = ChatMiddlewarePipeline(middleware)
@@ -1358,9 +1386,9 @@ class TestStreamingScenarios:
chunks_processed: list[str] = []
class ChatStreamProcessingMiddleware(ChatMiddleware):
async def process(self, context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def process(self, context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
chunks_processed.append("before_stream")
await next(context)
await call_next(context)
chunks_processed.append("after_stream")
middleware = ChatStreamProcessingMiddleware()
@@ -1408,24 +1436,24 @@ class FunctionTestArgs(BaseModel):
class TestAgentMiddleware(AgentMiddleware):
"""Test implementation of AgentMiddleware."""
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
await next(context)
async def process(self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]) -> None:
await call_next(context)
class TestFunctionMiddleware(FunctionMiddleware):
"""Test implementation of FunctionMiddleware."""
async def process(
self, context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
self, context: FunctionInvocationContext, call_next: Callable[[FunctionInvocationContext], Awaitable[None]]
) -> None:
await next(context)
await call_next(context)
class TestChatMiddleware(ChatMiddleware):
"""Test implementation of ChatMiddleware."""
async def process(self, context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
await next(context)
async def process(self, context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
await call_next(context)
class MockFunctionArgs(BaseModel):
@@ -1441,7 +1469,9 @@ class TestMiddlewareExecutionControl:
"""Test that when agent middleware doesn't call next(), no execution happens."""
class NoNextMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
# Don't call next() - this should prevent any execution
pass
@@ -1468,7 +1498,9 @@ class TestMiddlewareExecutionControl:
"""Test that when agent middleware doesn't call next(), no streaming execution happens."""
class NoNextStreamingMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
# Don't call next() - this should prevent any execution
pass
@@ -1505,7 +1537,7 @@ class TestMiddlewareExecutionControl:
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
# Don't call next() - this should prevent any execution
pass
@@ -1534,14 +1566,18 @@ class TestMiddlewareExecutionControl:
execution_order: list[str] = []
class FirstMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
execution_order.append("first")
# Don't call next() - this should stop the pipeline
class SecondMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
execution_order.append("second")
await next(context)
await call_next(context)
pipeline = AgentMiddlewarePipeline(FirstMiddleware(), SecondMiddleware())
messages = [ChatMessage(role="user", text="test")]
@@ -1565,7 +1601,7 @@ class TestMiddlewareExecutionControl:
"""Test that when chat middleware doesn't call next(), no execution happens."""
class NoNextChatMiddleware(ChatMiddleware):
async def process(self, context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def process(self, context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
# Don't call next() - this should prevent any execution
pass
@@ -1593,7 +1629,7 @@ class TestMiddlewareExecutionControl:
"""Test that when chat middleware doesn't call next(), no streaming execution happens."""
class NoNextStreamingChatMiddleware(ChatMiddleware):
async def process(self, context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def process(self, context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
# Don't call next() - this should prevent any execution
pass
@@ -1634,14 +1670,14 @@ class TestMiddlewareExecutionControl:
execution_order: list[str] = []
class FirstChatMiddleware(ChatMiddleware):
async def process(self, context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def process(self, context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
execution_order.append("first")
# Don't call next() - this should stop the pipeline
class SecondChatMiddleware(ChatMiddleware):
async def process(self, context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def process(self, context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
execution_order.append("second")
await next(context)
await call_next(context)
pipeline = ChatMiddlewarePipeline(FirstChatMiddleware(), SecondChatMiddleware())
messages = [ChatMessage(role="user", text="test")]
python/packages/core/tests/core/test_middleware_context_result.py
+36 -22
View File
@@ -43,9 +43,11 @@ class TestResultOverrideMiddleware:
override_response = AgentResponse(messages=[ChatMessage(role="assistant", text="overridden response")])
class ResponseOverrideMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
# Execute the pipeline first, then override the response
await next(context)
await call_next(context)
context.result = override_response
middleware = ResponseOverrideMiddleware()
@@ -77,9 +79,11 @@ class TestResultOverrideMiddleware:
yield AgentResponseUpdate(contents=[Content.from_text(text=" stream")])
class StreamResponseOverrideMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
# Execute the pipeline first, then override the response stream
await next(context)
await call_next(context)
context.result = ResponseStream(override_stream())
middleware = StreamResponseOverrideMiddleware()
@@ -111,10 +115,10 @@ class TestResultOverrideMiddleware:
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
# Execute the pipeline first, then override the result
await next(context)
await call_next(context)
context.result = override_result
middleware = ResultOverrideMiddleware()
@@ -141,9 +145,11 @@ class TestResultOverrideMiddleware:
mock_chat_client = MockChatClient()
class ChatAgentResponseOverrideMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
# Always call next() first to allow execution
await next(context)
await call_next(context)
# Then conditionally override based on content
if any("special" in msg.text for msg in context.messages if msg.text):
context.result = AgentResponse(
@@ -178,13 +184,15 @@ class TestResultOverrideMiddleware:
yield AgentResponseUpdate(contents=[Content.from_text(text=" response!")])
class ChatAgentStreamOverrideMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
# Check if we want to override BEFORE calling next to avoid creating unused streams
if any("custom stream" in msg.text for msg in context.messages if msg.text):
context.result = ResponseStream(custom_stream())
return # Don't call next() - we're overriding the entire result
# Normal case - let the agent handle it
await next(context)
await call_next(context)
# Create ChatAgent with override middleware
middleware = ChatAgentStreamOverrideMiddleware()
@@ -215,10 +223,12 @@ class TestResultOverrideMiddleware:
"""Test that when agent middleware conditionally doesn't call next(), no execution happens."""
class ConditionalNoNextMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], 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)
await call_next(context)
# Otherwise, don't call next() - no execution should happen
middleware = ConditionalNoNextMiddleware()
@@ -259,13 +269,13 @@ class TestResultOverrideMiddleware:
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_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)
await call_next(context)
# Otherwise, don't call next() - no execution should happen
middleware = ConditionalNoNextFunctionMiddleware()
@@ -308,12 +318,14 @@ class TestResultObservability:
observed_responses: list[AgentResponse] = []
class ObservabilityMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
# Context should be empty before next()
assert context.result is None
# Call next to execute
await next(context)
await call_next(context)
# Context should now contain the response for observability
assert context.result is not None
@@ -343,13 +355,13 @@ class TestResultObservability:
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
# Context should be empty before next()
assert context.result is None
# Call next to execute
await next(context)
await call_next(context)
# Context should now contain the result for observability
assert context.result is not None
@@ -374,9 +386,11 @@ class TestResultObservability:
"""Test that middleware can override response after observing execution."""
class PostExecutionOverrideMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
# Call next to execute first
await next(context)
await call_next(context)
# Now observe and conditionally override
assert context.result is not None
@@ -409,10 +423,10 @@ class TestResultObservability:
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
# Call next to execute first
await next(context)
await call_next(context)
# Now observe and conditionally override
assert context.result is not None
python/packages/core/tests/core/test_middleware_with_agent.py
+145 -115
View File
@@ -44,9 +44,11 @@ class TestChatAgentClassBasedMiddleware:
def __init__(self, name: str):
self.name = name
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
execution_order.append(f"{self.name}_before")
await next(context)
await call_next(context)
execution_order.append(f"{self.name}_after")
# Create ChatAgent with middleware
@@ -74,9 +76,9 @@ class TestChatAgentClassBasedMiddleware:
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
await next(context)
await call_next(context)
middleware = TrackingFunctionMiddleware()
ChatAgent(chat_client=chat_client, middleware=[middleware])
@@ -94,10 +96,10 @@ class TestChatAgentClassBasedMiddleware:
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
execution_order.append(f"{self.name}_before")
await next(context)
await call_next(context)
execution_order.append(f"{self.name}_after")
middleware = TrackingFunctionMiddleware("function_middleware")
@@ -120,11 +122,13 @@ class TestChatAgentFunctionBasedMiddleware:
execution_order: list[str] = []
class PreTerminationMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
execution_order.append("middleware_before")
raise MiddlewareTermination
# Code after raise is unreachable
await next(context)
await call_next(context)
execution_order.append("middleware_after")
# Create ChatAgent with terminating middleware
@@ -149,9 +153,11 @@ class TestChatAgentFunctionBasedMiddleware:
execution_order: list[str] = []
class PostTerminationMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
execution_order.append("middleware_before")
await next(context)
await call_next(context)
execution_order.append("middleware_after")
context.terminate = True
@@ -187,12 +193,12 @@ class TestChatAgentFunctionBasedMiddleware:
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
execution_order.append("middleware_before")
context.terminate = True
# We call next() but since terminate=True, subsequent middleware and handler should not execute
await next(context)
await call_next(context)
execution_order.append("middleware_after")
ChatAgent(chat_client=chat_client, middleware=[PreTerminationFunctionMiddleware()], tools=[])
@@ -205,10 +211,10 @@ class TestChatAgentFunctionBasedMiddleware:
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
execution_order.append("middleware_before")
await next(context)
await call_next(context)
execution_order.append("middleware_after")
context.terminate = True
@@ -219,10 +225,10 @@ class TestChatAgentFunctionBasedMiddleware:
execution_order: list[str] = []
async def tracking_agent_middleware(
context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]
context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
execution_order.append("agent_function_before")
await next(context)
await call_next(context)
execution_order.append("agent_function_after")
# Create ChatAgent with function middleware
@@ -246,9 +252,9 @@ class TestChatAgentFunctionBasedMiddleware:
"""Test function-based function middleware with ChatAgent."""
async def tracking_function_middleware(
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
context: FunctionInvocationContext, call_next: Callable[[FunctionInvocationContext], Awaitable[None]]
) -> None:
await next(context)
await call_next(context)
ChatAgent(chat_client=chat_client, middleware=[tracking_function_middleware])
@@ -259,10 +265,10 @@ class TestChatAgentFunctionBasedMiddleware:
execution_order: list[str] = []
async def tracking_function_middleware(
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
context: FunctionInvocationContext, call_next: Callable[[FunctionInvocationContext], Awaitable[None]]
) -> None:
execution_order.append("function_function_before")
await next(context)
await call_next(context)
execution_order.append("function_function_after")
agent = ChatAgent(chat_client=chat_client_base, middleware=[tracking_function_middleware])
@@ -284,10 +290,12 @@ class TestChatAgentStreamingMiddleware:
streaming_flags: list[bool] = []
class StreamingTrackingMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
execution_order.append("middleware_before")
streaming_flags.append(context.stream)
await next(context)
await call_next(context)
execution_order.append("middleware_after")
# Create ChatAgent with middleware
@@ -326,9 +334,11 @@ class TestChatAgentStreamingMiddleware:
streaming_flags: list[bool] = []
class FlagTrackingMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
streaming_flags.append(context.stream)
await next(context)
await call_next(context)
# Create ChatAgent with middleware
middleware = FlagTrackingMiddleware()
@@ -358,9 +368,11 @@ class TestChatAgentMultipleMiddlewareOrdering:
def __init__(self, name: str):
self.name = name
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
execution_order.append(f"{self.name}_before")
await next(context)
await call_next(context)
execution_order.append(f"{self.name}_after")
# Create multiple middleware
@@ -388,33 +400,35 @@ class TestChatAgentMultipleMiddlewareOrdering:
execution_order: list[str] = []
class ClassAgentMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
execution_order.append("class_agent_before")
await next(context)
await call_next(context)
execution_order.append("class_agent_after")
async def function_agent_middleware(
context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]
context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
execution_order.append("function_agent_before")
await next(context)
await call_next(context)
execution_order.append("function_agent_after")
class ClassFunctionMiddleware(FunctionMiddleware):
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
execution_order.append("class_function_before")
await next(context)
await call_next(context)
execution_order.append("class_function_after")
async def function_function_middleware(
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
context: FunctionInvocationContext, call_next: Callable[[FunctionInvocationContext], Awaitable[None]]
) -> None:
execution_order.append("function_function_before")
await next(context)
await call_next(context)
execution_order.append("function_function_after")
agent = ChatAgent(
@@ -433,23 +447,25 @@ class TestChatAgentMultipleMiddlewareOrdering:
execution_order: list[str] = []
class ClassAgentMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
execution_order.append("class_agent_before")
await next(context)
await call_next(context)
execution_order.append("class_agent_after")
async def function_agent_middleware(
context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]
context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
execution_order.append("function_agent_before")
await next(context)
await call_next(context)
execution_order.append("function_agent_after")
async def function_function_middleware(
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
context: FunctionInvocationContext, call_next: Callable[[FunctionInvocationContext], Awaitable[None]]
) -> None:
execution_order.append("function_function_before")
await next(context)
await call_next(context)
execution_order.append("function_function_after")
agent = ChatAgent(
@@ -505,10 +521,10 @@ class TestChatAgentFunctionMiddlewareWithTools:
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
execution_order.append(f"{self.name}_before")
await next(context)
await call_next(context)
execution_order.append(f"{self.name}_after")
# Set up mock to return a function call first, then a regular response
@@ -567,10 +583,10 @@ class TestChatAgentFunctionMiddlewareWithTools:
execution_order: list[str] = []
async def tracking_function_middleware(
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
context: FunctionInvocationContext, call_next: Callable[[FunctionInvocationContext], Awaitable[None]]
) -> None:
execution_order.append("function_middleware_before")
await next(context)
await call_next(context)
execution_order.append("function_middleware_after")
# Set up mock to return a function call first, then a regular response
@@ -631,20 +647,20 @@ class TestChatAgentFunctionMiddlewareWithTools:
async def process(
self,
context: AgentContext,
next: Callable[[AgentContext], Awaitable[None]],
call_next: Callable[[AgentContext], Awaitable[None]],
) -> None:
execution_order.append("agent_middleware_before")
await next(context)
await call_next(context)
execution_order.append("agent_middleware_after")
class TrackingFunctionMiddleware(FunctionMiddleware):
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
execution_order.append("function_middleware_before")
await next(context)
await call_next(context)
execution_order.append("function_middleware_after")
# Set up mock to return a function call first, then a regular response
@@ -712,7 +728,7 @@ class TestChatAgentFunctionMiddlewareWithTools:
@function_middleware
async def kwargs_middleware(
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
context: FunctionInvocationContext, call_next: Callable[[FunctionInvocationContext], Awaitable[None]]
) -> None:
nonlocal middleware_called
middleware_called = True
@@ -732,7 +748,7 @@ class TestChatAgentFunctionMiddlewareWithTools:
modified_kwargs["new_param"] = context.kwargs.get("new_param")
modified_kwargs["custom_param"] = context.kwargs.get("custom_param")
await next(context)
await call_next(context)
chat_client_base.run_responses = [
ChatResponse(
@@ -785,9 +801,9 @@ class TestMiddlewareDynamicRebuild:
self.name = name
self.execution_log = execution_log
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]) -> None:
self.execution_log.append(f"{self.name}_start")
await next(context)
await call_next(context)
self.execution_log.append(f"{self.name}_end")
async def test_middleware_dynamic_rebuild_non_streaming(self, chat_client: "MockChatClient") -> None:
@@ -908,9 +924,9 @@ class TestRunLevelMiddleware:
self.name = name
self.execution_log = execution_log
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]) -> None:
self.execution_log.append(f"{self.name}_start")
await next(context)
await call_next(context)
self.execution_log.append(f"{self.name}_end")
async def test_run_level_middleware_isolation(self, chat_client: "MockChatClient") -> None:
@@ -960,25 +976,29 @@ class TestRunLevelMiddleware:
def __init__(self, name: str):
self.name = name
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
execution_log.append(f"{self.name}_start")
# Set metadata to pass information to run middleware
context.metadata[f"{self.name}_key"] = f"{self.name}_value"
await next(context)
await call_next(context)
execution_log.append(f"{self.name}_end")
class MetadataRunMiddleware(AgentMiddleware):
def __init__(self, name: str):
self.name = name
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
execution_log.append(f"{self.name}_start")
# Read metadata set by agent middleware
for key, value in context.metadata.items():
metadata_log.append(f"{self.name}_reads_{key}:{value}")
# Set run-level metadata
context.metadata[f"{self.name}_key"] = f"{self.name}_value"
await next(context)
await call_next(context)
execution_log.append(f"{self.name}_end")
# Create agent with agent-level middleware
@@ -1029,10 +1049,12 @@ class TestRunLevelMiddleware:
def __init__(self, name: str):
self.name = name
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
execution_log.append(f"{self.name}_start")
streaming_flags.append(context.stream)
await next(context)
await call_next(context)
execution_log.append(f"{self.name}_end")
# Create agent without agent-level middleware
@@ -1071,44 +1093,48 @@ class TestRunLevelMiddleware:
# Agent-level middleware
class AgentLevelAgentMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
execution_log.append("agent_level_agent_start")
context.metadata["agent_level_agent"] = "processed"
await next(context)
await call_next(context)
execution_log.append("agent_level_agent_end")
class AgentLevelFunctionMiddleware(FunctionMiddleware):
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
execution_log.append("agent_level_function_start")
context.metadata["agent_level_function"] = "processed"
await next(context)
await call_next(context)
execution_log.append("agent_level_function_end")
# Run-level middleware
class RunLevelAgentMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
execution_log.append("run_level_agent_start")
# Verify agent-level middleware metadata is available
assert "agent_level_agent" in context.metadata
context.metadata["run_level_agent"] = "processed"
await next(context)
await call_next(context)
execution_log.append("run_level_agent_end")
class RunLevelFunctionMiddleware(FunctionMiddleware):
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
execution_log.append("run_level_function_start")
# Verify agent-level function middleware metadata is available
assert "agent_level_function" in context.metadata
context.metadata["run_level_function"] = "processed"
await next(context)
await call_next(context)
execution_log.append("run_level_function_end")
# Create tool function for testing function middleware
@@ -1192,17 +1218,17 @@ class TestMiddlewareDecoratorLogic:
@agent_middleware
async def matching_agent_middleware(
context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]
context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
execution_order.append("decorator_type_match_agent")
await next(context)
await call_next(context)
@function_middleware
async def matching_function_middleware(
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
context: FunctionInvocationContext, call_next: Callable[[FunctionInvocationContext], Awaitable[None]]
) -> None:
execution_order.append("decorator_type_match_function")
await next(context)
await call_next(context)
# Create tool function for testing function middleware
def custom_tool(message: str) -> str:
@@ -1254,9 +1280,9 @@ class TestMiddlewareDecoratorLogic:
@agent_middleware # type: ignore[arg-type]
async def mismatched_middleware(
context: FunctionInvocationContext, # Wrong type for @agent_middleware
next: Any,
call_next: Any,
) -> None:
await next(context)
await call_next(context)
agent = ChatAgent(chat_client=chat_client, middleware=[mismatched_middleware])
await agent.run([ChatMessage(role="user", text="test")])
@@ -1266,14 +1292,14 @@ class TestMiddlewareDecoratorLogic:
execution_order: list[str] = []
@agent_middleware
async def decorator_only_agent(context: Any, next: Any) -> None: # No type annotation
async def decorator_only_agent(context: Any, call_next: Any) -> None: # No type annotation
execution_order.append("decorator_only_agent")
await next(context)
await call_next(context)
@function_middleware
async def decorator_only_function(context: Any, next: Any) -> None: # No type annotation
async def decorator_only_function(context: Any, call_next: Any) -> None: # No type annotation
execution_order.append("decorator_only_function")
await next(context)
await call_next(context)
# Create tool function for testing function middleware
def custom_tool(message: str) -> str:
@@ -1320,16 +1346,16 @@ class TestMiddlewareDecoratorLogic:
execution_order: list[str] = []
# No decorator
async def type_only_agent(context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def type_only_agent(context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]) -> None:
execution_order.append("type_only_agent")
await next(context)
await call_next(context)
# No decorator
async def type_only_function(
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
context: FunctionInvocationContext, call_next: Callable[[FunctionInvocationContext], Awaitable[None]]
) -> None:
execution_order.append("type_only_function")
await next(context)
await call_next(context)
# Create tool function for testing function middleware
def custom_tool(message: str) -> str:
@@ -1372,8 +1398,8 @@ class TestMiddlewareDecoratorLogic:
async def test_neither_decorator_nor_type(self, chat_client: Any) -> None:
"""Neither decorator nor parameter type specified - should throw exception."""
async def no_info_middleware(context: Any, next: Any) -> None: # No decorator, no type
await next(context)
async def no_info_middleware(context: Any, call_next: Any) -> None: # No decorator, no type
await call_next(context)
# Should raise MiddlewareException
with pytest.raises(MiddlewareException, match="Cannot determine middleware type"):
@@ -1398,11 +1424,11 @@ class TestMiddlewareDecoratorLogic:
"""Test that decorator markers are properly set on functions."""
@agent_middleware
async def test_agent_middleware(context: Any, next: Any) -> None:
async def test_agent_middleware(context: Any, call_next: Any) -> None:
pass
@function_middleware
async def test_function_middleware(context: Any, next: Any) -> None:
async def test_function_middleware(context: Any, call_next: Any) -> None:
pass
# Check that decorator markers were set
@@ -1421,7 +1447,9 @@ class TestChatAgentThreadBehavior:
thread_states: list[dict[str, Any]] = []
class ThreadTrackingMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def process(
self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
# Capture state before next() call
thread_messages = []
if context.thread and context.thread.message_store:
@@ -1436,7 +1464,7 @@ class TestChatAgentThreadBehavior:
}
thread_states.append(before_state)
await next(context)
await call_next(context)
# Capture state after next() call
thread_messages_after = []
@@ -1532,9 +1560,9 @@ class TestChatAgentChatMiddleware:
execution_order: list[str] = []
class TrackingChatMiddleware(ChatMiddleware):
async def process(self, context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def process(self, context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
execution_order.append("chat_middleware_before")
await next(context)
await call_next(context)
execution_order.append("chat_middleware_after")
# Create ChatAgent with chat middleware
@@ -1561,10 +1589,10 @@ class TestChatAgentChatMiddleware:
execution_order: list[str] = []
async def tracking_chat_middleware(
context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]
context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]
) -> None:
execution_order.append("chat_middleware_before")
await next(context)
await call_next(context)
execution_order.append("chat_middleware_after")
# Create ChatAgent with function-based chat middleware
@@ -1590,7 +1618,7 @@ class TestChatAgentChatMiddleware:
@chat_middleware
async def message_modifier_middleware(
context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]
context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]
) -> None:
# Modify the first message by adding a prefix
if context.messages:
@@ -1600,7 +1628,7 @@ class TestChatAgentChatMiddleware:
original_text = msg.text or ""
context.messages[idx] = ChatMessage(role=msg.role, text=f"MODIFIED: {original_text}")
break
await next(context)
await call_next(context)
# Create ChatAgent with message-modifying middleware
chat_client = MockBaseChatClient()
@@ -1619,7 +1647,7 @@ class TestChatAgentChatMiddleware:
@chat_middleware
async def response_override_middleware(
context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]
context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]
) -> None:
# Override the response without calling next()
context.result = ChatResponse(
@@ -1647,15 +1675,15 @@ class TestChatAgentChatMiddleware:
execution_order: list[str] = []
@chat_middleware
async def first_middleware(context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def first_middleware(context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
execution_order.append("first_before")
await next(context)
await call_next(context)
execution_order.append("first_after")
@chat_middleware
async def second_middleware(context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def second_middleware(context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
execution_order.append("second_before")
await next(context)
await call_next(context)
execution_order.append("second_after")
# Create ChatAgent with multiple chat middleware
@@ -1681,10 +1709,10 @@ class TestChatAgentChatMiddleware:
streaming_flags: list[bool] = []
class StreamingTrackingChatMiddleware(ChatMiddleware):
async def process(self, context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def process(self, context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
execution_order.append("streaming_chat_before")
streaming_flags.append(context.stream)
await next(context)
await call_next(context)
execution_order.append("streaming_chat_after")
# Create ChatAgent with chat middleware
@@ -1721,13 +1749,13 @@ class TestChatAgentChatMiddleware:
execution_order: list[str] = []
class PreTerminationChatMiddleware(ChatMiddleware):
async def process(self, context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def process(self, context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
execution_order.append("middleware_before")
# Set a custom response since we're terminating
context.result = ChatResponse(messages=[ChatMessage(role="assistant", text="Terminated by middleware")])
raise MiddlewareTermination
# We call next() but since terminate=True, execution should stop
await next(context)
await call_next(context)
execution_order.append("middleware_after")
# Create ChatAgent with terminating middleware
@@ -1749,9 +1777,9 @@ class TestChatAgentChatMiddleware:
execution_order: list[str] = []
class PostTerminationChatMiddleware(ChatMiddleware):
async def process(self, context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def process(self, context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
execution_order.append("middleware_before")
await next(context)
await call_next(context)
execution_order.append("middleware_after")
context.terminate = True
@@ -1776,21 +1804,21 @@ class TestChatAgentChatMiddleware:
"""Test ChatAgent with combined middleware types."""
execution_order: list[str] = []
async def agent_middleware(context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def agent_middleware(context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]) -> None:
execution_order.append("agent_middleware_before")
await next(context)
await call_next(context)
execution_order.append("agent_middleware_after")
async def chat_middleware(context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def chat_middleware(context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
execution_order.append("chat_middleware_before")
await next(context)
await call_next(context)
execution_order.append("chat_middleware_after")
async def function_middleware(
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
context: FunctionInvocationContext, call_next: Callable[[FunctionInvocationContext], Awaitable[None]]
) -> None:
execution_order.append("function_middleware_before")
await next(context)
await call_next(context)
execution_order.append("function_middleware_after")
# Create ChatAgent with function middleware and tools
@@ -1814,7 +1842,9 @@ class TestChatAgentChatMiddleware:
modified_kwargs: dict[str, Any] = {}
@agent_middleware
async def kwargs_middleware(context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def kwargs_middleware(
context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
) -> None:
# Capture the original kwargs
captured_kwargs.update(context.kwargs)
@@ -1826,7 +1856,7 @@ class TestChatAgentChatMiddleware:
# Store modified kwargs for verification
modified_kwargs.update(context.kwargs)
await next(context)
await call_next(context)
# Create ChatAgent with agent middleware
chat_client = MockBaseChatClient()
@@ -1865,10 +1895,10 @@ class TestChatAgentChatMiddleware:
# class TrackingMiddleware(AgentMiddleware):
# async def process(
# self, context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]
# self, context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]
# ) -> None:
# execution_order.append("before")
# await next(context)
# await call_next(context)
# execution_order.append("after")
# @use_agent_middleware
python/packages/core/tests/core/test_middleware_with_chat.py
+33 -27
View File
@@ -32,10 +32,10 @@ class TestChatMiddleware:
async def process(
self,
context: ChatContext,
next: Callable[[ChatContext], Awaitable[None]],
call_next: Callable[[ChatContext], Awaitable[None]],
) -> None:
execution_order.append("chat_middleware_before")
await next(context)
await call_next(context)
execution_order.append("chat_middleware_after")
# Add middleware to chat client
@@ -58,9 +58,11 @@ class TestChatMiddleware:
execution_order: list[str] = []
@chat_middleware
async def logging_chat_middleware(context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def logging_chat_middleware(
context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]
) -> None:
execution_order.append("function_middleware_before")
await next(context)
await call_next(context)
execution_order.append("function_middleware_after")
# Add middleware to chat client
@@ -83,13 +85,13 @@ class TestChatMiddleware:
@chat_middleware
async def message_modifier_middleware(
context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]
context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]
) -> None:
# Modify the first message by adding a prefix
if context.messages and len(context.messages) > 0:
original_text = context.messages[0].text or ""
context.messages[0] = ChatMessage(role=context.messages[0].role, text=f"MODIFIED: {original_text}")
await next(context)
await call_next(context)
# Add middleware to chat client
chat_client_base.chat_middleware = [message_modifier_middleware]
@@ -109,7 +111,7 @@ class TestChatMiddleware:
@chat_middleware
async def response_override_middleware(
context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]
context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]
) -> None:
# Override the response without calling next()
context.result = ChatResponse(
@@ -136,15 +138,15 @@ class TestChatMiddleware:
execution_order: list[str] = []
@chat_middleware
async def first_middleware(context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def first_middleware(context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
execution_order.append("first_before")
await next(context)
await call_next(context)
execution_order.append("first_after")
@chat_middleware
async def second_middleware(context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def second_middleware(context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
execution_order.append("second_before")
await next(context)
await call_next(context)
execution_order.append("second_after")
# Add middleware to chat client (order should be preserved)
@@ -172,10 +174,10 @@ class TestChatMiddleware:
@chat_middleware
async def agent_level_chat_middleware(
context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]
context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]
) -> None:
execution_order.append("agent_chat_middleware_before")
await next(context)
await call_next(context)
execution_order.append("agent_chat_middleware_after")
chat_client = MockBaseChatClient()
@@ -203,15 +205,15 @@ class TestChatMiddleware:
execution_order: list[str] = []
@chat_middleware
async def first_middleware(context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def first_middleware(context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
execution_order.append("first_before")
await next(context)
await call_next(context)
execution_order.append("first_after")
@chat_middleware
async def second_middleware(context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def second_middleware(context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
execution_order.append("second_before")
await next(context)
await call_next(context)
execution_order.append("second_after")
# Create ChatAgent with multiple chat middleware
@@ -238,7 +240,9 @@ class TestChatMiddleware:
execution_order: list[str] = []
@chat_middleware
async def streaming_middleware(context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def streaming_middleware(
context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]
) -> None:
execution_order.append("streaming_before")
# Verify it's a streaming context
assert context.stream is True
@@ -250,7 +254,7 @@ class TestChatMiddleware:
return update
context.stream_transform_hooks.append(upper_case_update)
await next(context)
await call_next(context)
execution_order.append("streaming_after")
# Add middleware to chat client
@@ -274,9 +278,11 @@ class TestChatMiddleware:
execution_count = {"count": 0}
@chat_middleware
async def counting_middleware(context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def counting_middleware(
context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]
) -> None:
execution_count["count"] += 1
await next(context)
await call_next(context)
# First call with run-level middleware
messages = [ChatMessage(role="user", text="first message")]
@@ -304,7 +310,7 @@ class TestChatMiddleware:
modified_kwargs: dict[str, Any] = {}
@chat_middleware
async def kwargs_middleware(context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def kwargs_middleware(context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
# Capture the original kwargs
captured_kwargs.update(context.kwargs)
@@ -316,7 +322,7 @@ class TestChatMiddleware:
# Store modified kwargs for verification
modified_kwargs.update(context.kwargs)
await next(context)
await call_next(context)
# Add middleware to chat client
chat_client_base.chat_middleware = [kwargs_middleware]
@@ -349,11 +355,11 @@ class TestChatMiddleware:
@function_middleware
async def test_function_middleware(
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
context: FunctionInvocationContext, call_next: Callable[[FunctionInvocationContext], Awaitable[None]]
) -> None:
nonlocal execution_order
execution_order.append(f"function_middleware_before_{context.function.name}")
await next(context)
await call_next(context)
execution_order.append(f"function_middleware_after_{context.function.name}")
# Define a simple tool function
@@ -415,10 +421,10 @@ class TestChatMiddleware:
@function_middleware
async def run_level_function_middleware(
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
context: FunctionInvocationContext, call_next: Callable[[FunctionInvocationContext], Awaitable[None]]
) -> None:
execution_order.append("run_level_function_middleware_before")
await next(context)
await call_next(context)
execution_order.append("run_level_function_middleware_after")
# Define a simple tool function