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

Moved middleware functionality to decorator

Browse Source
This commit is contained in:
Dmytro Struk
2025-09-16 09:48:35 -07:00
Unverified
parent a7b460d724
commit 63ece868f3
6 changed files with 268 additions and 388 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/main/agent_framework/_agents.py
+49 -359
View File
@@ -1,6 +1,5 @@
# Copyright (c) Microsoft. All rights reserved.
import inspect
import sys
from collections.abc import AsyncIterable, Callable, MutableMapping, Sequence
from contextlib import AbstractAsyncContextManager, AsyncExitStack
@@ -13,17 +12,7 @@ from ._clients import BaseChatClient, ChatClientProtocol
from ._logging import get_logger
from ._mcp import MCPTool
from ._memory import AggregateContextProvider, Context, ContextProvider
from ._middleware import (
AgentInvocationContext,
AgentMiddleware,
AgentMiddlewareCallable,
AgentMiddlewarePipeline,
FunctionInvocationContext,
FunctionMiddleware,
FunctionMiddlewareCallable,
FunctionMiddlewarePipeline,
MiddlewareType,
)
from ._middleware import MiddlewareType, use_agent_middleware
from ._pydantic import AFBaseModel
from ._threads import AgentThread, ChatMessageStore, deserialize_thread_state, thread_on_new_messages
from ._tools import FUNCTION_INVOKING_CHAT_CLIENT_MARKER, ToolProtocol
@@ -150,16 +139,14 @@ class BaseAgent(AFBaseModel):
description: The description of the agent.
display_name: The display name of the agent, which is either the name or id.
context_providers: The collection of multiple context providers to include during agent invocation.
middlewares: List of middleware to intercept agent and function invocations.
middleware: List of middleware to intercept agent and function invocations.
"""
id: str = Field(default_factory=lambda: str(uuid4()))
name: str | None = None
description: str | None = None
context_providers: AggregateContextProvider | None = None
middlewares: MiddlewareType | list[MiddlewareType] | None = None
_agent_middleware_pipeline: AgentMiddlewarePipeline = PrivateAttr(default_factory=AgentMiddlewarePipeline)
_function_middleware_pipeline: FunctionMiddlewarePipeline = PrivateAttr(default_factory=FunctionMiddlewarePipeline)
middleware: MiddlewareType | list[MiddlewareType] | None = None
async def _notify_thread_of_new_messages(
self, thread: AgentThread, new_messages: ChatMessage | Sequence[ChatMessage]
@@ -186,199 +173,6 @@ class BaseAgent(AFBaseModel):
await deserialize_thread_state(thread, serialized_thread, **kwargs)
return thread
def _initialize_middleware_pipelines(self, middlewares: MiddlewareType | list[MiddlewareType] | None) -> None:
"""Initialize agent and function middleware pipelines from the provided middlewares.
This method classifies middleware by type and creates appropriate pipelines.
Args:
middlewares: List of middleware to classify and initialize into pipelines.
"""
# Separate middlewares by type
agent_middlewares: list[AgentMiddleware | AgentMiddlewareCallable] = []
function_middlewares: list[FunctionMiddleware | FunctionMiddlewareCallable] = []
if middlewares:
middlewares_list = middlewares if isinstance(middlewares, list) else [middlewares]
for middleware in middlewares_list:
# Classify middleware by type checking
if isinstance(middleware, AgentMiddleware):
agent_middlewares.append(middleware)
elif isinstance(middleware, FunctionMiddleware):
function_middlewares.append(middleware)
elif callable(middleware) and inspect.iscoroutinefunction(middleware):
# It's a function, classify by signature inspection
sig = inspect.signature(middleware)
params = list(sig.parameters.values())
# Look at the first parameter's type annotation
if len(params) >= 1:
first_param_annotation = params[0].annotation
if first_param_annotation != inspect.Parameter.empty:
# Check if it's an AgentInvocationContext or FunctionInvocationContext
if first_param_annotation is AgentInvocationContext:
agent_middlewares.append(middleware) # type: ignore
elif first_param_annotation is FunctionInvocationContext:
function_middlewares.append(middleware) # type: ignore
self._agent_middleware_pipeline = AgentMiddlewarePipeline(agent_middlewares)
self._function_middleware_pipeline = FunctionMiddlewarePipeline(function_middlewares)
async def _run_impl(
self,
messages: list[ChatMessage],
*,
thread: AgentThread | None = None,
**kwargs: Any,
) -> AgentRunResponse:
"""Core agent execution logic that subclasses must implement.
This method contains the actual agent logic and is called by the public run method
after middleware initialization and message normalization.
Args:
messages: Normalized list of ChatMessage objects.
thread: The conversation thread associated with the messages.
kwargs: Additional keyword arguments.
Returns:
An agent response.
Raises:
NotImplementedError: This method must be implemented by subclasses.
"""
raise NotImplementedError("Subclasses must implement _run_impl method")
def _run_stream_impl(
self,
messages: list[ChatMessage],
*,
thread: AgentThread | None = None,
**kwargs: Any,
) -> AsyncIterable[AgentRunResponseUpdate]:
"""Core agent streaming execution logic that subclasses must implement.
This method contains the actual agent streaming logic and is called by the public
run_stream method after middleware initialization and message normalization.
Args:
messages: Normalized list of ChatMessage objects.
thread: The conversation thread associated with the messages.
kwargs: Additional keyword arguments.
Yields:
Agent response updates.
Raises:
NotImplementedError: This method must be implemented by subclasses.
"""
raise NotImplementedError("Subclasses must implement _run_stream_impl method")
async def run(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
) -> AgentRunResponse:
"""Get a response from the agent with automatic middleware execution.
This method handles middleware initialization, message normalization, and execution
automatically. Custom agents should implement _run_impl instead of overriding this method.
Args:
messages: The message(s) to send to the agent.
thread: The conversation thread associated with the message(s).
kwargs: Additional keyword arguments.
Returns:
An agent response item.
"""
# Initialize middleware pipelines if not already done
if self.middlewares and not (
self._agent_middleware_pipeline.has_middlewares or self._function_middleware_pipeline.has_middlewares
):
self._initialize_middleware_pipelines(self.middlewares)
normalized_messages = self._normalize_messages(messages)
# Execute with middleware if available
if self._agent_middleware_pipeline.has_middlewares:
from ._middleware import AgentInvocationContext
context = AgentInvocationContext(
agent=self, # type: ignore[arg-type]
messages=normalized_messages,
is_streaming=False,
)
async def _execute_handler(ctx: AgentInvocationContext) -> AgentRunResponse:
return await self._run_impl(ctx.messages, thread=thread, **kwargs)
response = await self._agent_middleware_pipeline.execute(
self, # type: ignore[arg-type]
normalized_messages,
context,
_execute_handler,
)
return response if response else AgentRunResponse()
# No middleware, execute directly
return await self._run_impl(normalized_messages, thread=thread, **kwargs)
def run_stream(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
**kwargs: Any,
) -> AsyncIterable[AgentRunResponseUpdate]:
"""Run the agent as a stream with automatic middleware execution.
This method handles middleware initialization, message normalization, and execution
automatically. Custom agents should implement _run_stream_impl instead of overriding this method.
Args:
messages: The message(s) to send to the agent.
thread: The conversation thread associated with the message(s).
kwargs: Additional keyword arguments.
Yields:
Agent response items.
"""
# Initialize middleware pipelines if not already done
if self.middlewares and not (
self._agent_middleware_pipeline.has_middlewares or self._function_middleware_pipeline.has_middlewares
):
self._initialize_middleware_pipelines(self.middlewares)
normalized_messages = self._normalize_messages(messages)
# Execute with middleware if available
if self._agent_middleware_pipeline.has_middlewares:
from ._middleware import AgentInvocationContext
context = AgentInvocationContext(
agent=self, # type: ignore[arg-type]
messages=normalized_messages,
is_streaming=True,
)
async def _execute_stream_handler(ctx: AgentInvocationContext) -> AsyncIterable[AgentRunResponseUpdate]:
async for update in self._run_stream_impl(ctx.messages, thread=thread, **kwargs):
yield update
return self._agent_middleware_pipeline.execute_stream(
self, # type: ignore[arg-type]
normalized_messages,
context,
_execute_stream_handler,
)
# No middleware, execute directly
return self._run_stream_impl(normalized_messages, thread=thread, **kwargs)
def _normalize_messages(
self,
messages: str | ChatMessage | Sequence[str] | Sequence[ChatMessage] | None = None,
@@ -398,6 +192,7 @@ class BaseAgent(AFBaseModel):
# region ChatAgent
@use_agent_middleware
@use_agent_telemetry
class ChatAgent(BaseAgent):
"""A Chat Client Agent."""
@@ -440,7 +235,7 @@ class ChatAgent(BaseAgent):
additional_properties: dict[str, Any] | None = None,
chat_message_store_factory: Callable[[], ChatMessageStore] | None = None,
context_providers: ContextProvider | list[ContextProvider] | AggregateContextProvider | None = None,
middlewares: MiddlewareType | list[MiddlewareType] | None = None,
middleware: MiddlewareType | list[MiddlewareType] | None = None,
**kwargs: Any,
) -> None:
"""Create a ChatAgent.
@@ -476,8 +271,7 @@ class ChatAgent(BaseAgent):
chat_message_store_factory: factory function to create an instance of ChatMessageStore. If not provided,
the default in-memory store will be used.
context_providers: The collection of multiple context providers to include during agent invocation.
middlewares: List of middleware to intercept agent and function invocations.
Can include class instances implementing middleware protocols or pure functions.
middleware: List of middleware to intercept agent and function invocations.
kwargs: any additional keyword arguments.
Unused, can be used by subclasses of this Agent.
"""
@@ -499,7 +293,7 @@ class ChatAgent(BaseAgent):
"chat_client": chat_client,
"chat_message_store_factory": chat_message_store_factory,
"context_providers": aggregate_context_providers,
"middlewares": middlewares,
"middleware": middleware,
"chat_options": ChatOptions(
ai_model_id=model,
frequency_penalty=frequency_penalty,
@@ -532,9 +326,6 @@ class ChatAgent(BaseAgent):
self._update_agent_name()
self._local_mcp_tools = local_mcp_tools # type: ignore[assignment]
# Initialize middleware pipelines
self._initialize_middleware_pipelines(middlewares)
async def __aenter__(self) -> "Self":
"""Async context manager entry.
@@ -569,9 +360,9 @@ class ChatAgent(BaseAgent):
if hasattr(self.chat_client, "_update_agent_name") and callable(self.chat_client._update_agent_name): # type: ignore[reportAttributeAccessIssue, attr-defined]
self.chat_client._update_agent_name(self.name) # type: ignore[reportAttributeAccessIssue, attr-defined]
async def _run_impl(
async def run(
self,
messages: list[ChatMessage],
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
frequency_penalty: float | None = None,
@@ -598,10 +389,16 @@ class ChatAgent(BaseAgent):
additional_properties: dict[str, Any] | None = None,
**kwargs: Any,
) -> AgentRunResponse:
"""Core ChatAgent execution logic called by the base run method after middleware processing.
"""Run the agent with the given messages and options.
Remarks:
Since you won't always call the agent.run directly, but it get's called
through orchestration, it is advised to set your default values for
all the chat client parameters in the agent constructor.
If both parameters are used, the ones passed to the run methods take precedence.
Args:
messages: Normalized list of ChatMessage objects.
messages: The messages to process.
thread: The thread to use for the agent.
frequency_penalty: the frequency penalty to use.
logit_bias: the logit bias to use.
@@ -621,60 +418,7 @@ class ChatAgent(BaseAgent):
additional_properties: additional properties to include in the request.
kwargs: Additional keyword arguments for the agent.
will only be passed to functions that are called.
Returns:
An agent response.
"""
return await self._run_internal(
messages=messages,
thread=thread,
frequency_penalty=frequency_penalty,
logit_bias=logit_bias,
max_tokens=max_tokens,
metadata=metadata,
model=model,
presence_penalty=presence_penalty,
response_format=response_format,
seed=seed,
stop=stop,
store=store,
temperature=temperature,
tool_choice=tool_choice,
tools=tools, # type: ignore[arg-type]
top_p=top_p,
user=user,
additional_properties=additional_properties,
**kwargs,
)
async def _run_internal(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
frequency_penalty: float | None = None,
logit_bias: dict[str | int, float] | None = None,
max_tokens: int | None = None,
metadata: dict[str, Any] | None = None,
model: str | None = None,
presence_penalty: float | None = None,
response_format: type[BaseModel] | None = None,
seed: int | None = None,
stop: str | Sequence[str] | None = None,
store: bool | None = None,
temperature: float | None = None,
tool_choice: ChatToolMode | Literal["auto", "required", "none"] | dict[str, Any] | None = None,
tools: ToolProtocol
| Callable[..., Any]
| dict[str, Any]
| list[ToolProtocol | Callable[..., Any] | dict[str, Any]]
| None = None,
top_p: float | None = None,
user: str | None = None,
additional_properties: dict[str, Any] | None = None,
**kwargs: Any,
) -> AgentRunResponse:
"""Internal run method that performs the actual agent execution without middleware."""
input_messages = self._normalize_messages(messages)
context = await self.context_providers.model_invoking(input_messages) if self.context_providers else None
thread, thread_messages = await self._prepare_thread_and_messages(
@@ -695,10 +439,6 @@ class ChatAgent(BaseAgent):
for mcp_server in self._local_mcp_tools:
final_tools.extend(mcp_server.functions)
# Add function middleware pipeline to kwargs if available
if self._function_middleware_pipeline.has_middlewares:
kwargs["_function_middleware_pipeline"] = self._function_middleware_pipeline
response = await self.chat_client.get_response(
messages=thread_messages,
chat_options=self.chat_options
@@ -749,83 +489,7 @@ class ChatAgent(BaseAgent):
additional_properties=response.additional_properties,
)
def _run_stream_impl(
self,
messages: list[ChatMessage],
*,
thread: AgentThread | None = None,
frequency_penalty: float | None = None,
logit_bias: dict[str | int, float] | None = None,
max_tokens: int | None = None,
metadata: dict[str, Any] | None = None,
model: str | None = None,
presence_penalty: float | None = None,
response_format: type[BaseModel] | None = None,
seed: int | None = None,
stop: str | Sequence[str] | None = None,
store: bool | None = None,
temperature: float | None = None,
tool_choice: ChatToolMode | Literal["auto", "required", "none"] | dict[str, Any] | None = None,
tools: ToolProtocol
| Callable[..., Any]
| MutableMapping[str, Any]
| list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
| None = None,
top_p: float | None = None,
user: str | None = None,
additional_properties: dict[str, Any] | None = None,
**kwargs: Any,
) -> AsyncIterable[AgentRunResponseUpdate]:
"""Core ChatAgent streaming execution logic called by the base run_stream method after middleware processing.
Args:
messages: Normalized list of ChatMessage objects.
thread: The thread to use for the agent.
frequency_penalty: the frequency penalty to use.
logit_bias: the logit bias to use.
max_tokens: The maximum number of tokens to generate.
metadata: additional metadata to include in the request.
model: The model to use for the agent.
presence_penalty: the presence penalty to use.
response_format: the format of the response.
seed: the random seed to use.
stop: the stop sequence(s) for the request.
store: whether to store the response.
temperature: the sampling temperature to use.
tool_choice: the tool choice for the request.
tools: the tools to use for the request.
top_p: the nucleus sampling probability to use.
user: the user to associate with the request.
additional_properties: additional properties to include in the request.
kwargs: any additional keyword arguments.
will only be passed to functions that are called.
Yields:
Agent response updates.
"""
return self._run_stream_internal(
messages=messages,
thread=thread,
frequency_penalty=frequency_penalty,
logit_bias=logit_bias,
max_tokens=max_tokens,
metadata=metadata,
model=model,
presence_penalty=presence_penalty,
response_format=response_format,
seed=seed,
stop=stop,
store=store,
temperature=temperature,
tool_choice=tool_choice,
tools=tools, # type: ignore[arg-type]
top_p=top_p,
user=user,
additional_properties=additional_properties,
**kwargs,
)
async def _run_stream_internal(
async def run_stream(
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
@@ -852,7 +516,37 @@ class ChatAgent(BaseAgent):
additional_properties: dict[str, Any] | None = None,
**kwargs: Any,
) -> AsyncIterable[AgentRunResponseUpdate]:
"""Internal run_stream method that performs the actual agent execution without middleware."""
"""Stream the agent with the given messages and options.
Remarks:
Since you won't always call the agent.run_stream directly, but it get's called
through orchestration, it is advised to set your default values for
all the chat client parameters in the agent constructor.
If both parameters are used, the ones passed to the run methods take precedence.
Args:
messages: The messages to process.
thread: The thread to use for the agent.
frequency_penalty: the frequency penalty to use.
logit_bias: the logit bias to use.
max_tokens: The maximum number of tokens to generate.
metadata: additional metadata to include in the request.
model: The model to use for the agent.
presence_penalty: the presence penalty to use.
response_format: the format of the response.
seed: the random seed to use.
stop: the stop sequence(s) for the request.
store: whether to store the response.
temperature: the sampling temperature to use.
tool_choice: the tool choice for the request.
tools: the tools to use for the request.
top_p: the nucleus sampling probability to use.
user: the user to associate with the request.
additional_properties: additional properties to include in the request.
kwargs: any additional keyword arguments.
will only be passed to functions that are called.
"""
input_messages = self._normalize_messages(messages)
context = await self.context_providers.model_invoking(input_messages) if self.context_providers else None
thread, thread_messages = await self._prepare_thread_and_messages(
@@ -874,10 +568,6 @@ class ChatAgent(BaseAgent):
for mcp_server in self._local_mcp_tools:
final_tools.extend(mcp_server.functions)
# Add function middleware pipeline to kwargs if available
if self._function_middleware_pipeline.has_middlewares:
kwargs["_function_middleware_pipeline"] = self._function_middleware_pipeline
async for update in self.chat_client.get_streaming_response(
messages=thread_messages,
chat_options=self.chat_options
python/packages/main/agent_framework/_clients.py
+3 -3
View File
@@ -474,7 +474,7 @@ class BaseChatClient(AFBaseModel, ABC):
| None = None,
chat_message_store_factory: Callable[[], ChatMessageStore] | None = None,
context_providers: ContextProvider | list[ContextProvider] | AggregateContextProvider | None = None,
middlewares: MiddlewareType | list[MiddlewareType] | None = None,
middleware: MiddlewareType | list[MiddlewareType] | None = None,
**kwargs: Any,
) -> "ChatAgent":
"""Create an agent with the given name and instructions.
@@ -486,7 +486,7 @@ class BaseChatClient(AFBaseModel, ABC):
chat_message_store_factory: Factory function to create an instance of ChatMessageStore. If not provided,
the default in-memory store will be used.
context_providers: Context providers to include during agent invocation.
middlewares: List of middleware to intercept agent and function invocations.
middleware: List of middleware to intercept agent and function invocations.
**kwargs: Additional keyword arguments to pass to the agent.
See ChatAgent for all the available options.
@@ -502,7 +502,7 @@ class BaseChatClient(AFBaseModel, ABC):
tools=tools,
chat_message_store_factory=chat_message_store_factory,
context_providers=context_providers,
middlewares=middlewares,
middleware=middleware,
**kwargs,
)
python/packages/main/agent_framework/_middleware.py
+206 -9
View File
@@ -2,14 +2,17 @@
from abc import ABC, abstractmethod
from collections.abc import AsyncIterable, Awaitable, Callable
from typing import TYPE_CHECKING, Any
from typing import TYPE_CHECKING, Any, TypeVar
from ._types import AgentRunResponse, AgentRunResponseUpdate, ChatMessage
if TYPE_CHECKING:
from pydantic import BaseModel
from ._agents import AgentProtocol
from ._tools import AIFunction
from ._types import AgentRunResponse, AgentRunResponseUpdate, ChatMessage
TAgent = TypeVar("TAgent", bound="AgentProtocol")
__all__ = [
"AgentInvocationContext",
@@ -17,6 +20,7 @@ __all__ = [
"FunctionInvocationContext",
"FunctionMiddleware",
"MiddlewareType",
"use_agent_middleware",
]
@@ -33,7 +37,7 @@ class AgentInvocationContext:
def __init__(
self,
agent: "AgentProtocol",
messages: list["ChatMessage"],
messages: list[ChatMessage],
is_streaming: bool = False,
metadata: dict[str, Any] | None = None,
) -> None:
@@ -197,10 +201,10 @@ class AgentMiddlewarePipeline:
async def execute(
self,
agent: "AgentProtocol",
messages: list["ChatMessage"],
messages: list[ChatMessage],
context: AgentInvocationContext,
final_handler: Callable[[AgentInvocationContext], Awaitable["AgentRunResponse"]],
) -> "AgentRunResponse | None":
final_handler: Callable[[AgentInvocationContext], Awaitable[AgentRunResponse]],
) -> AgentRunResponse | None:
"""Execute the agent middleware pipeline for non-streaming.
Args:
@@ -248,10 +252,10 @@ class AgentMiddlewarePipeline:
async def execute_stream(
self,
agent: "AgentProtocol",
messages: list["ChatMessage"],
messages: list[ChatMessage],
context: AgentInvocationContext,
final_handler: Callable[[AgentInvocationContext], AsyncIterable["AgentRunResponseUpdate"]],
) -> AsyncIterable["AgentRunResponseUpdate"]:
final_handler: Callable[[AgentInvocationContext], AsyncIterable[AgentRunResponseUpdate]],
) -> AsyncIterable[AgentRunResponseUpdate]:
"""Execute the agent middleware pipeline for streaming.
Args:
@@ -386,3 +390,196 @@ class FunctionMiddlewarePipeline:
def has_middlewares(self) -> bool:
"""Check if there are any middlewares registered."""
return bool(self._middlewares)
# Decorator for adding middleware support to agent classes
def use_agent_middleware(agent_class: type[TAgent]) -> type[TAgent]:
"""Class decorator that adds middleware support to an agent class.
This decorator adds middleware functionality to any agent class.
It wraps the run() and run_stream() methods to provide middleware execution.
Args:
agent_class: The agent class to add middleware support to.
Returns:
The modified agent class with middleware support.
"""
import inspect
# Store original methods
original_run = agent_class.run # type: ignore[attr-defined]
original_run_stream = agent_class.run_stream # type: ignore[attr-defined]
def _initialize_middleware_pipelines(self: Any, middlewares: MiddlewareType | list[MiddlewareType] | None) -> None:
"""Initialize agent and function middleware pipelines from the provided middleware list."""
if not middlewares:
return
middleware_list: list[MiddlewareType] = middlewares if isinstance(middlewares, list) else [middlewares] # type: ignore
# Separate agent and function middleware using isinstance checks
agent_middlewares: list[AgentMiddleware | AgentMiddlewareCallable] = []
function_middlewares: list[FunctionMiddleware | FunctionMiddlewareCallable] = []
for middleware in middleware_list:
if isinstance(middleware, AgentMiddleware):
agent_middlewares.append(middleware)
elif isinstance(middleware, FunctionMiddleware):
function_middlewares.append(middleware)
elif callable(middleware): # type: ignore[arg-type]
# Check function signature to determine type
try:
sig = inspect.signature(middleware)
params = list(sig.parameters.values())
if len(params) >= 1:
first_param = params[0]
# Check if first parameter is AgentInvocationContext or FunctionInvocationContext
if (
hasattr(first_param.annotation, "__name__")
and first_param.annotation.__name__ == "AgentInvocationContext"
):
agent_middlewares.append(middleware) # type: ignore
elif (
hasattr(first_param.annotation, "__name__")
and first_param.annotation.__name__ == "FunctionInvocationContext"
):
function_middlewares.append(middleware) # type: ignore
else:
# Default to agent middleware if uncertain
agent_middlewares.append(middleware) # type: ignore
else:
agent_middlewares.append(middleware) # type: ignore
except Exception:
# If signature inspection fails, assume it's an agent middleware
agent_middlewares.append(middleware) # type: ignore
else:
# Fallback
agent_middlewares.append(middleware) # type: ignore
self._agent_middleware_pipeline = AgentMiddlewarePipeline(agent_middlewares)
self._function_middleware_pipeline = FunctionMiddlewarePipeline(function_middlewares)
async def middleware_enabled_run(
self: Any,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
thread: Any = None,
**kwargs: Any,
) -> AgentRunResponse:
"""Middleware-enabled run method."""
# Initialize middleware pipelines if not already done
if (
hasattr(self, "middleware")
and self.middleware
and not (
hasattr(self, "_agent_middleware_pipeline")
and hasattr(self, "_function_middleware_pipeline")
and (
self._agent_middleware_pipeline.has_middlewares
or self._function_middleware_pipeline.has_middlewares
)
)
):
_initialize_middleware_pipelines(self, self.middleware)
# Ensure pipelines exist even if empty
if not hasattr(self, "_agent_middleware_pipeline"):
self._agent_middleware_pipeline = AgentMiddlewarePipeline()
if not hasattr(self, "_function_middleware_pipeline"):
self._function_middleware_pipeline = FunctionMiddlewarePipeline()
# Add function middleware pipeline to kwargs if available
if self._function_middleware_pipeline.has_middlewares:
kwargs["_function_middleware_pipeline"] = self._function_middleware_pipeline
normalized_messages = self._normalize_messages(messages)
# Execute with middleware if available
if self._agent_middleware_pipeline.has_middlewares:
context = AgentInvocationContext(
agent=self, # type: ignore[arg-type]
messages=normalized_messages,
is_streaming=False,
)
async def _execute_handler(ctx: AgentInvocationContext) -> AgentRunResponse:
return await original_run(self, ctx.messages, thread=thread, **kwargs) # type: ignore
response = await self._agent_middleware_pipeline.execute(
self, # type: ignore[arg-type]
normalized_messages,
context,
_execute_handler,
)
return response if response else AgentRunResponse()
# No middleware, execute directly
return await original_run(self, normalized_messages, thread=thread, **kwargs) # type: ignore[return-value]
def middleware_enabled_run_stream(
self: Any,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
thread: Any = None,
**kwargs: Any,
) -> AsyncIterable[AgentRunResponseUpdate]:
"""Middleware-enabled run_stream method."""
# Initialize middleware pipelines if not already done
if (
hasattr(self, "middleware")
and self.middleware
and not (
hasattr(self, "_agent_middleware_pipeline")
and hasattr(self, "_function_middleware_pipeline")
and (
self._agent_middleware_pipeline.has_middlewares
or self._function_middleware_pipeline.has_middlewares
)
)
):
_initialize_middleware_pipelines(self, self.middleware)
# Ensure pipelines exist even if empty
if not hasattr(self, "_agent_middleware_pipeline"):
self._agent_middleware_pipeline = AgentMiddlewarePipeline()
if not hasattr(self, "_function_middleware_pipeline"):
self._function_middleware_pipeline = FunctionMiddlewarePipeline()
# Add function middleware pipeline to kwargs if available
if self._function_middleware_pipeline.has_middlewares:
kwargs["_function_middleware_pipeline"] = self._function_middleware_pipeline
normalized_messages = self._normalize_messages(messages)
# Execute with middleware if available
if self._agent_middleware_pipeline.has_middlewares:
context = AgentInvocationContext(
agent=self, # type: ignore[arg-type]
messages=normalized_messages,
is_streaming=True,
)
async def _execute_stream_handler(ctx: AgentInvocationContext) -> AsyncIterable[AgentRunResponseUpdate]:
async for update in original_run_stream(self, ctx.messages, thread=thread, **kwargs): # type: ignore[misc]
yield update
async def _stream_generator() -> AsyncIterable[AgentRunResponseUpdate]:
async for update in self._agent_middleware_pipeline.execute_stream(
self, # type: ignore[arg-type]
normalized_messages,
context,
_execute_stream_handler,
):
yield update
return _stream_generator()
# No middleware, execute directly
return original_run_stream(self, normalized_messages, thread=thread, **kwargs) # type: ignore
agent_class.run = middleware_enabled_run # type: ignore
agent_class.run_stream = middleware_enabled_run_stream # type: ignore
return agent_class
python/packages/workflow/agent_framework_workflow/_agent.py
+8 -15
View File
@@ -107,7 +107,7 @@ class WorkflowAgent(BaseAgent):
thread = thread or self.get_new_thread()
response_id = str(uuid.uuid4())
async for update in self._run_stream_impl(input_messages, thread=thread, response_id=response_id):
async for update in self._run_stream_impl(input_messages, response_id):
response_updates.append(update)
# Convert updates to final response.
@@ -141,7 +141,7 @@ class WorkflowAgent(BaseAgent):
response_updates: list[AgentRunResponseUpdate] = []
response_id = str(uuid.uuid4())
async for update in self._run_stream_impl(input_messages, thread=thread, response_id=response_id):
async for update in self._run_stream_impl(input_messages, response_id):
response_updates.append(update)
yield update
@@ -154,25 +154,18 @@ class WorkflowAgent(BaseAgent):
async def _run_stream_impl(
self,
messages: list[ChatMessage],
*,
thread: AgentThread | None = None,
**kwargs: Any,
input_messages: list[ChatMessage],
response_id: str,
) -> AsyncIterable[AgentRunResponseUpdate]:
"""Internal implementation of streaming execution.
Args:
messages: Normalized input messages to process.
thread: The conversation thread (unused in workflow agent).
kwargs: Additional keyword arguments, including response_id.
input_messages: Normalized input messages to process.
response_id: The unique response ID for this workflow execution.
Yields:
AgentRunResponseUpdate objects representing the workflow execution progress.
"""
# Extract response_id from kwargs
response_id = kwargs.get("response_id")
if not response_id:
raise ValueError("response_id must be provided in kwargs")
# Determine the event stream based on whether we have function responses
if bool(self.pending_requests):
# This is a continuation - use send_responses_streaming to send function responses back
@@ -181,7 +174,7 @@ class WorkflowAgent(BaseAgent):
# Extract function responses from input messages, and ensure that
# only function responses are present in messages if there is any
# pending request.
function_responses = self._extract_function_responses(messages)
function_responses = self._extract_function_responses(input_messages)
# Pop pending requests if fulfilled.
for request_id in list(self.pending_requests.keys()):
@@ -195,7 +188,7 @@ class WorkflowAgent(BaseAgent):
else:
# Execute workflow with streaming (initial run or no function responses)
# Pass the new input messages directly to the workflow
event_stream = self.workflow.run_stream(messages)
event_stream = self.workflow.run_stream(input_messages)
# Process events from the stream
async for event in event_stream:
python/samples/getting_started/middleware/class_based_middleware.py
+1 -1
View File
@@ -75,7 +75,7 @@ async def main() -> None:
name="WeatherAgent",
instructions="You are a helpful weather assistant.",
tools=get_weather,
middlewares=[SecurityAgentMiddleware(), LoggingFunctionMiddleware()],
middleware=[SecurityAgentMiddleware(), LoggingFunctionMiddleware()],
) as agent,
):
# Test with normal query
python/samples/getting_started/middleware/function_based_middleware.py
+1 -1
View File
@@ -72,7 +72,7 @@ async def main() -> None:
name="WeatherAgent",
instructions="You are a helpful weather assistant.",
tools=get_weather,
middlewares=[security_agent_middleware, logging_function_middleware],
middleware=[security_agent_middleware, logging_function_middleware],
) as agent,
):
# Test with normal query