diff --git a/python/packages/main/agent_framework/_agents.py b/python/packages/main/agent_framework/_agents.py index 9adaf8bafb..f23de00604 100644 --- a/python/packages/main/agent_framework/_agents.py +++ b/python/packages/main/agent_framework/_agents.py @@ -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 diff --git a/python/packages/main/agent_framework/_clients.py b/python/packages/main/agent_framework/_clients.py index 1595833273..7ed740a8c5 100644 --- a/python/packages/main/agent_framework/_clients.py +++ b/python/packages/main/agent_framework/_clients.py @@ -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, ) diff --git a/python/packages/main/agent_framework/_middleware.py b/python/packages/main/agent_framework/_middleware.py index 5d3b281422..9b0e1debfc 100644 --- a/python/packages/main/agent_framework/_middleware.py +++ b/python/packages/main/agent_framework/_middleware.py @@ -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 diff --git a/python/packages/workflow/agent_framework_workflow/_agent.py b/python/packages/workflow/agent_framework_workflow/_agent.py index 73a013bb41..94e1d6555e 100644 --- a/python/packages/workflow/agent_framework_workflow/_agent.py +++ b/python/packages/workflow/agent_framework_workflow/_agent.py @@ -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: diff --git a/python/samples/getting_started/middleware/class_based_middleware.py b/python/samples/getting_started/middleware/class_based_middleware.py index af36aba3c6..7a1bb324d0 100644 --- a/python/samples/getting_started/middleware/class_based_middleware.py +++ b/python/samples/getting_started/middleware/class_based_middleware.py @@ -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 diff --git a/python/samples/getting_started/middleware/function_based_middleware.py b/python/samples/getting_started/middleware/function_based_middleware.py index 2f99146277..5c31f79936 100644 --- a/python/samples/getting_started/middleware/function_based_middleware.py +++ b/python/samples/getting_started/middleware/function_based_middleware.py @@ -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