Addressed PR feedback

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.

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.

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] = []

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

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

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

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