From 65633f6f798d090418fa6bbee4e606478e5b985c Mon Sep 17 00:00:00 2001 From: Dmytro Struk <13853051+dmytrostruk@users.noreply.github.com> Date: Wed, 17 Sep 2025 23:48:33 -0700 Subject: [PATCH] Addressed PR feedback --- .../packages/main/agent_framework/_agents.py | 6 +- .../packages/main/agent_framework/_clients.py | 14 +++-- .../main/agent_framework/_middleware.py | 59 ++++++------------- .../middleware/class_based_middleware.py | 14 +++++ .../exception_handling_with_middleware.py | 14 +++++ .../middleware/function_based_middleware.py | 15 +++++ .../override_result_with_middleware.py | 15 +++++ 7 files changed, 88 insertions(+), 49 deletions(-) diff --git a/python/packages/main/agent_framework/_agents.py b/python/packages/main/agent_framework/_agents.py index f23de00604..6977388588 100644 --- a/python/packages/main/agent_framework/_agents.py +++ b/python/packages/main/agent_framework/_agents.py @@ -12,7 +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 MiddlewareType, use_agent_middleware +from ._middleware import Middleware, 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 @@ -146,7 +146,7 @@ class BaseAgent(AFBaseModel): name: str | None = None description: str | None = None context_providers: AggregateContextProvider | None = None - middleware: MiddlewareType | list[MiddlewareType] | None = None + middleware: Middleware | list[Middleware] | None = None async def _notify_thread_of_new_messages( self, thread: AgentThread, new_messages: ChatMessage | Sequence[ChatMessage] @@ -235,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, - middleware: MiddlewareType | list[MiddlewareType] | None = None, + middleware: Middleware | list[Middleware] | None = None, **kwargs: Any, ) -> None: """Create a ChatAgent. diff --git a/python/packages/main/agent_framework/_clients.py b/python/packages/main/agent_framework/_clients.py index 7ed740a8c5..fede0de4d8 100644 --- a/python/packages/main/agent_framework/_clients.py +++ b/python/packages/main/agent_framework/_clients.py @@ -10,7 +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 ._middleware import Middleware from ._pydantic import AFBaseModel from ._threads import ChatMessageStore from ._tools import ToolProtocol @@ -347,7 +347,10 @@ class BaseChatClient(AFBaseModel, ABC): self._prepare_tool_choice(chat_options=chat_options) # Remove middleware pipeline from kwargs as it's only used by function invocation wrappers - filtered_kwargs = {k: v for k, v in kwargs.items() if k != "_function_middleware_pipeline"} + if "_function_middleware_pipeline" in kwargs: + filtered_kwargs = {k: v for k, v in kwargs.items() if k != "_function_middleware_pipeline"} + else: + filtered_kwargs = kwargs return await self._inner_get_response(messages=prepped_messages, chat_options=chat_options, **filtered_kwargs) @@ -431,7 +434,10 @@ class BaseChatClient(AFBaseModel, ABC): self._prepare_tool_choice(chat_options=chat_options) # Remove middleware pipeline from kwargs as it's only used by function invocation wrappers - filtered_kwargs = {k: v for k, v in kwargs.items() if k != "_function_middleware_pipeline"} + if "_function_middleware_pipeline" in kwargs: + filtered_kwargs = {k: v for k, v in kwargs.items() if k != "_function_middleware_pipeline"} + else: + filtered_kwargs = kwargs async for update in self._inner_get_streaming_response( messages=prepped_messages, chat_options=chat_options, **filtered_kwargs @@ -474,7 +480,7 @@ class BaseChatClient(AFBaseModel, ABC): | None = None, chat_message_store_factory: Callable[[], ChatMessageStore] | None = None, context_providers: ContextProvider | list[ContextProvider] | AggregateContextProvider | None = None, - middleware: MiddlewareType | list[MiddlewareType] | None = None, + middleware: Middleware | list[Middleware] | None = None, **kwargs: Any, ) -> "ChatAgent": """Create an agent with the given name and instructions. diff --git a/python/packages/main/agent_framework/_middleware.py b/python/packages/main/agent_framework/_middleware.py index 23dc2110b5..351ecab64c 100644 --- a/python/packages/main/agent_framework/_middleware.py +++ b/python/packages/main/agent_framework/_middleware.py @@ -2,7 +2,8 @@ from abc import ABC, abstractmethod from collections.abc import AsyncIterable, Awaitable, Callable -from typing import TYPE_CHECKING, Any, TypeVar +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, TypeAlias, TypeVar from ._types import AgentRunResponse, AgentRunResponseUpdate, ChatMessage @@ -19,11 +20,12 @@ __all__ = [ "AgentRunContext", "FunctionInvocationContext", "FunctionMiddleware", - "MiddlewareType", + "Middleware", "use_agent_middleware", ] +@dataclass class AgentRunContext: """Context object for agent middleware invocations. @@ -38,28 +40,14 @@ class AgentRunContext: For streaming: should be AsyncIterable[AgentRunResponseUpdate] """ - def __init__( - self, - agent: "AgentProtocol", - messages: list[ChatMessage], - is_streaming: bool = False, - 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. - metadata: Metadata dictionary. - """ - self.agent = agent - self.messages = messages - self.is_streaming = is_streaming - self.metadata = metadata or {} - self.result: AgentRunResponse | AsyncIterable[AgentRunResponseUpdate] | None = None + agent: "AgentProtocol" + messages: list[ChatMessage] + is_streaming: bool = False + metadata: dict[str, Any] = field(default_factory=lambda: {}) + result: AgentRunResponse | AsyncIterable[AgentRunResponseUpdate] | None = None +@dataclass class FunctionInvocationContext: """Context object for function middleware invocations. @@ -71,23 +59,10 @@ class FunctionInvocationContext: to see the actual execution result or can be set to override the execution result. """ - def __init__( - self, - function: "AIFunction[Any, Any]", - arguments: "BaseModel", - metadata: dict[str, Any] | None = None, - ) -> None: - """Initialize function invocation context. - - Args: - function: The function being invoked. - arguments: The validated arguments for the function. - metadata: Metadata dictionary. - """ - self.function = function - self.arguments = arguments - self.metadata = metadata or {} - self.result: Any = None + function: "AIFunction[Any, Any]" + arguments: "BaseModel" + metadata: dict[str, Any] = field(default_factory=lambda: {}) + result: Any = None class AgentMiddleware(ABC): @@ -153,7 +128,7 @@ FunctionMiddlewareCallable = Callable[ ] # Type alias for all middleware types -MiddlewareType = AgentMiddleware | AgentMiddlewareCallable | FunctionMiddleware | FunctionMiddlewareCallable +Middleware: TypeAlias = AgentMiddleware | AgentMiddlewareCallable | FunctionMiddleware | FunctionMiddlewareCallable class AgentMiddlewareWrapper(AgentMiddleware): @@ -451,12 +426,12 @@ def use_agent_middleware(agent_class: type[TAgent]) -> type[TAgent]: 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: + def _initialize_middleware_pipelines(self: Any, middlewares: Middleware | list[Middleware] | 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 + middleware_list: list[Middleware] = middlewares if isinstance(middlewares, list) else [middlewares] # type: ignore # Separate agent and function middleware using isinstance checks agent_middlewares: list[AgentMiddleware | AgentMiddlewareCallable] = [] diff --git a/python/samples/getting_started/middleware/class_based_middleware.py b/python/samples/getting_started/middleware/class_based_middleware.py index e8e501d560..d29b79c4fa 100644 --- a/python/samples/getting_started/middleware/class_based_middleware.py +++ b/python/samples/getting_started/middleware/class_based_middleware.py @@ -1,5 +1,19 @@ # Copyright (c) Microsoft. All rights reserved. +""" +Class-based Middleware Example + +This sample demonstrates how to implement middleware using class-based approach by inheriting +from AgentMiddleware and FunctionMiddleware base classes. The example includes: + +- SecurityAgentMiddleware: Checks for security violations in user queries and blocks requests + containing sensitive information like passwords or secrets +- LoggingFunctionMiddleware: Logs function execution details including timing and parameters + +This approach is useful when you need stateful middleware or complex logic that benefits +from object-oriented design patterns. +""" + import asyncio import time from collections.abc import Awaitable, Callable diff --git a/python/samples/getting_started/middleware/exception_handling_with_middleware.py b/python/samples/getting_started/middleware/exception_handling_with_middleware.py index 7646ff68ed..737dde3a13 100644 --- a/python/samples/getting_started/middleware/exception_handling_with_middleware.py +++ b/python/samples/getting_started/middleware/exception_handling_with_middleware.py @@ -1,5 +1,19 @@ # Copyright (c) Microsoft. All rights reserved. +""" +Exception Handling with Middleware + +This sample demonstrates how to use middleware for centralized exception handling in function calls. +The example shows: + +- How to catch exceptions thrown by functions and provide graceful error responses +- Overriding function results when errors occur to provide user-friendly messages +- Using middleware to implement retry logic, fallback mechanisms, or error reporting + +The middleware catches TimeoutError from an unstable data service and replaces it with +a helpful message for the user, preventing raw exceptions from reaching the end user. +""" + import asyncio from collections.abc import Awaitable, Callable from typing import Annotated diff --git a/python/samples/getting_started/middleware/function_based_middleware.py b/python/samples/getting_started/middleware/function_based_middleware.py index da0cfd6985..7b2593031b 100644 --- a/python/samples/getting_started/middleware/function_based_middleware.py +++ b/python/samples/getting_started/middleware/function_based_middleware.py @@ -1,5 +1,20 @@ # Copyright (c) Microsoft. All rights reserved. +""" +Function-based Middleware Example + +This sample demonstrates how to implement middleware using simple async functions instead of classes. +The example includes: + +- Security middleware that validates agent requests for sensitive information +- Logging middleware that tracks function execution timing and parameters +- Performance monitoring to measure execution duration + +Function-based middleware is ideal for simple, stateless operations and provides a more +lightweight approach compared to class-based middleware. Both agent and function middleware +can be implemented as async functions that accept context and next parameters. +""" + import asyncio import time from collections.abc import Awaitable, Callable diff --git a/python/samples/getting_started/middleware/override_result_with_middleware.py b/python/samples/getting_started/middleware/override_result_with_middleware.py index d038dc9d8b..af23705f50 100644 --- a/python/samples/getting_started/middleware/override_result_with_middleware.py +++ b/python/samples/getting_started/middleware/override_result_with_middleware.py @@ -1,5 +1,20 @@ # Copyright (c) Microsoft. All rights reserved. +""" +Result Override with Middleware + +This sample demonstrates how to use middleware to intercept and modify function results +after execution. The example shows: + +- How to execute the original function first and then modify its result +- Replacing function outputs with custom messages or transformed data +- Using middleware for result filtering, formatting, or enhancement + +The weather override middleware lets the original weather function execute normally, +then replaces its result with a custom "perfect weather" message, demonstrating +how middleware can be used for content filtering, A/B testing, or result enhancement. +""" + import asyncio from collections.abc import Awaitable, Callable from random import randint