From eb56b307359ebabe95aef5be9d1d1ced30f2ee01 Mon Sep 17 00:00:00 2001 From: Dmytro Struk <13853051+dmytrostruk@users.noreply.github.com> Date: Thu, 18 Sep 2025 09:57:35 -0700 Subject: [PATCH] Addressed PR feedback --- python/DEV_SETUP.md | 17 --- .../main/agent_framework/_middleware.py | 107 ++++++++++++------ python/pytest.ini | 9 -- .../middleware/class_based_middleware.py | 40 ++++--- .../exception_handling_with_middleware.py | 18 +-- .../middleware/function_based_middleware.py | 28 ++--- .../override_result_with_middleware.py | 21 ++-- 7 files changed, 132 insertions(+), 108 deletions(-) delete mode 100644 python/pytest.ini diff --git a/python/DEV_SETUP.md b/python/DEV_SETUP.md index ad33153d09..a041712443 100644 --- a/python/DEV_SETUP.md +++ b/python/DEV_SETUP.md @@ -75,23 +75,6 @@ Open the `python` folder in [VSCode](https://code.visualstudio.com/docs/editor/w Open any of the `.py` files in the project and run the `Python: Select Interpreter` command from the command palette. Make sure the virtual env (default path is `.venv`) created by `uv` is selected. -### Configuring Unit Testing in VSCode - -- We have removed the strict dependency on forcing `pytest` usage via the `.vscode/settings.json` file. -- Developers are free to set up unit tests using their preferred framework, whether it is `pytest` or `unittest`. -- If needed, adjust VSCode's local `settings.json` (accessed via the Command Palette(`Ctrl+Shift+P`) and type `Preferences: Open User Settings (JSON)`) to configure the test framework. For example: - - ```json - "python.testing.unittestEnabled": false, - "python.testing.pytestEnabled": true, - ``` - - Or, for `unittest`: - - ```json - "python.testing.unittestEnabled": true, - "python.testing.pytestEnabled": false, - ## LLM setup Make sure you have an diff --git a/python/packages/main/agent_framework/_middleware.py b/python/packages/main/agent_framework/_middleware.py index 351ecab64c..a069463fa2 100644 --- a/python/packages/main/agent_framework/_middleware.py +++ b/python/packages/main/agent_framework/_middleware.py @@ -18,6 +18,7 @@ TAgent = TypeVar("TAgent", bound="AgentProtocol") __all__ = [ "AgentMiddleware", "AgentRunContext", + "BaseMiddlewarePipeline", "FunctionInvocationContext", "FunctionMiddleware", "Middleware", @@ -159,7 +160,65 @@ class FunctionMiddlewareWrapper(FunctionMiddleware): await self.func(context, next) -class AgentMiddlewarePipeline: +class BaseMiddlewarePipeline(ABC): + """Base class for middleware pipeline execution.""" + + def __init__(self) -> None: + """Initialize the base middleware pipeline.""" + self._middlewares: list[Any] = [] + + @abstractmethod + def _register_middleware(self, middleware: Any) -> None: + """Register a middleware item. Must be implemented by subclasses.""" + ... + + @property + def has_middlewares(self) -> bool: + """Check if there are any middlewares registered.""" + return bool(self._middlewares) + + def _create_handler_chain( + self, + context: Any, + final_handler: Callable[[Any], Awaitable[Any]], + result_container: dict[str, Any], + result_key: str = "result", + ) -> Callable[[Any], Awaitable[None]]: + """Create a chain of middleware handlers. + + Args: + context: The execution context + final_handler: The final handler to execute + result_container: Container to store the result + result_key: Key to use in the result container + + Returns: + The first handler in the chain + """ + + def create_next_handler(index: int) -> Callable[[Any], Awaitable[None]]: + if index >= len(self._middlewares): + + async def final_wrapper(c: Any) -> None: + # Execute actual handler and populate context for observability + result = await final_handler(c) + result_container[result_key] = result + c.result = result + + return final_wrapper + + middleware = self._middlewares[index] + next_handler = create_next_handler(index + 1) + + async def current_handler(c: Any) -> None: + await middleware.process(c, next_handler) + + return current_handler + + return create_next_handler(0) + + +class AgentMiddlewarePipeline(BaseMiddlewarePipeline): """Executes agent middleware in a chain.""" def __init__(self, middlewares: list[AgentMiddleware | AgentMiddlewareCallable] | None = None): @@ -168,6 +227,7 @@ class AgentMiddlewarePipeline: Args: middlewares: List of agent middleware to include in the pipeline. """ + super().__init__() self._middlewares: list[AgentMiddleware] = [] if middlewares: @@ -312,13 +372,8 @@ class AgentMiddlewarePipeline: async for update in result_stream: yield update - @property - def has_middlewares(self) -> bool: - """Check if there are any middlewares registered.""" - return bool(self._middlewares) - -class FunctionMiddlewarePipeline: +class FunctionMiddlewarePipeline(BaseMiddlewarePipeline): """Executes function middleware in a chain.""" def __init__(self, middlewares: list[FunctionMiddleware | FunctionMiddlewareCallable] | None = None): @@ -327,6 +382,7 @@ class FunctionMiddlewarePipeline: Args: middlewares: List of function middleware to include in the pipeline. """ + super().__init__() self._middlewares: list[FunctionMiddleware] = [] if middlewares: @@ -369,31 +425,15 @@ class FunctionMiddlewarePipeline: # 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): + # Custom final handler that handles pre-existing results + async def function_final_handler(c: FunctionInvocationContext) -> Any: + # If result was set before calling next(), skip execution + if c.result is not None: + return c.result + # Execute actual handler and populate context for observability + return await final_handler(c) - async def final_wrapper(c: FunctionInvocationContext) -> None: - # If result was set before calling next(), skip execution - if c.result is not None: - result_container["result"] = c.result - return - - # Execute actual handler and populate context for observability - result = await final_handler(c) - result_container["result"] = result - c.result = result - - 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) + first_handler = self._create_handler_chain(context, function_final_handler, result_container, "result") await first_handler(context) # Return the result from result container or overridden result @@ -401,11 +441,6 @@ class FunctionMiddlewarePipeline: return context.result return result_container["result"] - @property - 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]: diff --git a/python/pytest.ini b/python/pytest.ini deleted file mode 100644 index 5b3e31cdaa..0000000000 --- a/python/pytest.ini +++ /dev/null @@ -1,9 +0,0 @@ -[pytest] -testpaths = packages/main/tests packages/azure/tests packages/foundry/tests packages/copilotstudio/tests packages/mem0/tests packages/runtime/tests -python_files = test_*.py *_test.py -python_classes = Test* -python_functions = test_* -addopts = -ra -q -r fEX --tb=short -asyncio_mode = auto -asyncio_default_fixture_loop_scope = function -timeout = 120 \ No newline at end of file diff --git a/python/samples/getting_started/middleware/class_based_middleware.py b/python/samples/getting_started/middleware/class_based_middleware.py index d29b79c4fa..4958e3d094 100644 --- a/python/samples/getting_started/middleware/class_based_middleware.py +++ b/python/samples/getting_started/middleware/class_based_middleware.py @@ -1,5 +1,24 @@ # 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 ( + AgentMiddleware, + AgentRunContext, + AgentRunResponse, + ChatMessage, + FunctionInvocationContext, + FunctionMiddleware, + Role, +) +from agent_framework.foundry import FoundryChatClient +from azure.identity.aio import AzureCliCredential +from pydantic import Field + """ Class-based Middleware Example @@ -14,17 +33,6 @@ This approach is useful when you need stateful middleware or complex logic that from object-oriented design patterns. """ -import asyncio -import time -from collections.abc import Awaitable, Callable -from random import randint -from typing import Annotated - -from agent_framework import AgentMiddleware, AgentRunContext, 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.")], @@ -49,6 +57,12 @@ class SecurityAgentMiddleware(AgentMiddleware): query = last_message.text if "password" in query.lower() or "secret" in query.lower(): print("[SecurityAgentMiddleware] Security Warning: Detected sensitive information, blocking request.") + # Override the result with warning message + context.result = AgentRunResponse( + messages=[ + ChatMessage(role=Role.ASSISTANT, text="Detected sensitive information, the request is blocked.") + ] + ) # Simply don't call next() to prevent execution return @@ -97,14 +111,14 @@ async def main() -> None: query = "What's the weather like in Seattle?" print(f"User: {query}") result = await agent.run(query) - print(f"Agent: {result.text if result.text else 'No response'}\n") + print(f"Agent: {result.text}\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.text if result.text else 'No response'}\n") + print(f"Agent: {result.text}\n") if __name__ == "__main__": 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 737dde3a13..ca37a75bd1 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,14 @@ # Copyright (c) Microsoft. All rights reserved. +import asyncio +from collections.abc import Awaitable, Callable +from typing import Annotated + +from agent_framework import FunctionInvocationContext +from agent_framework.foundry import FoundryChatClient +from azure.identity.aio import AzureCliCredential +from pydantic import Field + """ Exception Handling with Middleware @@ -14,15 +23,6 @@ The middleware catches TimeoutError from an unstable data service and replaces i 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 - -from agent_framework import FunctionInvocationContext -from agent_framework.foundry import FoundryChatClient -from azure.identity.aio import AzureCliCredential -from pydantic import Field - def unstable_data_service( query: Annotated[str, Field(description="The data query to execute.")], diff --git a/python/samples/getting_started/middleware/function_based_middleware.py b/python/samples/getting_started/middleware/function_based_middleware.py index 7b2593031b..e6b92494b3 100644 --- a/python/samples/getting_started/middleware/function_based_middleware.py +++ b/python/samples/getting_started/middleware/function_based_middleware.py @@ -1,5 +1,19 @@ # 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 ( + AgentRunContext, + FunctionInvocationContext, +) +from agent_framework.foundry import FoundryChatClient +from azure.identity.aio import AzureCliCredential +from pydantic import Field + """ Function-based Middleware Example @@ -15,20 +29,6 @@ lightweight approach compared to class-based middleware. Both agent and function can be implemented as async functions that accept context and next parameters. """ -import asyncio -import time -from collections.abc import Awaitable, Callable -from random import randint -from typing import Annotated - -from agent_framework import ( - AgentRunContext, - 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.")], 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 af23705f50..b56cb6c78b 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,15 @@ # Copyright (c) Microsoft. All rights reserved. +import asyncio +from collections.abc import Awaitable, Callable +from random import randint +from typing import Annotated + +from agent_framework import FunctionInvocationContext +from agent_framework.foundry import FoundryChatClient +from azure.identity.aio import AzureCliCredential +from pydantic import Field + """ Result Override with Middleware @@ -15,16 +25,6 @@ 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 -from typing import Annotated - -from agent_framework import 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.")], @@ -48,6 +48,7 @@ async def weather_override_middleware( print(f"[WeatherOverrideMiddleware] Original result: {original_result}") # Override with a custom message + # It's also possible to override the result before "next()" call if needed custom_message = ( "Weather Advisory - due to special atmospheric conditions, " "all locations are experiencing perfect weather today! "