Introduce middleware (#69)

Author: szykol

splunklib/ai/README.md

@@ -386,6 +386,163 @@ async with Agent(
 
 **Note**: Currently input schemas can only be used by subagents, not by regular agents.
 
+## Middleware
+
+Middleware lets you intercept model, tool, and subagent calls in a request/handler chain.
+Each middleware can inspect input, call `handler(request)`, and modify the returned response.
+
+Available decorators:
+
+- `model_middleware`
+- `tool_middleware`
+- `subagent_middleware`
+
+Class-based middleware:
+
+```py
+from typing import override
+from splunklib.ai.middleware import (
+    AgentMiddleware,
+    ModelMiddlewareHandler,
+    ModelRequest,
+    SubagentMiddlewareHandler,
+    SubagentRequest,
+    SubagentResponse,
+    ToolMiddlewareHandler,
+    ToolRequest,
+    ToolResponse,
+)
+from splunklib.ai.messages import AIMessage
+
+
+class ExampleMiddleware(AgentMiddleware):
+    @override
+    async def model_middleware(
+        self, request: ModelRequest, handler: ModelMiddlewareHandler
+    ) -> AIMessage:
+        request.system_message = request.system_message.replace("SECRET", "[REDACTED]")
+        return await handler(request)
+
+    @override
+    async def tool_middleware(
+        self, request: ToolRequest, handler: ToolMiddlewareHandler
+    ) -> ToolResponse:
+        if request.call.name == "temperature":
+            return ToolResponse(content="25.0")
+        return await handler(request)
+
+    @override
+    async def subagent_middleware(
+        self, request: SubagentRequest, handler: SubagentMiddlewareHandler
+    ) -> SubagentResponse:
+        if request.call.name == "SummaryAgent":
+            return SubagentResponse(
+                content="Executive summary: no critical incidents detected."
+            )
+        return await handler(request)
+```
+
+Example model middleware:
+
+```py
+from splunklib.ai.middleware import (
+    model_middleware,
+    ModelMiddlewareHandler,
+    ModelRequest,
+)
+from splunklib.ai.messages import AIMessage
+
+
+@model_middleware
+async def redact_system_prompt(
+    request: ModelRequest, handler: ModelMiddlewareHandler
+) -> AIMessage:
+    request.system_message = request.system_message.replace("SECRET", "[REDACTED]")
+    return await handler(request)
+```
+
+Example tool middleware:
+
+```py
+from splunklib.ai.middleware import (
+    tool_middleware,
+    ToolMiddlewareHandler,
+    ToolRequest,
+    ToolResponse,
+)
+
+
+@tool_middleware
+async def mock_temperature(
+    request: ToolRequest, handler: ToolMiddlewareHandler
+) -> ToolResponse:
+    if request.call.name == "temperature":
+        return ToolResponse(content="25.0")
+    return await handler(request)
+```
+
+Example subagent middleware:
+
+```py
+from splunklib.ai.middleware import (
+    subagent_middleware,
+    SubagentMiddlewareHandler,
+    SubagentRequest,
+    SubagentResponse,
+)
+
+
+@subagent_middleware
+async def mock_subagent(
+    request: SubagentRequest, handler: SubagentMiddlewareHandler
+) -> SubagentResponse:
+    if request.call.name == "SummaryAgent":
+        return SubagentResponse(
+            content="Executive summary: no critical incidents detected."
+        )
+    return await handler(request)
+```
+
+Retry pattern (bounded retries):
+
+```py
+from splunklib.ai.middleware import (
+    tool_middleware,
+    ToolMiddlewareHandler,
+    ToolRequest,
+    ToolResponse,
+)
+
+
+class RetryableToolError(Exception): pass
+
+
+@tool_middleware
+async def retry_transient_tool_failures(
+    request: ToolRequest, handler: ToolMiddlewareHandler
+) -> ToolResponse:
+    last_error: Exception | None = None
+    for _ in range(3):
+        try:
+            return await handler(request)
+        except RetryableToolError as e:
+            last_error = e
+
+    assert last_error is not None
+    raise last_error
+```
+
+Pass middleware to `Agent`:
+
+```py
+async with Agent(
+    model=model,
+    service=service,
+    system_prompt="...",
+    middleware=[redact_system_prompt, mock_temperature, mock_subagent],
+) as agent: ...
+```
+
 ## Hooks
 
 Hooks are user-defined callback functions that can be registered to execute at specific points

splunklib/ai/agent.py

@@ -12,10 +12,10 @@
 # License for the specific language governing permissions and limitations
 # under the License.
 
-from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager
-from logging import Logger
 import os
 from collections.abc import AsyncGenerator, Sequence
+from contextlib import AbstractAsyncContextManager, AsyncExitStack, asynccontextmanager
+from logging import Logger
 from typing import Self, final, override
 
 from pydantic import BaseModel
@@ -25,6 +25,7 @@
 from splunklib.ai.core.backend_registry import get_backend
 from splunklib.ai.hooks import AgentHook
 from splunklib.ai.messages import AgentResponse, BaseMessage, OutputT
+from splunklib.ai.middleware import AgentMiddleware
 from splunklib.ai.model import PredefinedModel
 from splunklib.ai.tool_filtering import ToolFilters, filter_tools
 from splunklib.ai.tools import (
@@ -130,6 +131,7 @@ def __init__(
         output_schema: type[OutputT] | None = None,
         input_schema: type[BaseModel] | None = None,  # Only used by Subagents
         hooks: Sequence[AgentHook] | None = None,
+        middleware: Sequence[AgentMiddleware] | None = None,
         name: str = "",  # Only used by Subagents
         description: str = "",  # Only used by Subagents
         logger: Logger | None = None,
@@ -143,6 +145,7 @@ def __init__(
             input_schema=input_schema,
             output_schema=output_schema,
             hooks=hooks,
+            middleware=middleware,
             logger=logger,
         )
 

splunklib/ai/base_agent.py

@@ -12,16 +12,17 @@
 # License for the specific language governing permissions and limitations
 # under the License.
 
-from abc import ABC, abstractmethod
-from collections.abc import Sequence
 import logging
 import secrets
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
 from typing import Generic
 
 from pydantic import BaseModel
 
 from splunklib.ai.hooks import AgentHook
 from splunklib.ai.messages import AgentResponse, BaseMessage, OutputT
+from splunklib.ai.middleware import AgentMiddleware
 from splunklib.ai.model import PredefinedModel
 from splunklib.ai.tools import Tool
 
@@ -36,6 +37,7 @@ class BaseAgent(Generic[OutputT], ABC):
     _input_schema: type[BaseModel] | None = None
     _output_schema: type[OutputT] | None = None
     _hooks: Sequence[AgentHook] | None = None
+    _middleware: Sequence[AgentMiddleware] | None = None
     _trace_id: str
     _logger: logging.Logger
 
@@ -50,6 +52,7 @@ def __init__(
         input_schema: type[BaseModel] | None = None,
         output_schema: type[OutputT] | None = None,
         hooks: Sequence[AgentHook] | None = None,
+        middleware: Sequence[AgentMiddleware] | None = None,
         logger: logging.Logger | None = None,
     ) -> None:
         self._system_prompt = system_prompt
@@ -61,6 +64,7 @@ def __init__(
         self._input_schema = input_schema
         self._output_schema = output_schema
         self._hooks = tuple(hooks) if hooks else ()
+        self._middleware = tuple(middleware) if middleware else ()
         self._trace_id = secrets.token_hex(16)  # 32 Hex characters
 
         if logger is None:
@@ -112,6 +116,10 @@ def output_schema(self) -> type[OutputT] | None:
     def hooks(self) -> Sequence[AgentHook] | None:
         return self._hooks
 
+    @property
+    def middleware(self) -> Sequence[AgentMiddleware] | None:
+        return self._middleware
+
     @property
     def trace_id(self) -> str:
         return self._trace_id