aip-agents-binary 0.0.0__py3-none-any.whl

aip_agents/middleware/manager.py

@@ -0,0 +1,291 @@
+"""Middleware orchestration and lifecycle management.
+
+This module provides MiddlewareManager, which coordinates multiple middleware
+components and manages their lifecycle hooks.
+
+Architecture Relationship
+-------------------------
+
+MiddlewareManager works with AgentMiddleware (defined in base.py):
+
+- AgentMiddleware: Protocol that individual middleware classes implement
+- MiddlewareManager: Orchestrates multiple middleware instances
+
+The manager is used by agents to interact with the middleware stack:
+
+    agent -> MiddlewareManager -> [Middleware1, Middleware2, ...]
+                │
+                ├── get_all_tools()          → Collects tools from all middleware
+                ├── build_system_prompt()    → Builds combined system prompt
+                ├── abefore_run()            → Calls before_run on all middleware
+                ├── abefore_model()          → Calls before_model on all middleware
+                ├── modify_model_request()   → Chains request modifications
+                ├── aafter_model()           → Calls after_model on all middleware
+                └── aafter_run()             → Calls after_run on all middleware
+
+Hook Execution Order
+--------------------
+
+- Forward order (first → last): before_run, before_model, modify_model_request, after_run, on_final_response
+- Reverse order (last → first): after_model
+
+Error Handling
+--------------
+
+All hook execution is "best-effort":
+- Exceptions are logged but don't stop execution
+- Middleware failures don't block agent responses
+- Each middleware hook failure is isolated
+"""
+
+from typing import Any
+
+from langchain_core.tools import BaseTool
+
+from aip_agents.middleware.base import AgentMiddleware, ModelRequest
+from aip_agents.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class MiddlewareManager:
+    """Orchestrates multiple middleware components and manages hook execution.
+
+    The manager collects tools from all middleware, builds enhanced system prompts,
+    and executes lifecycle hooks in the configured order:
+    forward for run-level hooks and before/modify model hooks, reverse for after_model.
+
+    Attributes:
+        middleware: List of middleware components in registration order.
+    """
+
+    def __init__(self, middleware: list[AgentMiddleware]) -> None:
+        """Initialize the middleware manager.
+
+        Args:
+            middleware: List of middleware components to manage. Order matters:
+                       hooks execute forward (first to last) for run-level hooks and
+                       before/modify model hooks; after_model executes in reverse order.
+        """
+        self.middleware = middleware
+
+    def __bool__(self) -> bool:
+        """Return True when at least one middleware is configured."""
+        return bool(self.middleware)
+
+    def get_all_tools(self) -> list[BaseTool]:
+        """Collect tools from all registered middleware.
+
+        Returns:
+            Combined list of all tools contributed by all middleware components.
+            Empty list if no middleware or if middleware provide no tools.
+        """
+        tools: list[BaseTool] = []
+        for mw in self.middleware:
+            tools.extend(mw.tools)
+        return tools
+
+    def build_system_prompt(self, base_instruction: str) -> str:
+        """Build enhanced system prompt by concatenating base instruction with middleware additions.
+
+        Args:
+            base_instruction: The base system prompt for the agent.
+
+        Returns:
+            Enhanced system prompt with all middleware additions appended.
+            If no middleware provide additions, returns base_instruction unchanged.
+        """
+        parts = [base_instruction]
+
+        for mw in self.middleware:
+            if mw.system_prompt_additions:
+                parts.append(mw.system_prompt_additions)
+
+        return "\n\n".join(parts)
+
+    def modify_model_request(self, request: ModelRequest, state: dict[str, Any]) -> ModelRequest:
+        """Execute modify_model_request hooks for all middleware in forward order.
+
+        Each middleware receives the request modified by previous middleware,
+        allowing them to build on each other's changes.
+
+        Args:
+            request: The model request to be modified.
+            state: Current agent state for context.
+
+        Returns:
+            Final modified request after all middleware have processed it.
+        """
+        current_request = request
+
+        for mw in self.middleware:
+            current_request = mw.modify_model_request(current_request, state)
+
+        return current_request
+
+    def before_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """Execute before_model hooks for all middleware in forward order.
+
+        Hooks run first to last, allowing earlier middleware to prepare state
+        for later middleware.
+
+        Args:
+            state: Current agent state.
+
+        Returns:
+            Merged dictionary of all state updates from all middleware.
+            Updates are accumulated in order of execution.
+        """
+        state_updates: dict[str, Any] = {}
+
+        for mw in self.middleware:
+            mw_updates = mw.before_model(state)
+            if mw_updates:
+                state_updates.update(mw_updates)
+
+        return state_updates
+
+    async def abefore_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """Asynchronously execute before_model hooks for all middleware."""
+        state_updates: dict[str, Any] = {}
+
+        for mw in self.middleware:
+            mw_updates = await mw.abefore_model(state)
+            if mw_updates:
+                state_updates.update(mw_updates)
+
+        return state_updates
+
+    def after_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """Execute after_model hooks for all middleware in reverse order.
+
+        Hooks run last to first (reverse of registration order), allowing
+        proper cleanup and unwinding of middleware operations.
+
+        Args:
+            state: Current agent state after model invocation.
+
+        Returns:
+            Merged dictionary of all state updates from all middleware.
+            Updates are accumulated in reverse order of execution.
+        """
+        updates: dict[str, Any] = {}
+
+        # Execute in reverse order
+        for mw in reversed(self.middleware):
+            mw_updates = mw.after_model(state)
+            if mw_updates:
+                updates.update(mw_updates)
+
+        return updates
+
+    async def aafter_model(self, state: dict[str, Any]) -> dict[str, Any]:
+        """Asynchronously execute after_model hooks for all middleware in reverse order."""
+        updates: dict[str, Any] = {}
+
+        for mw in reversed(self.middleware):
+            mw_updates = await mw.aafter_model(state)
+            if mw_updates:
+                updates.update(mw_updates)
+
+        return updates
+
+    # ------------------------------------------------------------------
+    # Run-level lifecycle hooks (Milestone 1: Middleware Run Lifecycle Hooks)
+    # ------------------------------------------------------------------
+
+    async def abefore_run(self, state: dict[str, Any], config: dict[str, Any] | None) -> dict[str, Any]:
+        """Execute before_run hooks for all middleware in forward order.
+
+        Hooks are optional; only executed if middleware defines them.
+        Best-effort execution: exceptions are logged and execution continues.
+
+        Args:
+            state: Current agent state.
+            config: Optional run configuration dict.
+
+        Returns:
+            Merged dictionary of all state updates from all middleware.
+        """
+        state_updates: dict[str, Any] = {}
+
+        for mw in self.middleware:
+            # Keep duck-typed compatibility for middleware that may not fully
+            # implement optional run hooks despite protocol defaults.
+            mw_any: Any = mw
+            hook = getattr(mw_any, "abefore_run", None)
+            if hook is not None:
+                try:
+                    mw_updates_any = await hook(state, config)
+                    if isinstance(mw_updates_any, dict) and mw_updates_any:
+                        state_updates.update(mw_updates_any)
+                except Exception as exc:
+                    logger.warning(
+                        "Middleware %s.abefore_run failed: %s",
+                        type(mw).__name__,
+                        exc,
+                        exc_info=True,
+                    )
+
+        return state_updates
+
+    async def aafter_run(
+        self,
+        *,
+        final_state: dict[str, Any],
+        output: str | None,
+        config: dict[str, Any] | None,
+        error: Exception | None,
+    ) -> None:
+        """Execute after_run hooks for all middleware in forward order.
+
+        Hooks are optional; only executed if middleware defines them.
+        Best-effort execution: exceptions are logged and execution continues.
+
+        Args:
+            final_state: Final agent state after run completion.
+            output: The final output string from the agent, or None if failed.
+            config: Optional run configuration dict.
+            error: Exception if the run failed, None if successful.
+        """
+        for mw in self.middleware:
+            mw_any: Any = mw
+            hook = getattr(mw_any, "aafter_run", None)
+            if hook is not None:
+                try:
+                    await hook(final_state=final_state, output=output, config=config, error=error)
+                except Exception as exc:
+                    logger.warning(
+                        "Middleware %s.aafter_run failed: %s",
+                        type(mw).__name__,
+                        exc,
+                        exc_info=True,
+                    )
+
+    async def aon_final_response(self, content: str, context: dict[str, Any]) -> None:
+        """Invoke streaming final-response hooks (best-effort).
+
+        Called when an A2A FINAL_RESPONSE event is captured during streaming.
+
+        Args:
+            content: The final response content.
+            context: Context dict (thread_id, original_query, memory_user_id, etc.).
+        """
+        for mw in self.middleware:
+            mw_any: Any = mw
+            hook = getattr(mw_any, "aon_final_response", None)
+            if hook is None:
+                hook = getattr(mw_any, "on_final_response", None)  # pragma: no cover
+            if hook is None:
+                continue  # pragma: no cover
+            try:
+                result = hook(content, context)
+                if hasattr(result, "__await__"):
+                    await result
+            except Exception as exc:
+                logger.warning(
+                    "Middleware %s.on_final_response failed: %s",
+                    type(mw).__name__,
+                    exc,
+                    exc_info=True,
+                )