Add agent middleware (#78)

Author: mateusz834

splunklib/ai/README.md

@@ -393,16 +393,18 @@ Each middleware can inspect input, call `handler(request)`, and modify the retur
 
 Available decorators:
 
+- `agent_middleware`
 - `model_middleware`
 - `tool_middleware`
 - `subagent_middleware`
 
 Class-based middleware:
 
 ```py
-from typing import override
+from typing import Any, override
 from splunklib.ai.middleware import (
-    AgentMiddleware,
+	AgentMiddlewareHandler,
+    AgentRequest,
     ModelMiddlewareHandler,
     ModelRequest,
     SubagentMiddlewareHandler,
@@ -412,10 +414,20 @@ from splunklib.ai.middleware import (
     ToolRequest,
     ToolResponse,
 )
-from splunklib.ai.messages import AIMessage
+from splunklib.ai.messages import AIMessage, AgentResponse, ToolCall
 
 
 class ExampleMiddleware(AgentMiddleware):
+    @override
+    async def agent_middleware(
+        self, request: AgentRequest, handler: AgentMiddlewareHandler
+    ) -> AgentResponse[Any | None]:
+        # Keep retrying until the agent makes at least one tool call.
+        resp = await handler(request)
+        while not any(m for m in resp.messages if isinstance(m, ToolCall)):
+            resp = await handler(request)
+        return resp
+
     @override
     async def model_middleware(
         self, request: ModelRequest, handler: ModelMiddlewareHandler
@@ -442,6 +454,29 @@ class ExampleMiddleware(AgentMiddleware):
         return await handler(request)
 ```
 
+Example agent middleware:
+
+```py
+from typing import Any
+from splunklib.ai.middleware import (
+    agent_middleware,
+    AgentMiddlewareHandler,
+    AgentRequest,
+)
+from splunklib.ai.messages import AgentResponse, ToolCall
+
+
+@agent_middleware
+async def force_tool_call(
+    request: AgentRequest, handler: AgentMiddlewareHandler
+) -> AgentResponse[Any | None]:
+    # Keep retrying until the agent makes at least one tool call.
+    resp = await handler(request)
+    while not any(m for m in resp.messages if isinstance(m, ToolCall)):
+        resp = await handler(request)
+    return resp
+```
+
 Example model middleware:
 
 ```py

splunklib/ai/engines/langchain.py

@@ -79,8 +79,10 @@
     ToolMessage,
 )
 from splunklib.ai.middleware import (
+    AgentMiddlewareHandler,
     AgentState,
     AgentMiddleware,
+    AgentRequest,
     ModelMiddlewareHandler,
     ModelRequest,
     SubagentMiddlewareHandler,
@@ -122,55 +124,128 @@ class LangChainAgentImpl(AgentImpl[OutputT]):
     _thread_id: uuid.UUID
     _config: RunnableConfig
     _output_schema: type[OutputT] | None
+    _middleware: Sequence[AgentMiddleware]
 
     def __init__(
         self,
         system_prompt: str,
         model: BaseChatModel,
         tools: list[BaseTool],
         output_schema: type[OutputT] | None,
-        middleware: Sequence[LC_AgentMiddleware] | None = None,
+        lcmiddleware: Sequence[LC_AgentMiddleware] | None = None,
+        middleware: Sequence[AgentMiddleware] | None = None,
     ) -> None:
         super().__init__()
         self._output_schema = output_schema
         self._thread_id = uuid.uuid4()
         self._config = {"configurable": {"thread_id": self._thread_id}}
+        self._middleware = middleware or []
 
         checkpointer = InMemorySaver()
-        middleware = middleware or []
 
         self._agent = create_agent(
             model=model,
             tools=tools,
             system_prompt=system_prompt,
             checkpointer=checkpointer,
             response_format=output_schema,
-            middleware=middleware,
+            middleware=lcmiddleware or [],
         )
 
+    def _with_agent_middleware(
+        self,
+        agent_invoke: Callable[[AgentRequest], Awaitable[AgentResponse[Any | None]]],
+    ) -> Callable[[AgentRequest], Awaitable[AgentResponse[Any | None]]]:
+        # When provided with a list of middlewares, e.g. [m1, m2, m3],
+        # they are executed in the following order:
+        #
+        # m1 -> m2 -> m3 -> agent_invoke
+        #
+        # Each middleware wraps the next one in the chain.
+        #
+        # - m1's handler calls m2.agent_middleware(...)
+        # - m2's handler calls m3.agent_middleware(...)
+        # - m3's handler eventually calls agent_invoke(...)
+        #
+        # We build the chain by iterating in reverse order.
+        # Each middleware wraps the previously constructed handler,
+        # so the first middleware in the list becomes the outermost one.
+
+        invoke = agent_invoke
+        for middleware in reversed(self._middleware):
+
+            def make_next(
+                m: AgentMiddleware, h: AgentMiddlewareHandler
+            ) -> AgentMiddlewareHandler:
+                async def next(r: AgentRequest) -> AgentResponse[Any | None]:
+                    return await m.agent_middleware(r, h)
+
+                return next
+
+            invoke = make_next(middleware, invoke)
+
+        return invoke
+
     @override
     async def invoke(self, messages: list[BaseMessage]) -> AgentResponse[OutputT]:
-        langchain_msgs = [_map_message_to_langchain(m) for m in messages]
+        async def invoke_agent(req: AgentRequest) -> AgentResponse[Any | None]:
+            langchain_msgs = [_map_message_to_langchain(m) for m in req.messages]
 
-        # call the langchain agent
-        result = await self._agent.ainvoke(
-            {"messages": langchain_msgs},
-            config=self._config,
-        )
+            # call the langchain agent
+            result = await self._agent.ainvoke(
+                {"messages": langchain_msgs},
+                config=self._config,
+            )
+
+            sdk_msgs = [_map_message_from_langchain(m) for m in result["messages"]]
+
+            # NOTE: Agent responses will always conform to output schema. Verifying
+            # if an LLM made any mistakes or not is _always_ up to the developer.
+
+            assert (
+                self._output_schema is None
+                or type(result["structured_response"]) is self._output_schema
+            )
+
+            if self._output_schema:
+                return AgentResponse(
+                    structured_output=result["structured_response"],
+                    messages=sdk_msgs,
+                )
+            else:
+                return AgentResponse(structured_output=None, messages=sdk_msgs)
 
-        sdk_msgs = [_map_message_from_langchain(m) for m in result["messages"]]
+        result = await self._with_agent_middleware(invoke_agent)(
+            AgentRequest(
+                messages=messages,
+            )
+        )
 
-        # NOTE: Agent responses will always conform to output schema. Verifying
-        # if an LLM made any mistakes or not is _always_ up to the developer.
         if self._output_schema:
-            return AgentResponse(
-                structured_output=result["structured_response"],
-                messages=sdk_msgs,
+            if result.structured_output is None:
+                raise AssertionError("Agent middleware discarded a structured output")
+
+            if type(result.structured_output) is not self._output_schema:
+                raise AssertionError(
+                    f"Agent middleware returned an invalid structured_output type: {type(result.structured_output)}, want: {self._output_schema}"
+                )
+
+            return AgentResponse[OutputT](
+                messages=result.messages,
+                structured_output=result.structured_output,
             )
+        else:
+            if result.structured_output is not None:
+                raise AssertionError(
+                    "Agent middleware unexpectedly included a structured output"
+                )
 
-        # HACK: This let's us put None in the structured_output field. It also shows
-        # None as the field type if no `output_schema`was provided to the Agent class.
-        return AgentResponse(structured_output=cast(OutputT, None), messages=sdk_msgs)
+            return AgentResponse[OutputT](
+                messages=result.messages,
+                # HACK: This let's us put None in the structured_output field. It also shows
+                # None as the field type if no `output_schema`was provided to the Agent class.
+                structured_output=cast(OutputT, None),
+            )
 
 
 @final
@@ -229,7 +304,8 @@ async def create_agent(
             model=model_impl,
             tools=tools,
             output_schema=agent.output_schema,
-            middleware=middleware,
+            lcmiddleware=middleware,
+            middleware=agent.middleware,
         )
 
 

splunklib/ai/messages.py

@@ -35,6 +35,7 @@ class SubagentCall:
     args: dict[str, Any]
     id: str | None  # TODO: can be None?
 
+
 @dataclass(frozen=True)
 class BaseMessage:
     role: str = ""

splunklib/ai/middleware.py

@@ -19,6 +19,7 @@
 from splunklib.ai.messages import (
     AIMessage,
     AgentResponse,
+    BaseMessage,
     SubagentCall,
     ToolCall,
 )
@@ -75,6 +76,14 @@ class ModelRequest:
 ModelMiddlewareHandler = Callable[[ModelRequest], Awaitable[AIMessage]]
 
 
+@dataclass
+class AgentRequest:
+    messages: list[BaseMessage]
+
+
+AgentMiddlewareHandler = Callable[[AgentRequest], Awaitable[AgentResponse[Any | None]]]
+
+
 class AgentMiddleware:
     async def tool_middleware(
         self,
@@ -103,6 +112,15 @@ async def model_middleware(
 
         return await handler(request)
 
+    async def agent_middleware(
+        self,
+        request: AgentRequest,
+        handler: AgentMiddlewareHandler,
+    ) -> AgentResponse[Any | None]:
+        """Executed in between invoke"""
+
+        return await handler(request)
+
 
 def tool_middleware(
     func: Callable[[ToolRequest, ToolMiddlewareHandler], Awaitable[ToolResponse]],
@@ -149,3 +167,20 @@ async def model_middleware(
             return await func(request, handler)
 
     return _CustomMiddleware()
+
+
+def agent_middleware(
+    func: Callable[
+        [AgentRequest, AgentMiddlewareHandler], Awaitable[AgentResponse[Any | None]]
+    ],
+) -> AgentMiddleware:
+    class _CustomMiddleware(AgentMiddleware):
+        @override
+        async def agent_middleware(
+            self,
+            request: AgentRequest,
+            handler: AgentMiddlewareHandler,
+        ) -> AgentResponse[Any | None]:
+            return await func(request, handler)
+
+    return _CustomMiddleware()