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/AGENTS.md
+3 -4
View File
@@ -117,11 +117,10 @@ agent = OpenAIChatClient().as_agent(
from agent_framework import ChatAgent, AgentMiddleware, AgentContext
class LoggingMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next) -> AgentResponse:
async def process(self, context: AgentContext, call_next) -> None:
print(f"Input: {context.messages}")
response = await next(context)
print(f"Output: {response}")
return response
await call_next(context)
print(f"Output: {context.result}")
agent = ChatAgent(..., middleware=[LoggingMiddleware()])
```
python/packages/core/agent_framework/_middleware.py
+37 -37
View File
@@ -122,7 +122,7 @@ class AgentContext:
options: The options for the agent invocation as a dict.
stream: Whether this is a streaming invocation.
metadata: Metadata dictionary for sharing data between agent middleware.
result: Agent execution result. Can be observed after calling ``next()``
result: Agent execution result. Can be observed after calling ``call_next()``
to see the actual execution result or can be set to override the execution result.
For non-streaming: should be AgentResponse.
For streaming: should be ResponseStream[AgentResponseUpdate, AgentResponse].
@@ -135,7 +135,7 @@ class AgentContext:
class LoggingMiddleware(AgentMiddleware):
async def process(self, context: AgentContext, next):
async def process(self, context: AgentContext, call_next):
print(f"Agent: {context.agent.name}")
print(f"Messages: {len(context.messages)}")
print(f"Thread: {context.thread}")
@@ -145,7 +145,7 @@ class AgentContext:
context.metadata["start_time"] = time.time()
# Continue execution
await next(context)
await call_next(context)
# Access result after execution
print(f"Result: {context.result}")
@@ -208,7 +208,7 @@ class FunctionInvocationContext:
function: The function being invoked.
arguments: The validated arguments for the function.
metadata: Metadata dictionary for sharing data between function middleware.
result: Function execution result. Can be observed after calling ``next()``
result: Function execution result. Can be observed after calling ``call_next()``
to see the actual execution result or can be set to override the execution result.
kwargs: Additional keyword arguments passed to the chat method that invoked this function.
@@ -220,7 +220,7 @@ class FunctionInvocationContext:
class ValidationMiddleware(FunctionMiddleware):
async def process(self, context: FunctionInvocationContext, next):
async def process(self, context: FunctionInvocationContext, call_next):
print(f"Function: {context.function.name}")
print(f"Arguments: {context.arguments}")
@@ -229,7 +229,7 @@ class FunctionInvocationContext:
raise MiddlewareTermination("Validation failed")
# Continue execution
await next(context)
await call_next(context)
"""
def __init__(
@@ -268,7 +268,7 @@ class ChatContext:
options: The options for the chat request as a dict.
stream: Whether this is a streaming invocation.
metadata: Metadata dictionary for sharing data between chat middleware.
result: Chat execution result. Can be observed after calling ``next()``
result: Chat execution result. Can be observed after calling ``call_next()``
to see the actual execution result or can be set to override the execution result.
For non-streaming: should be ChatResponse.
For streaming: should be ResponseStream[ChatResponseUpdate, ChatResponse].
@@ -284,7 +284,7 @@ class ChatContext:
class TokenCounterMiddleware(ChatMiddleware):
async def process(self, context: ChatContext, next):
async def process(self, context: ChatContext, call_next):
print(f"Chat client: {context.chat_client.__class__.__name__}")
print(f"Messages: {len(context.messages)}")
print(f"Model: {context.options.get('model_id')}")
@@ -293,7 +293,7 @@ class ChatContext:
context.metadata["input_tokens"] = self.count_tokens(context.messages)
# Continue execution
await next(context)
await call_next(context)
# Access result and count output tokens
if context.result:
@@ -363,9 +363,9 @@ class AgentMiddleware(ABC):
def __init__(self, max_retries: int = 3):
self.max_retries = max_retries
async def process(self, context: AgentContext, next):
async def process(self, context: AgentContext, call_next):
for attempt in range(self.max_retries):
await next(context)
await call_next(context)
if context.result and not context.result.is_error:
break
print(f"Retry {attempt + 1}/{self.max_retries}")
@@ -379,7 +379,7 @@ class AgentMiddleware(ABC):
async def process(
self,
context: AgentContext,
next: Callable[[AgentContext], Awaitable[None]],
call_next: Callable[[AgentContext], Awaitable[None]],
) -> None:
"""Process an agent invocation.
@@ -387,16 +387,16 @@ class AgentMiddleware(ABC):
context: Agent invocation context containing agent, messages, and metadata.
Use context.stream to determine if this is a streaming call.
MiddlewareTypes can set context.result to override execution, or observe
the actual execution result after calling next().
the actual execution result after calling call_next().
For non-streaming: AgentResponse
For streaming: AsyncIterable[AgentResponseUpdate]
next: Function to call the next middleware or final agent execution.
call_next: Function to call the next middleware or final agent execution.
Does not return anything - all data flows through the context.
Note:
MiddlewareTypes should not return anything. All data manipulation should happen
within the context object. Set context.result to override execution,
or observe context.result after calling next() for actual results.
or observe context.result after calling call_next() for actual results.
"""
...
@@ -422,7 +422,7 @@ class FunctionMiddleware(ABC):
def __init__(self):
self.cache = {}
async def process(self, context: FunctionInvocationContext, next):
async def process(self, context: FunctionInvocationContext, call_next):
cache_key = f"{context.function.name}:{context.arguments}"
# Check cache
@@ -431,7 +431,7 @@ class FunctionMiddleware(ABC):
raise MiddlewareTermination()
# Execute function
await next(context)
await call_next(context)
# Cache result
if context.result:
@@ -446,21 +446,21 @@ class FunctionMiddleware(ABC):
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
"""Process a function invocation.
Args:
context: Function invocation context containing function, arguments, and metadata.
MiddlewareTypes can set context.result to override execution, or observe
the actual execution result after calling next().
next: Function to call the next middleware or final function execution.
the actual execution result after calling call_next().
call_next: Function to call the next middleware or final function execution.
Does not return anything - all data flows through the context.
Note:
MiddlewareTypes should not return anything. All data manipulation should happen
within the context object. Set context.result to override execution,
or observe context.result after calling next() for actual results.
or observe context.result after calling call_next() for actual results.
"""
...
@@ -486,14 +486,14 @@ class ChatMiddleware(ABC):
def __init__(self, system_prompt: str):
self.system_prompt = system_prompt
async def process(self, context: ChatContext, next):
async def process(self, context: ChatContext, call_next):
# Add system prompt to messages
from agent_framework import ChatMessage
context.messages.insert(0, ChatMessage(role="system", text=self.system_prompt))
# Continue execution
await next(context)
await call_next(context)
# Use with an agent
@@ -508,7 +508,7 @@ class ChatMiddleware(ABC):
async def process(
self,
context: ChatContext,
next: Callable[[ChatContext], Awaitable[None]],
call_next: Callable[[ChatContext], Awaitable[None]],
) -> None:
"""Process a chat client request.
@@ -516,16 +516,16 @@ class ChatMiddleware(ABC):
context: Chat invocation context containing chat client, messages, options, and metadata.
Use context.stream to determine if this is a streaming call.
MiddlewareTypes can set context.result to override execution, or observe
the actual execution result after calling next().
the actual execution result after calling call_next().
For non-streaming: ChatResponse
For streaming: ResponseStream[ChatResponseUpdate, ChatResponse]
next: Function to call the next middleware or final chat execution.
call_next: Function to call the next middleware or final chat execution.
Does not return anything - all data flows through the context.
Note:
MiddlewareTypes should not return anything. All data manipulation should happen
within the context object. Set context.result to override execution,
or observe context.result after calling next() for actual results.
or observe context.result after calling call_next() for actual results.
"""
...
@@ -576,9 +576,9 @@ def agent_middleware(func: AgentMiddlewareCallable) -> AgentMiddlewareCallable:
@agent_middleware
async def logging_middleware(context: AgentContext, next):
async def logging_middleware(context: AgentContext, call_next):
print(f"Before: {context.agent.name}")
await next(context)
await call_next(context)
print(f"After: {context.result}")
@@ -609,9 +609,9 @@ def function_middleware(func: FunctionMiddlewareCallable) -> FunctionMiddlewareC
@function_middleware
async def logging_middleware(context: FunctionInvocationContext, next):
async def logging_middleware(context: FunctionInvocationContext, call_next):
print(f"Calling: {context.function.name}")
await next(context)
await call_next(context)
print(f"Result: {context.result}")
@@ -642,9 +642,9 @@ def chat_middleware(func: ChatMiddlewareCallable) -> ChatMiddlewareCallable:
@chat_middleware
async def logging_middleware(context: ChatContext, next):
async def logging_middleware(context: ChatContext, call_next):
print(f"Messages: {len(context.messages)}")
await next(context)
await call_next(context)
print(f"Response: {context.result}")
@@ -669,8 +669,8 @@ class MiddlewareWrapper(Generic[TContext]):
def __init__(self, func: Callable[[TContext, Callable[[TContext], Awaitable[None]]], Awaitable[None]]) -> None:
self.func = func
async def process(self, context: TContext, next: Callable[[TContext], Awaitable[None]]) -> None:
await self.func(context, next)
async def process(self, context: TContext, call_next: Callable[[TContext], Awaitable[None]]) -> None:
await self.func(context, call_next)
class BaseMiddlewarePipeline(ABC):
@@ -1226,7 +1226,7 @@ def _determine_middleware_type(middleware: Any) -> MiddlewareType:
sig = inspect.signature(middleware)
params = list(sig.parameters.values())
# Must have at least 2 parameters (context and next)
# Must have at least 2 parameters (context and call_next)
if len(params) >= 2:
first_param = params[0]
if hasattr(first_param.annotation, "__name__"):
@@ -1240,7 +1240,7 @@ def _determine_middleware_type(middleware: Any) -> MiddlewareType:
else:
# Not enough parameters - can't be valid middleware
raise MiddlewareException(
f"MiddlewareTypes function must have at least 2 parameters (context, next), "
f"Middleware function must have at least 2 parameters (context, call_next), "
f"but {middleware.__name__} has {len(params)}"
)
except Exception as e:
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
python/packages/ollama/tests/test_ollama_chat_client.py
+2 -2
View File
@@ -207,8 +207,8 @@ def test_serialize(ollama_unit_test_env: dict[str, str]) -> None:
def test_chat_middleware(ollama_unit_test_env: dict[str, str]) -> None:
@chat_middleware
async def sample_middleware(context, next):
await next(context)
async def sample_middleware(context, call_next):
await call_next(context)
ollama_chat_client = OllamaChatClient(middleware=[sample_middleware])
assert len(ollama_chat_client.middleware) == 1
python/packages/orchestrations/agent_framework_orchestrations/_handoff.py
+2 -2
View File
@@ -129,11 +129,11 @@ class _AutoHandoffMiddleware(FunctionMiddleware):
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
"""Intercept matching handoff tool calls and inject synthetic results."""
if context.function.name not in self._handoff_functions:
await next(context)
await call_next(context)
return
from agent_framework._middleware import MiddlewareTermination
python/packages/purview/agent_framework_purview/_middleware.py
+4 -4
View File
@@ -48,7 +48,7 @@ class PurviewPolicyMiddleware(AgentMiddleware):
async def process(
self,
context: AgentContext,
next: Callable[[AgentContext], Awaitable[None]],
call_next: Callable[[AgentContext], Awaitable[None]],
) -> None: # type: ignore[override]
resolved_user_id: str | None = None
try:
@@ -74,7 +74,7 @@ class PurviewPolicyMiddleware(AgentMiddleware):
if not self._settings.ignore_exceptions:
raise
await next(context)
await call_next(context)
try:
# Post (response) check only if we have a normal AgentResponse
@@ -140,7 +140,7 @@ class PurviewChatPolicyMiddleware(ChatMiddleware):
async def process(
self,
context: ChatContext,
next: Callable[[ChatContext], Awaitable[None]],
call_next: Callable[[ChatContext], Awaitable[None]],
) -> None: # type: ignore[override]
resolved_user_id: str | None = None
try:
@@ -164,7 +164,7 @@ class PurviewChatPolicyMiddleware(ChatMiddleware):
if not self._settings.ignore_exceptions:
raise
await next(context)
await call_next(context)
try:
# Post (response) evaluation only if non-streaming and we have messages result shape