diff --git a/python/packages/main/agent_framework/__init__.py b/python/packages/main/agent_framework/__init__.py index ca82fa412d..b7963138cc 100644 --- a/python/packages/main/agent_framework/__init__.py +++ b/python/packages/main/agent_framework/__init__.py @@ -13,6 +13,7 @@ from ._clients import * # noqa: F403 from ._logging import * # noqa: F403 from ._mcp import * # noqa: F403 from ._memory import * # noqa: F403 +from ._middleware import * # noqa: F403 from ._threads import * # noqa: F403 from ._tools import * # noqa: F403 from ._types import * # noqa: F403 diff --git a/python/packages/main/agent_framework/_agents.py b/python/packages/main/agent_framework/_agents.py index 0981112521..c1183a75b4 100644 --- a/python/packages/main/agent_framework/_agents.py +++ b/python/packages/main/agent_framework/_agents.py @@ -1,5 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. +import inspect import sys from collections.abc import AsyncIterable, Callable, MutableMapping, Sequence from contextlib import AbstractAsyncContextManager, AsyncExitStack @@ -12,6 +13,17 @@ 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 ._pydantic import AFBaseModel from ._threads import AgentThread, ChatMessageStore, deserialize_thread_state, thread_on_new_messages from ._tools import FUNCTION_INVOKING_CHAT_CLIENT_MARKER, ToolProtocol @@ -138,12 +150,16 @@ 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. """ 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) async def _notify_thread_of_new_messages( self, thread: AgentThread, new_messages: ChatMessage | Sequence[ChatMessage] @@ -170,6 +186,223 @@ 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) + + # Normalize messages to ChatMessage list + if messages is None: + normalized_messages: list[ChatMessage] = [] + elif isinstance(messages, str): + normalized_messages = [ChatMessage(role=Role.USER, text=messages)] + elif isinstance(messages, ChatMessage): + normalized_messages = [messages] + elif isinstance(messages, list): + normalized_messages = [] + for msg in messages: + if isinstance(msg, str): + normalized_messages.append(ChatMessage(role=Role.USER, text=msg)) + elif isinstance(msg, ChatMessage): + normalized_messages.append(msg) + + # 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) + + return await self._agent_middleware_pipeline.execute( + self, # type: ignore[arg-type] + normalized_messages, + context, + _execute_handler, + ) + + # 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) + + # Normalize messages to ChatMessage list + if messages is None: + normalized_messages: list[ChatMessage] = [] + elif isinstance(messages, str): + normalized_messages = [ChatMessage(role=Role.USER, text=messages)] + elif isinstance(messages, ChatMessage): + normalized_messages = [messages] + elif isinstance(messages, list): + normalized_messages = [] + for msg in messages: + if isinstance(msg, str): + normalized_messages.append(ChatMessage(role=Role.USER, text=msg)) + elif isinstance(msg, ChatMessage): + normalized_messages.append(msg) + + # 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) + # region ChatAgent @@ -216,6 +449,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, **kwargs: Any, ) -> None: """Create a ChatAgent. @@ -251,6 +485,8 @@ 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. kwargs: any additional keyword arguments. Unused, can be used by subclasses of this Agent. """ @@ -272,6 +508,7 @@ class ChatAgent(BaseAgent): "chat_client": chat_client, "chat_message_store_factory": chat_message_store_factory, "context_providers": aggregate_context_providers, + "middlewares": middlewares, "chat_options": ChatOptions( ai_model_id=model, frequency_penalty=frequency_penalty, @@ -304,6 +541,9 @@ 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. @@ -338,9 +578,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( + async def _run_impl( self, - messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None, + messages: list[ChatMessage], *, thread: AgentThread | None = None, frequency_penalty: float | None = None, @@ -367,16 +607,10 @@ class ChatAgent(BaseAgent): additional_properties: dict[str, Any] | None = None, **kwargs: Any, ) -> AgentRunResponse: - """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. + """Core ChatAgent execution logic called by the base run method after middleware processing. Args: - messages: The messages to process. + 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. @@ -396,7 +630,60 @@ 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( @@ -467,7 +754,83 @@ class ChatAgent(BaseAgent): additional_properties=response.additional_properties, ) - async def run_stream( + 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( self, messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None, *, @@ -494,37 +857,7 @@ class ChatAgent(BaseAgent): additional_properties: dict[str, Any] | None = None, **kwargs: Any, ) -> AsyncIterable[AgentRunResponseUpdate]: - """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. - - """ + """Internal run_stream 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( diff --git a/python/packages/main/agent_framework/_clients.py b/python/packages/main/agent_framework/_clients.py index 4d5775c96b..63284f53f7 100644 --- a/python/packages/main/agent_framework/_clients.py +++ b/python/packages/main/agent_framework/_clients.py @@ -10,6 +10,7 @@ from pydantic import BaseModel, Field from ._logging import get_logger from ._mcp import MCPTool from ._memory import AggregateContextProvider, ContextProvider +from ._middleware import MiddlewareType from ._pydantic import AFBaseModel from ._threads import ChatMessageStore from ._tools import ToolProtocol @@ -465,6 +466,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, **kwargs: Any, ) -> "ChatAgent": """Create an agent with the given name and instructions. @@ -476,6 +478,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. **kwargs: Additional keyword arguments to pass to the agent. See ChatAgent for all the available options. @@ -491,6 +494,7 @@ class BaseChatClient(AFBaseModel, ABC): tools=tools, chat_message_store_factory=chat_message_store_factory, context_providers=context_providers, + middlewares=middlewares, **kwargs, ) diff --git a/python/packages/main/agent_framework/_middleware.py b/python/packages/main/agent_framework/_middleware.py new file mode 100644 index 0000000000..28e8859c6f --- /dev/null +++ b/python/packages/main/agent_framework/_middleware.py @@ -0,0 +1,414 @@ +# Copyright (c) Microsoft. All rights reserved. + +from abc import ABC, abstractmethod +from collections.abc import AsyncIterable, Awaitable, Callable +from typing import TYPE_CHECKING, Any +from uuid import uuid4 + +if TYPE_CHECKING: + from pydantic import BaseModel + + from ._agents import AgentProtocol + from ._tools import AIFunction + from ._types import AgentRunResponse, AgentRunResponseUpdate, ChatMessage + +__all__ = [ + "AgentInvocationContext", + "AgentMiddleware", + "FunctionInvocationContext", + "FunctionMiddleware", + "MiddlewareType", +] + + +class AgentInvocationContext: + """Context object for agent middleware invocations. + + Attributes: + agent: The agent being invoked. + messages: The messages being sent to the agent. + is_streaming: Whether this is a streaming invocation. + request_id: Unique identifier for the current request. + metadata: Metadata dictionary for sharing data between agent middleware. + """ + + def __init__( + self, + agent: "AgentProtocol", + messages: list["ChatMessage"], + is_streaming: bool = False, + request_id: str | None = None, + metadata: dict[str, Any] | None = None, + ) -> None: + """Initialize agent invocation context. + + Args: + agent: The agent being invoked. + messages: The messages being sent to the agent. + is_streaming: Whether this is a streaming invocation. + request_id: Unique identifier for the request. Auto-generated if None. + metadata: Metadata dictionary. + """ + self.agent = agent + self.messages = messages + self.is_streaming = is_streaming + self.request_id = request_id or str(uuid4()) + self.metadata = metadata or {} + + +class FunctionInvocationContext: + """Context object for function middleware invocations. + + Attributes: + function: The function being invoked. + arguments: The validated arguments for the function. + request_id: Unique identifier for the current request. + metadata: Metadata dictionary for sharing data between function middleware. + """ + + def __init__( + self, + function: "AIFunction[Any, Any]", + arguments: "BaseModel", + request_id: str | None = None, + metadata: dict[str, Any] | None = None, + ) -> None: + """Initialize function invocation context. + + Args: + function: The function being invoked. + arguments: The validated arguments for the function. + request_id: Unique identifier for the request. Auto-generated if None. + metadata: Metadata dictionary. + """ + self.function = function + self.arguments = arguments + self.request_id = request_id or str(uuid4()) + self.metadata = metadata or {} + + +class AgentMiddleware(ABC): + """Abstract base class for agent middleware that can intercept agent invocations.""" + + @abstractmethod + async def process( + self, + context: AgentInvocationContext, + next: Callable[[AgentInvocationContext], Awaitable[None]], + ) -> None: + """Process an agent invocation. + + Args: + context: Agent invocation context containing agent, messages, and metadata. + Use context.is_streaming to determine if this is a streaming call. + Middleware can set context.should_skip=True and provide context.response + or context.response_stream to override the agent execution. + next: Function to call the next middleware or final agent execution. + Does not return anything - all data flows through the context. + + Note: + Middleware should not return anything. All data manipulation should happen + within the context object. Set context.should_skip=True and provide + context.response or context.response_stream to override execution. + """ + ... + + +class FunctionMiddleware(ABC): + """Abstract base class for function middleware that can intercept function invocations.""" + + @abstractmethod + async def process( + self, + context: FunctionInvocationContext, + next: Callable[[FunctionInvocationContext], Awaitable[None]], + ) -> None: + """Process a function invocation. + + Args: + context: Function invocation context containing function, arguments, and metadata. + Middleware can set context.should_skip=True and provide context.result + to override the function execution. + next: Function to call the next middleware or final function execution. + Does not return anything - all data flows through the context. + + Note: + Middleware should not return anything. All data manipulation should happen + within the context object. Set context.should_skip=True and provide + context.result to override execution. + """ + ... + + +# Pure function type definitions for convenience +AgentMiddlewareCallable = Callable[ + [AgentInvocationContext, Callable[[AgentInvocationContext], Awaitable[None]]], Awaitable[None] +] + +FunctionMiddlewareCallable = Callable[ + [FunctionInvocationContext, Callable[[FunctionInvocationContext], Awaitable[None]]], Awaitable[None] +] + +# Type alias for all middleware types +MiddlewareType = AgentMiddleware | AgentMiddlewareCallable | FunctionMiddleware | FunctionMiddlewareCallable + + +class AgentMiddlewareWrapper: + """Wrapper to convert pure functions into AgentMiddleware protocol objects.""" + + def __init__(self, func: AgentMiddlewareCallable): + self.func = func + + async def process( + self, + context: AgentInvocationContext, + next: Callable[[AgentInvocationContext], Awaitable[None]], + ) -> None: + await self.func(context, next) + + +class FunctionMiddlewareWrapper: + """Wrapper to convert pure functions into FunctionMiddleware protocol objects.""" + + def __init__(self, func: FunctionMiddlewareCallable): + self.func = func + + async def process( + self, + context: FunctionInvocationContext, + next: Callable[[FunctionInvocationContext], Awaitable[None]], + ) -> None: + await self.func(context, next) + + +class AgentMiddlewarePipeline: + """Executes agent middleware in a chain.""" + + def __init__(self, middlewares: list[AgentMiddleware | AgentMiddlewareCallable] | None = None): + """Initialize the agent middleware pipeline. + + Args: + middlewares: List of agent middleware to include in the pipeline. + """ + self._middlewares: list[AgentMiddleware] = [] + + if middlewares: + for middleware in middlewares: + self._register_middleware(middleware) + + def _register_middleware(self, middleware: AgentMiddleware | AgentMiddlewareCallable) -> None: + """Register an agent middleware item.""" + if callable(middleware): + # Check if it's already a protocol implementation + if callable(middleware) and not hasattr(middleware, "func"): + # It's a class instance implementing the protocol + self._middlewares.append(middleware) # type: ignore + else: + # It's a pure function, wrap it + self._middlewares.append(AgentMiddlewareWrapper(middleware)) # type: ignore + else: + self._middlewares.append(middleware) # type: ignore + + async def execute( + self, + agent: "AgentProtocol", + messages: list["ChatMessage"], + context: AgentInvocationContext, + final_handler: Callable[[AgentInvocationContext], Awaitable["AgentRunResponse"]], + ) -> "AgentRunResponse": + """Execute the agent middleware pipeline for non-streaming. + + Args: + agent: The agent being invoked. + messages: The messages to send to the agent. + context: The agent invocation context. + final_handler: The final handler that performs the actual agent execution. + + Returns: + The agent response after processing through all middleware. + """ + # Update context with agent and messages + context.agent = agent + context.messages = messages + context.is_streaming = False + + if not self._middlewares: + return await final_handler(context) + + # Store the final result + result_container: dict[str, AgentRunResponse | None] = {"response": None} + + def create_next_handler(index: int) -> Callable[[AgentInvocationContext], Awaitable[None]]: + if index >= len(self._middlewares): + + async def final_wrapper(c: AgentInvocationContext) -> None: + result_container["response"] = await final_handler(c) + + return final_wrapper + + middleware = self._middlewares[index] + next_handler = create_next_handler(index + 1) + + async def current_handler(c: AgentInvocationContext) -> None: + await middleware.process(c, next_handler) + + return current_handler + + first_handler = create_next_handler(0) + await first_handler(context) + + # Return the response from result container + response = result_container["response"] + if response is None: + raise RuntimeError("No response set after middleware execution") + return response + + async def execute_stream( + self, + agent: "AgentProtocol", + messages: list["ChatMessage"], + context: AgentInvocationContext, + final_handler: Callable[[AgentInvocationContext], AsyncIterable["AgentRunResponseUpdate"]], + ) -> AsyncIterable["AgentRunResponseUpdate"]: + """Execute the agent middleware pipeline for streaming. + + Args: + agent: The agent being invoked. + messages: The messages to send to the agent. + context: The agent invocation context. + final_handler: The final handler that performs the actual agent streaming execution. + + Yields: + Agent response updates after processing through all middleware. + """ + # Update context with agent and messages + context.agent = agent + context.messages = messages + context.is_streaming = True + + if not self._middlewares: + async for update in final_handler(context): + yield update + return + + # Store the final result + result_container: dict[str, AsyncIterable[AgentRunResponseUpdate] | None] = {"response_stream": None} + + def create_next_handler(index: int) -> Callable[[AgentInvocationContext], Awaitable[None]]: + if index >= len(self._middlewares): + + async def final_wrapper(c: AgentInvocationContext) -> None: # noqa: RUF029 + result_container["response_stream"] = final_handler(c) + + return final_wrapper + + middleware = self._middlewares[index] + next_handler = create_next_handler(index + 1) + + async def current_handler(c: AgentInvocationContext) -> None: + await middleware.process(c, next_handler) + + return current_handler + + first_handler = create_next_handler(0) + await first_handler(context) + + # Yield from the response stream in result container + response_stream = result_container["response_stream"] + if response_stream is None: + raise RuntimeError("No response stream set after middleware execution") + + async for update in response_stream: + yield update + + @property + def has_middlewares(self) -> bool: + """Check if there are any middlewares registered.""" + return bool(self._middlewares) + + +class FunctionMiddlewarePipeline: + """Executes function middleware in a chain.""" + + def __init__(self, middlewares: list[FunctionMiddleware | FunctionMiddlewareCallable] | None = None): + """Initialize the function middleware pipeline. + + Args: + middlewares: List of function middleware to include in the pipeline. + """ + self._middlewares: list[FunctionMiddleware] = [] + + if middlewares: + for middleware in middlewares: + self._register_middleware(middleware) + + def _register_middleware(self, middleware: FunctionMiddleware | FunctionMiddlewareCallable) -> None: + """Register a function middleware item.""" + if callable(middleware): + # Check if it's already a protocol implementation + if callable(middleware) and not hasattr(middleware, "func"): + # It's a class instance implementing the protocol + self._middlewares.append(middleware) # type: ignore + else: + # It's a pure function, wrap it + self._middlewares.append(FunctionMiddlewareWrapper(middleware)) # type: ignore + else: + self._middlewares.append(middleware) # type: ignore + + async def execute( + self, + function: Any, + arguments: "BaseModel", + context: FunctionInvocationContext, + final_handler: Callable[[FunctionInvocationContext], Awaitable[Any]], + ) -> Any: + """Execute the function middleware pipeline. + + Args: + function: The function being invoked. + arguments: The validated arguments for the function. + context: The function invocation context. + final_handler: The final handler that performs the actual function execution. + + Returns: + The function result after processing through all middleware. + """ + # Update context with function and arguments + context.function = function + context.arguments = arguments + + if not self._middlewares: + return await final_handler(context) + + # Store the final result + result_container: dict[str, Any] = {"result": None} + + def create_next_handler(index: int) -> Callable[[FunctionInvocationContext], Awaitable[None]]: + if index >= len(self._middlewares): + + async def final_wrapper(c: FunctionInvocationContext) -> None: + result_container["result"] = await final_handler(c) + + return final_wrapper + + middleware = self._middlewares[index] + next_handler = create_next_handler(index + 1) + + async def current_handler(c: FunctionInvocationContext) -> None: + await middleware.process(c, next_handler) + + return current_handler + + first_handler = create_next_handler(0) + await first_handler(context) + + # Return the result from result container + result = result_container["result"] + if result is None: + raise RuntimeError("No result set after middleware execution") + return result + + @property + def has_middlewares(self) -> bool: + """Check if there are any middlewares registered.""" + return bool(self._middlewares) diff --git a/python/packages/main/agent_framework/_tools.py b/python/packages/main/agent_framework/_tools.py index 9825af985a..6bff5f8cb5 100644 --- a/python/packages/main/agent_framework/_tools.py +++ b/python/packages/main/agent_framework/_tools.py @@ -571,6 +571,7 @@ async def _auto_invoke_function( tool_map: dict[str, AIFunction[BaseModel, Any]], sequence_index: int | None = None, request_index: int | None = None, + middleware_pipeline: Any = None, # Optional MiddlewarePipeline ) -> "Contents": """Invoke a function call requested by the agent, applying filters that are defined in the agent.""" from ._types import FunctionResultContent @@ -585,14 +586,43 @@ async def _auto_invoke_function( merged_args: dict[str, Any] = (custom_args or {}) | parsed_args args = tool.input_model.model_validate(merged_args) exception = None - try: - function_result = await tool.invoke( + + # Execute through middleware pipeline if available + if middleware_pipeline and hasattr(middleware_pipeline, "has_middlewares") and middleware_pipeline.has_middlewares: + from ._middleware import FunctionInvocationContext + + middleware_context = FunctionInvocationContext( + function=tool, arguments=args, - tool_call_id=function_call_content.call_id, - ) # type: ignore[arg-type] - except Exception as ex: - exception = ex - function_result = None + ) + + async def final_function_handler(context_obj: Any) -> Any: + return await tool.invoke( + arguments=context_obj.arguments, + tool_call_id=function_call_content.call_id, + ) + + try: + function_result = await middleware_pipeline.execute( + function=tool, + arguments=args, + context=middleware_context, + final_handler=final_function_handler, + ) + except Exception as ex: + exception = ex + function_result = None + else: + # No middleware - execute directly + try: + function_result = await tool.invoke( + arguments=args, + tool_call_id=function_call_content.call_id, + ) # type: ignore[arg-type] + except Exception as ex: + exception = ex + function_result = None + return FunctionResultContent( call_id=function_call_content.call_id, exception=exception, @@ -626,6 +656,7 @@ async def execute_function_calls( | Callable[..., Any] \ | MutableMapping[str, Any] \ | list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]", + middleware_pipeline: Any = None, # Optional MiddlewarePipeline to avoid circular imports ) -> list["Contents"]: tool_map = _get_tool_map(tools) # Run all function calls concurrently @@ -636,6 +667,7 @@ async def execute_function_calls( tool_map=tool_map, sequence_index=seq_idx, request_index=attempt_idx, + middleware_pipeline=middleware_pipeline, ) for seq_idx, function_call in enumerate(function_calls) ]) @@ -701,11 +733,14 @@ def _handle_function_calls_response( if not tools and (chat_options := kwargs.get("chat_options")) and isinstance(chat_options, ChatOptions): tools = chat_options.tools if function_calls and tools: + # Extract function middleware pipeline from kwargs if available + middleware_pipeline = kwargs.get("_function_middleware_pipeline") function_results = await execute_function_calls( custom_args=kwargs, attempt_idx=attempt_idx, function_calls=function_calls, tools=tools, # type: ignore + middleware_pipeline=middleware_pipeline, ) # add a single ChatMessage to the response with the results result_message = ChatMessage(role="tool", contents=function_results) # type: ignore[call-overload] @@ -810,11 +845,14 @@ def _handle_function_calls_streaming_response( tools = chat_options.tools if function_calls and tools: + # Extract function middleware pipeline from kwargs if available + middleware_pipeline = kwargs.get("_function_middleware_pipeline") function_results = await execute_function_calls( custom_args=kwargs, attempt_idx=attempt_idx, function_calls=function_calls, tools=tools, # type: ignore[reportArgumentType] + middleware_pipeline=middleware_pipeline, ) function_result_msg = ChatMessage(role="tool", contents=function_results) yield ChatResponseUpdate(contents=function_results, role="tool") diff --git a/python/packages/workflow/agent_framework_workflow/_agent.py b/python/packages/workflow/agent_framework_workflow/_agent.py index ffd9877679..e2c126969d 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, response_id): + async for update in self._run_stream_impl(input_messages, thread=thread, response_id=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, response_id): + async for update in self._run_stream_impl(input_messages, thread=thread, response_id=response_id): response_updates.append(update) yield update @@ -154,18 +154,25 @@ class WorkflowAgent(BaseAgent): async def _run_stream_impl( self, - input_messages: list[ChatMessage], - response_id: str, + messages: list[ChatMessage], + *, + thread: AgentThread | None = None, + **kwargs: Any, ) -> AsyncIterable[AgentRunResponseUpdate]: """Internal implementation of streaming execution. Args: - input_messages: Normalized input messages to process. - response_id: The unique response ID for this workflow execution. + messages: Normalized input messages to process. + thread: The conversation thread (unused in workflow agent). + kwargs: Additional keyword arguments, including response_id. 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 @@ -174,7 +181,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(input_messages) + function_responses = self._extract_function_responses(messages) # Pop pending requests if fulfilled. for request_id in list(self.pending_requests.keys()): @@ -188,7 +195,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(input_messages) + event_stream = self.workflow.run_stream(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 new file mode 100644 index 0000000000..73133b1307 --- /dev/null +++ b/python/samples/getting_started/middleware/class_based_middleware.py @@ -0,0 +1,97 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +import time +from collections.abc import Awaitable, Callable +from random import randint +from typing import Annotated + +from agent_framework import AgentInvocationContext, AgentMiddleware, FunctionInvocationContext, FunctionMiddleware +from agent_framework.foundry import FoundryChatClient +from azure.identity.aio import AzureCliCredential +from pydantic import Field + + +def get_weather( + location: Annotated[str, Field(description="The location to get the weather for.")], +) -> str: + """Get the weather for a given location.""" + conditions = ["sunny", "cloudy", "rainy", "stormy"] + return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}°C." + + +class SecurityAgentMiddleware(AgentMiddleware): + """Agent middleware that checks for security violations.""" + + async def process( + self, + context: AgentInvocationContext, + next: Callable[[AgentInvocationContext], Awaitable[None]], + ) -> None: + # Check for potential security violations in the query + # Look at the last user message + last_message = context.messages[-1] if context.messages else None + if last_message and last_message.text: + 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 + return + + print("[SecurityAgentMiddleware] Security check passed.") + await next(context) + + +class LoggingFunctionMiddleware(FunctionMiddleware): + """Function middleware that logs function calls.""" + + async def process( + self, + context: FunctionInvocationContext, + 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) + + end_time = time.time() + duration = end_time - start_time + + print(f"[LoggingFunctionMiddleware] Function {function_name} completed in {duration:.3f}s.") + + +async def main() -> None: + """Example demonstrating class-based middleware.""" + print("=== Class-based Middleware Example ===") + + # For authentication, run `az login` command in terminal or replace AzureCliCredential with preferred + # authentication option. + async with ( + AzureCliCredential() as credential, + FoundryChatClient(async_credential=credential).create_agent( + name="WeatherAgent", + instructions="You are a helpful weather assistant.", + tools=get_weather, + middlewares=[SecurityAgentMiddleware(), LoggingFunctionMiddleware()], + ) as agent, + ): + # Test with normal query + print("\n--- Normal Query ---") + query = "What's the weather like in Seattle?" + print(f"User: {query}") + result = await agent.run(query) + print(f"Agent: {result}\n") + + # Test with security-related query + print("--- Security Test ---") + query = "What's the password for the weather service?" + print(f"User: {query}") + result = await agent.run(query) + print(f"Agent: {result}") + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/python/samples/getting_started/middleware/function_based_middleware.py b/python/samples/getting_started/middleware/function_based_middleware.py new file mode 100644 index 0000000000..80828b3392 --- /dev/null +++ b/python/samples/getting_started/middleware/function_based_middleware.py @@ -0,0 +1,94 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +import time +from collections.abc import Callable +from random import randint +from typing import Annotated, Any + +from agent_framework import ( + AgentInvocationContext, + FunctionInvocationContext, +) +from agent_framework.foundry import FoundryChatClient +from azure.identity.aio import AzureCliCredential +from pydantic import Field + + +def get_weather( + location: Annotated[str, Field(description="The location to get the weather for.")], +) -> str: + """Get the weather for a given location.""" + conditions = ["sunny", "cloudy", "rainy", "stormy"] + return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}°C." + + +async def security_agent_middleware( + context: AgentInvocationContext, + next: Callable[[AgentInvocationContext], Any], +) -> None: + """Agent middleware that checks for security violations.""" + # Check for potential security violations in the query + # For this example, we'll check the last user message + last_message = context.messages[-1] if context.messages else None + if last_message and last_message.text: + query = last_message.text + if "password" in query.lower() or "secret" in query.lower(): + print("[SecurityAgentMiddleware] Security Warning: Detected potential sensitive information.") + # Simply don't call next() to prevent execution + return + + print("[SecurityAgentMiddleware] Security check passed.") + await next(context) + + +async def logging_function_middleware( + context: FunctionInvocationContext, + next: Callable[[FunctionInvocationContext], Any], +) -> None: + """Function middleware that logs function calls.""" + function_name = context.function.name + print(f"[LoggingFunctionMiddleware] About to call function: {function_name}.") + + start_time = time.time() + + await next(context) + + end_time = time.time() + duration = end_time - start_time + + print(f"[LoggingFunctionMiddleware] Function {function_name} completed in {duration:.3f}s.") + + +async def main() -> None: + """Example demonstrating function-based middleware.""" + print("=== Function-based Middleware Example ===") + + # For authentication, run `az login` command in terminal or replace AzureCliCredential with preferred + # authentication option. + async with ( + AzureCliCredential() as credential, + FoundryChatClient(async_credential=credential).create_agent( + name="WeatherAgent", + instructions="You are a helpful weather assistant.", + tools=get_weather, + middlewares=[security_agent_middleware, logging_function_middleware], + ) as agent, + ): + # Test with normal query + print("\n--- Normal Query ---") + query = "What's the weather like in Tokyo?" + print(f"User: {query}") + result = await agent.run(query) + print(f"Agent: {result}\n") + + # Test with security violation + print("--- Security Test ---") + query = "What's the secret weather password?" + print(f"User: {query}") + result = await agent.run(query) + print(f"Agent: {result}") + + +if __name__ == "__main__": + asyncio.run(main())