agent/agent-framework
1
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
python/samples/concepts/tools/README.md
+32 -32
View File
@@ -37,7 +37,7 @@ sequenceDiagram
AML->>AMP: execute(AgentContext)
loop Agent Middleware Chain
AMP->>AMP: middleware[i].process(context, next)
AMP->>AMP: middleware[i].process(context, call_next)
Note right of AMP: Can modify: messages, options, thread
end
@@ -60,7 +60,7 @@ sequenceDiagram
CML->>CMP: execute(ChatContext)
loop Chat Middleware Chain
CMP->>CMP: middleware[i].process(context, next)
CMP->>CMP: middleware[i].process(context, call_next)
Note right of CMP: Can modify: messages, options
end
@@ -81,7 +81,7 @@ sequenceDiagram
loop For each function_call
FIL->>FMP: execute(FunctionInvocationContext)
loop Function Middleware Chain
FMP->>FMP: middleware[i].process(context, next)
FMP->>FMP: middleware[i].process(context, call_next)
Note right of FMP: Can modify: arguments
end
FMP->>Tool: invoke(arguments)
@@ -137,7 +137,7 @@ sequenceDiagram
| `options` | `Mapping[str, Any]` | Chat options dict |
| `stream` | `bool` | Whether streaming is enabled |
| `metadata` | `dict` | Shared data between middleware |
| `result` | `AgentResponse \| None` | Set after `next()` is called |
| `result` | `AgentResponse \| None` | Set after `call_next()` is called |
| `kwargs` | `Mapping[str, Any]` | Additional run arguments |
**Key Operations:**
@@ -150,7 +150,7 @@ sequenceDiagram
- `context.messages` - Add, remove, or modify input messages
- `context.options` - Change model parameters, temperature, etc.
- `context.thread` - Replace or modify the thread
- `context.result` - Override the final response (after `next()`)
- `context.result` - Override the final response (after `call_next()`)
### 2. Chat Middleware Layer (`ChatMiddlewareLayer`)
@@ -165,7 +165,7 @@ sequenceDiagram
| `options` | `Mapping[str, Any]` | Chat options |
| `stream` | `bool` | Whether streaming |
| `metadata` | `dict` | Shared data between middleware |
| `result` | `ChatResponse \| None` | Set after `next()` is called |
| `result` | `ChatResponse \| None` | Set after `call_next()` is called |
| `kwargs` | `Mapping[str, Any]` | Additional arguments |
**Key Operations:**
@@ -176,7 +176,7 @@ sequenceDiagram
**What Can Be Modified:**
- `context.messages` - Inject system prompts, filter content
- `context.options` - Change model, temperature, tool_choice
- `context.result` - Override the response (after `next()`)
- `context.result` - Override the response (after `call_next()`)
### 3. Function Invocation Layer (`FunctionInvocationLayer`)
@@ -251,23 +251,23 @@ response = await client.get_response(
| `function` | `FunctionTool` | The function being invoked |
| `arguments` | `BaseModel` | Validated Pydantic arguments |
| `metadata` | `dict` | Shared data between middleware |
| `result` | `Any` | Set after `next()` is called |
| `result` | `Any` | Set after `call_next()` is called |
| `kwargs` | `Mapping[str, Any]` | Runtime kwargs |
**What Can Be Modified:**
- `context.arguments` - Modify validated arguments before execution
- `context.result` - Override the function result (after `next()`)
- `context.result` - Override the function result (after `call_next()`)
- Raise `MiddlewareTermination` to skip execution and terminate the function invocation loop
**Special Behavior:** When `MiddlewareTermination` is raised in function middleware, it signals that the function invocation loop should exit **without making another LLM call**. This is useful when middleware determines that no further processing is needed (e.g., a termination condition is met).
```python
class TerminatingMiddleware(FunctionMiddleware):
async def process(self, context: FunctionInvocationContext, next):
async def process(self, context: FunctionInvocationContext, call_next):
if self.should_terminate(context):
context.result = "terminated by middleware"
raise MiddlewareTermination # Exit function invocation loop
await next(context)
await call_next(context)
```
## Arguments Added/Altered at Each Layer
@@ -334,20 +334,20 @@ class TerminatingMiddleware(FunctionMiddleware):
There are three ways to exit a middleware's `process()` method:
### 1. Return Normally (with or without calling `next`)
### 1. Return Normally (with or without calling `call_next`)
Returns control to the upstream middleware, allowing its post-processing code to run.
```python
class CachingMiddleware(FunctionMiddleware):
async def process(self, context: FunctionInvocationContext, next):
# Option A: Return early WITHOUT calling next (skip downstream)
async def process(self, context: FunctionInvocationContext, call_next):
# Option A: Return early WITHOUT calling call_next (skip downstream)
if cached := self.cache.get(context.function.name):
context.result = cached
return # Upstream post-processing still runs
# Option B: Call next, then return normally
await next(context)
# Option B: Call call_next, then return normally
await call_next(context)
self.cache[context.function.name] = context.result
return # Normal completion
```
@@ -358,11 +358,11 @@ Immediately exits the entire middleware chain. Upstream middleware's post-proces
```python
class BlockedFunctionMiddleware(FunctionMiddleware):
async def process(self, context: FunctionInvocationContext, next):
async def process(self, context: FunctionInvocationContext, call_next):
if context.function.name in self.blocked_functions:
context.result = "Function blocked by policy"
raise MiddlewareTermination("Blocked") # Skips ALL post-processing
await next(context)
await call_next(context)
```
### 3. Raise Any Other Exception
@@ -371,10 +371,10 @@ Bubbles up to the caller. The middleware chain is aborted and the exception prop
```python
class ValidationMiddleware(FunctionMiddleware):
async def process(self, context: FunctionInvocationContext, next):
async def process(self, context: FunctionInvocationContext, call_next):
if not self.is_valid(context.arguments):
raise ValueError("Invalid arguments") # Bubbles up to user
await next(context)
await call_next(context)
```
## `return` vs `raise MiddlewareTermination`
@@ -383,13 +383,13 @@ The key difference is what happens to **upstream middleware's post-processing**:
```python
class MiddlewareA(AgentMiddleware):
async def process(self, context, next):
async def process(self, context, call_next):
print("A: before")
await next(context)
await call_next(context)
print("A: after") # Does this run?
class MiddlewareB(AgentMiddleware):
async def process(self, context, next):
async def process(self, context, call_next):
print("B: before")
context.result = "early result"
# Choose one:
@@ -408,14 +408,14 @@ With middleware registered as `[MiddlewareA, MiddlewareB]`:
**Use `raise MiddlewareTermination`** when you want to completely bypass all remaining processing (e.g., blocking a request, returning cached response without any modification).
## Calling `next()` or Not
## Calling `call_next()` or Not
The decision to call `next(context)` determines whether downstream middleware and the actual operation execute:
The decision to call `call_next(context)` determines whether downstream middleware and the actual operation execute:
### Without calling `next()` - Skip downstream
### Without calling `call_next()` - Skip downstream
```python
async def process(self, context, next):
async def process(self, context, call_next):
context.result = "replacement result"
return # Downstream middleware and actual execution are SKIPPED
```
@@ -425,12 +425,12 @@ async def process(self, context, next):
- Upstream middleware post-processing: ✅ Still runs (unless `MiddlewareTermination` raised)
- Result: Whatever you set in `context.result`
### With calling `next()` - Full execution
### With calling `call_next()` - Full execution
```python
async def process(self, context, next):
async def process(self, context, call_next):
# Pre-processing
await next(context) # Execute downstream + actual operation
await call_next(context) # Execute downstream + actual operation
# Post-processing (context.result now contains real result)
return
```
@@ -442,7 +442,7 @@ async def process(self, context, next):
### Summary Table
| Exit Method | Call `next()`? | Downstream Executes? | Actual Op Executes? | Upstream Post-Processing? |
| Exit Method | Call `call_next()`? | Downstream Executes? | Actual Op Executes? | Upstream Post-Processing? |
|-------------|----------------|---------------------|---------------------|--------------------------|
| `return` (or implicit) | Yes | ✅ | ✅ | ✅ Yes |
| `return` | No | ❌ | ❌ | ✅ Yes |
@@ -450,7 +450,7 @@ async def process(self, context, next):
| `raise MiddlewareTermination` | Yes | ✅ | ✅ | ❌ No |
| `raise OtherException` | Either | Depends | Depends | ❌ No (exception propagates) |
> **Note:** The first row (`return` after calling `next()`) is the default behavior. Python functions implicitly return `None` at the end, so simply calling `await next(context)` without an explicit `return` statement achieves this pattern.
> **Note:** The first row (`return` after calling `call_next()`) is the default behavior. Python functions implicitly return `None` at the end, so simply calling `await call_next(context)` without an explicit `return` statement achieves this pattern.
## Streaming vs Non-Streaming
python/samples/getting_started/agents/azure_ai/azure_ai_with_agent_as_tool.py
+2 -2
View File
@@ -20,13 +20,13 @@ multiple specialized agents, each focusing on specific tasks.
async def logging_middleware(
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
"""MiddlewareTypes that logs tool invocations to show the delegation flow."""
print(f"[Calling tool: {context.function.name}]")
print(f"[Request: {context.arguments}]")
await next(context)
await call_next(context)
print(f"[Response: {context.result}]")
python/samples/getting_started/agents/openai/openai_responses_client_basic.py
+2 -2
View File
@@ -28,7 +28,7 @@ response generation, showing both streaming and non-streaming responses.
@chat_middleware
async def security_and_override_middleware(
context: ChatContext,
next: Callable[[ChatContext], Awaitable[None]],
call_next: Callable[[ChatContext], Awaitable[None]],
) -> None:
"""Function-based middleware that implements security filtering and response override."""
print("[SecurityMiddleware] Processing input...")
@@ -59,7 +59,7 @@ async def security_and_override_middleware(
raise MiddlewareTermination
# Continue to next middleware or AI execution
await next(context)
await call_next(context)
print("[SecurityMiddleware] Response generated.")
print(type(context.result))
python/samples/getting_started/agents/openai/openai_responses_client_with_agent_as_tool.py
+2 -2
View File
@@ -19,13 +19,13 @@ multiple specialized agents, each focusing on specific tasks.
async def logging_middleware(
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
"""MiddlewareTypes that logs tool invocations to show the delegation flow."""
print(f"[Calling tool: {context.function.name}]")
print(f"[Request: {context.arguments}]")
await next(context)
await call_next(context)
print(f"[Response: {context.result}]")
python/samples/getting_started/devui/weather_agent_azure/agent.py
+4 -4
View File
@@ -37,7 +37,7 @@ def cleanup_resources():
@chat_middleware
async def security_filter_middleware(
context: ChatContext,
next: Callable[[ChatContext], Awaitable[None]],
call_next: Callable[[ChatContext], Awaitable[None]],
) -> None:
"""Chat middleware that blocks requests containing sensitive information."""
blocked_terms = ["password", "secret", "api_key", "token"]
@@ -76,13 +76,13 @@ async def security_filter_middleware(
raise MiddlewareTermination
await next(context)
await call_next(context)
@function_middleware
async def atlantis_location_filter_middleware(
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
"""Function middleware that blocks weather requests for Atlantis."""
# Check if location parameter is "atlantis"
@@ -94,7 +94,7 @@ async def atlantis_location_filter_middleware(
)
raise MiddlewareTermination
await next(context)
await call_next(context)
# NOTE: approval_mode="never_require" is for sample brevity. Use "always_require" in production; see samples/getting_started/tools/function_tool_with_approval.py and samples/getting_started/tools/function_tool_with_approval_and_threads.py.
python/samples/getting_started/middleware/agent_and_run_level_middleware.py
+14 -14
View File
@@ -49,7 +49,7 @@ def get_weather(
class SecurityAgentMiddleware(AgentMiddleware):
"""Agent-level security middleware that validates all requests."""
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:
print("[SecurityMiddleware] Checking security for all requests...")
# Check for security violations in the last user message
@@ -58,22 +58,22 @@ class SecurityAgentMiddleware(AgentMiddleware):
query = last_message.text.lower()
if any(word in query for word in ["password", "secret", "credentials"]):
print("[SecurityMiddleware] Security violation detected! Blocking request.")
return # Don't call next() to prevent execution
return # Don't call call_next() to prevent execution
print("[SecurityMiddleware] Security check passed.")
context.metadata["security_validated"] = True
await next(context)
await call_next(context)
async def performance_monitor_middleware(
context: AgentContext,
next: Callable[[AgentContext], Awaitable[None]],
call_next: Callable[[AgentContext], Awaitable[None]],
) -> None:
"""Agent-level performance monitoring for all runs."""
print("[PerformanceMonitor] Starting performance monitoring...")
start_time = time.time()
await next(context)
await call_next(context)
end_time = time.time()
duration = end_time - start_time
@@ -85,7 +85,7 @@ async def performance_monitor_middleware(
class HighPriorityMiddleware(AgentMiddleware):
"""Run-level middleware for high priority requests."""
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:
print("[HighPriority] Processing high priority request with expedited handling...")
# Read metadata set by agent-level middleware
@@ -96,13 +96,13 @@ class HighPriorityMiddleware(AgentMiddleware):
context.metadata["priority"] = "high"
context.metadata["expedited"] = True
await next(context)
await call_next(context)
print("[HighPriority] High priority processing completed")
async def debugging_middleware(
context: AgentContext,
next: Callable[[AgentContext], Awaitable[None]],
call_next: Callable[[AgentContext], Awaitable[None]],
) -> None:
"""Run-level debugging middleware for troubleshooting specific runs."""
print("[Debug] Debug mode enabled for this run")
@@ -115,7 +115,7 @@ async def debugging_middleware(
context.metadata["debug_enabled"] = True
await next(context)
await call_next(context)
print("[Debug] Debug information collected")
@@ -126,7 +126,7 @@ class CachingMiddleware(AgentMiddleware):
def __init__(self) -> None:
self.cache: dict[str, AgentResponse] = {}
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:
# Create a simple cache key from the last message
last_message = context.messages[-1] if context.messages else None
cache_key: str = last_message.text if last_message and last_message.text else "no_message"
@@ -134,12 +134,12 @@ class CachingMiddleware(AgentMiddleware):
if cache_key in self.cache:
print(f"[Cache] Cache HIT for: '{cache_key[:30]}...'")
context.result = self.cache[cache_key] # type: ignore
return # Don't call next(), return cached result
return # Don't call call_next(), return cached result
print(f"[Cache] Cache MISS for: '{cache_key[:30]}...'")
context.metadata["cache_key"] = cache_key
await next(context)
await call_next(context)
# Cache the result if we have one
if context.result:
@@ -149,14 +149,14 @@ class CachingMiddleware(AgentMiddleware):
async def function_logging_middleware(
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
"""Function middleware that logs all function calls."""
function_name = context.function.name
args = context.arguments
print(f"[FunctionLog] Calling function: {function_name} with args: {args}")
await next(context)
await call_next(context)
print(f"[FunctionLog] Function {function_name} completed")
python/samples/getting_started/middleware/chat_middleware.py
+4 -4
View File
@@ -57,7 +57,7 @@ class InputObserverMiddleware(ChatMiddleware):
async def process(
self,
context: ChatContext,
next: Callable[[ChatContext], Awaitable[None]],
call_next: Callable[[ChatContext], Awaitable[None]],
) -> None:
"""Observe and modify input messages before they are sent to AI."""
print("[InputObserverMiddleware] Observing input messages:")
@@ -91,7 +91,7 @@ class InputObserverMiddleware(ChatMiddleware):
context.messages[:] = modified_messages
# Continue to next middleware or AI execution
await next(context)
await call_next(context)
# Observe that processing is complete
print("[InputObserverMiddleware] Processing completed")
@@ -100,7 +100,7 @@ class InputObserverMiddleware(ChatMiddleware):
@chat_middleware
async def security_and_override_middleware(
context: ChatContext,
next: Callable[[ChatContext], Awaitable[None]],
call_next: Callable[[ChatContext], Awaitable[None]],
) -> None:
"""Function-based middleware that implements security filtering and response override."""
print("[SecurityMiddleware] Processing input...")
@@ -131,7 +131,7 @@ async def security_and_override_middleware(
raise MiddlewareTermination
# Continue to next middleware or AI execution
await next(context)
await call_next(context)
async def class_based_chat_middleware() -> None:
python/samples/getting_started/middleware/class_based_middleware.py
+5 -5
View File
@@ -50,7 +50,7 @@ class SecurityAgentMiddleware(AgentMiddleware):
async def process(
self,
context: AgentContext,
next: Callable[[AgentContext], Awaitable[None]],
call_next: Callable[[AgentContext], Awaitable[None]],
) -> None:
# Check for potential security violations in the query
# Look at the last user message
@@ -63,11 +63,11 @@ class SecurityAgentMiddleware(AgentMiddleware):
context.result = AgentResponse(
messages=[ChatMessage("assistant", ["Detected sensitive information, the request is blocked."])]
)
# Simply don't call next() to prevent execution
# Simply don't call call_next() to prevent execution
return
print("[SecurityAgentMiddleware] Security check passed.")
await next(context)
await call_next(context)
class LoggingFunctionMiddleware(FunctionMiddleware):
@@ -76,14 +76,14 @@ class LoggingFunctionMiddleware(FunctionMiddleware):
async def process(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
function_name = context.function.name
print(f"[LoggingFunctionMiddleware] About to call function: {function_name}.")
start_time = time.time()
await next(context)
await call_next(context)
end_time = time.time()
duration = end_time - start_time
python/samples/getting_started/middleware/decorator_middleware.py
+4 -4
View File
@@ -50,18 +50,18 @@ def get_current_time() -> str:
@agent_middleware # Decorator marks this as agent middleware - no type annotations needed
async def simple_agent_middleware(context, next): # type: ignore - parameters intentionally untyped to demonstrate decorator functionality
async def simple_agent_middleware(context, call_next): # type: ignore - parameters intentionally untyped to demonstrate decorator functionality
"""Agent middleware that runs before and after agent execution."""
print("[Agent MiddlewareTypes] Before agent execution")
await next(context)
await call_next(context)
print("[Agent MiddlewareTypes] After agent execution")
@function_middleware # Decorator marks this as function middleware - no type annotations needed
async def simple_function_middleware(context, next): # type: ignore - parameters intentionally untyped to demonstrate decorator functionality
async def simple_function_middleware(context, call_next): # type: ignore - parameters intentionally untyped to demonstrate decorator functionality
"""Function middleware that runs before and after function calls."""
print(f"[Function MiddlewareTypes] Before calling: {context.function.name}") # type: ignore
await next(context)
await call_next(context)
print(f"[Function MiddlewareTypes] After calling: {context.function.name}") # type: ignore
python/samples/getting_started/middleware/exception_handling_with_middleware.py
+2 -2
View File
@@ -35,13 +35,13 @@ def unstable_data_service(
async def exception_handling_middleware(
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
context: FunctionInvocationContext, call_next: Callable[[FunctionInvocationContext], Awaitable[None]]
) -> None:
function_name = context.function.name
try:
print(f"[ExceptionHandlingMiddleware] Executing function: {function_name}")
await next(context)
await call_next(context)
print(f"[ExceptionHandlingMiddleware] Function {function_name} completed successfully.")
except TimeoutError as e:
print(f"[ExceptionHandlingMiddleware] Caught TimeoutError: {e}")
python/samples/getting_started/middleware/function_based_middleware.py
+6 -6
View File
@@ -27,7 +27,7 @@ The example includes:
Function-based middleware is ideal for simple, stateless operations and provides a more
lightweight approach compared to class-based middleware. Both agent and function middleware
can be implemented as async functions that accept context and next parameters.
can be implemented as async functions that accept context and call_next parameters.
"""
@@ -43,7 +43,7 @@ def get_weather(
async def security_agent_middleware(
context: AgentContext,
next: Callable[[AgentContext], Awaitable[None]],
call_next: Callable[[AgentContext], Awaitable[None]],
) -> None:
"""Agent middleware that checks for security violations."""
# Check for potential security violations in the query
@@ -53,16 +53,16 @@ async def security_agent_middleware(
query = last_message.text
if "password" in query.lower() or "secret" in query.lower():
print("[SecurityAgentMiddleware] Security Warning: Detected sensitive information, blocking request.")
# Simply don't call next() to prevent execution
# Simply don't call call_next() to prevent execution
return
print("[SecurityAgentMiddleware] Security check passed.")
await next(context)
await call_next(context)
async def logging_function_middleware(
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
"""Function middleware that logs function calls."""
function_name = context.function.name
@@ -70,7 +70,7 @@ async def logging_function_middleware(
start_time = time.time()
await next(context)
await call_next(context)
end_time = time.time()
duration = end_time - start_time
python/samples/getting_started/middleware/middleware_termination.py
+5 -5
View File
@@ -23,7 +23,7 @@ MiddlewareTypes Termination Example
This sample demonstrates how middleware can terminate execution using the `context.terminate` flag.
The example includes:
- PreTerminationMiddleware: Terminates execution before calling next() to prevent agent processing
- PreTerminationMiddleware: Terminates execution before calling call_next() to prevent agent processing
- PostTerminationMiddleware: Allows processing to complete but terminates further execution
This is useful for implementing security checks, rate limiting, or early exit conditions.
@@ -49,7 +49,7 @@ class PreTerminationMiddleware(AgentMiddleware):
async def process(
self,
context: AgentContext,
next: Callable[[AgentContext], Awaitable[None]],
call_next: Callable[[AgentContext], Awaitable[None]],
) -> None:
# Check if the user message contains any blocked words
last_message = context.messages[-1] if context.messages else None
@@ -75,7 +75,7 @@ class PreTerminationMiddleware(AgentMiddleware):
# Set terminate flag to prevent further processing
raise MiddlewareTermination
await next(context)
await call_next(context)
class PostTerminationMiddleware(AgentMiddleware):
@@ -88,7 +88,7 @@ class PostTerminationMiddleware(AgentMiddleware):
async def process(
self,
context: AgentContext,
next: Callable[[AgentContext], Awaitable[None]],
call_next: Callable[[AgentContext], Awaitable[None]],
) -> None:
print(f"[PostTerminationMiddleware] Processing request (response count: {self.response_count})")
@@ -101,7 +101,7 @@ class PostTerminationMiddleware(AgentMiddleware):
raise MiddlewareTermination
# Allow the agent to process normally
await next(context)
await call_next(context)
# Increment response count after processing
self.response_count += 1
python/samples/getting_started/middleware/override_result_with_middleware.py
+6 -6
View File
@@ -48,11 +48,11 @@ def get_weather(
return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}°C."
async def weather_override_middleware(context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def weather_override_middleware(context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
"""Chat middleware that overrides weather results for both streaming and non-streaming cases."""
# Let the original agent execution complete first
await next(context)
await call_next(context)
# Check if there's a result to override (agent called weather function)
if context.result is not None:
@@ -83,9 +83,9 @@ async def weather_override_middleware(context: ChatContext, next: Callable[[Chat
context.result = ChatResponse(messages=[ChatMessage(role="assistant", text=custom_message)])
async def validate_weather_middleware(context: ChatContext, next: Callable[[ChatContext], Awaitable[None]]) -> None:
async def validate_weather_middleware(context: ChatContext, call_next: Callable[[ChatContext], Awaitable[None]]) -> None:
"""Chat middleware that simulates result validation for both streaming and non-streaming cases."""
await next(context)
await call_next(context)
validation_note = "Validation: weather data verified."
@@ -103,9 +103,9 @@ async def validate_weather_middleware(context: ChatContext, next: Callable[[Chat
context.result.messages.append(ChatMessage(role="assistant", text=validation_note))
async def agent_cleanup_middleware(context: AgentContext, next: Callable[[AgentContext], Awaitable[None]]) -> None:
async def agent_cleanup_middleware(context: AgentContext, call_next: Callable[[AgentContext], Awaitable[None]]) -> None:
"""Agent middleware that validates chat middleware effects and cleans the result."""
await next(context)
await call_next(context)
if context.result is None:
return
python/samples/getting_started/middleware/runtime_context_delegation.py
+8 -8
View File
@@ -54,7 +54,7 @@ class SessionContextContainer:
async def inject_context_middleware(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
"""MiddlewareTypes that extracts runtime context from kwargs and stores in container.
@@ -74,7 +74,7 @@ class SessionContextContainer:
print(f" - Session Metadata Keys: {list(self.session_metadata.keys())}")
# Continue to tool execution
await next(context)
await call_next(context)
# Create a container instance that will be shared via closure
@@ -278,19 +278,19 @@ async def pattern_2_hierarchical_with_kwargs_propagation() -> None:
@function_middleware
async def email_kwargs_tracker(
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
context: FunctionInvocationContext, call_next: Callable[[FunctionInvocationContext], Awaitable[None]]
) -> None:
email_agent_kwargs.update(context.kwargs)
print(f"[EmailAgent] Received runtime context: {list(context.kwargs.keys())}")
await next(context)
await call_next(context)
@function_middleware
async def sms_kwargs_tracker(
context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
context: FunctionInvocationContext, call_next: Callable[[FunctionInvocationContext], Awaitable[None]]
) -> None:
sms_agent_kwargs.update(context.kwargs)
print(f"[SMSAgent] Received runtime context: {list(context.kwargs.keys())}")
await next(context)
await call_next(context)
client = OpenAIChatClient(model_id="gpt-4o-mini")
@@ -359,7 +359,7 @@ class AuthContextMiddleware:
self.validated_tokens: list[str] = []
async def validate_and_track(
self, context: FunctionInvocationContext, next: Callable[[FunctionInvocationContext], Awaitable[None]]
self, context: FunctionInvocationContext, call_next: Callable[[FunctionInvocationContext], Awaitable[None]]
) -> None:
"""Validate API token and track usage."""
api_token = context.kwargs.get("api_token")
@@ -375,7 +375,7 @@ class AuthContextMiddleware:
else:
print("[AuthMiddleware] No API token provided")
await next(context)
await call_next(context)
@tool(approval_mode="never_require")
python/samples/getting_started/middleware/shared_state_middleware.py
+4 -4
View File
@@ -57,7 +57,7 @@ class MiddlewareContainer:
async def call_counter_middleware(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
"""First middleware: increments call count in shared state."""
# Increment the shared call count
@@ -66,18 +66,18 @@ class MiddlewareContainer:
print(f"[CallCounter] This is function call #{self.call_count}")
# Call the next middleware/function
await next(context)
await call_next(context)
async def result_enhancer_middleware(
self,
context: FunctionInvocationContext,
next: Callable[[FunctionInvocationContext], Awaitable[None]],
call_next: Callable[[FunctionInvocationContext], Awaitable[None]],
) -> None:
"""Second middleware: uses shared call count to enhance function results."""
print(f"[ResultEnhancer] Current total calls so far: {self.call_count}")
# Call the next middleware/function
await next(context)
await call_next(context)
# After function execution, enhance the result using shared state
if context.result:
python/samples/getting_started/middleware/thread_behavior_middleware.py
+7 -7
View File
@@ -21,14 +21,14 @@ The example shows:
- How AgentContext.thread property behaves across multiple runs
- How middleware can access conversation history through the thread
- The timing of when thread messages are populated (before vs after next() call)
- The timing of when thread messages are populated (before vs after call_next() call)
- How to track thread state changes across runs
Key behaviors demonstrated:
1. First run: context.messages is populated, context.thread is initially empty (before next())
2. After next(): thread contains input message + response from agent
1. First run: context.messages is populated, context.thread is initially empty (before call_next())
2. After call_next(): thread contains input message + response from agent
3. Second run: context.messages contains only current input, thread contains previous history
4. After next(): thread contains full conversation history (all previous + current messages)
4. After call_next(): thread contains full conversation history (all previous + current messages)
"""
@@ -46,7 +46,7 @@ def get_weather(
async def thread_tracking_middleware(
context: AgentContext,
next: Callable[[AgentContext], Awaitable[None]],
call_next: Callable[[AgentContext], Awaitable[None]],
) -> None:
"""MiddlewareTypes that tracks and logs thread behavior across runs."""
thread_messages = []
@@ -56,8 +56,8 @@ async def thread_tracking_middleware(
print(f"[MiddlewareTypes pre-execution] Current input messages: {len(context.messages)}")
print(f"[MiddlewareTypes pre-execution] Thread history messages: {len(thread_messages)}")
# Call next to execute the agent
await next(context)
# Call call_next to execute the agent
await call_next(context)
# Check thread state after agent execution
updated_thread_messages = []