Implement Hooks as middlewares (#80)

This change:

- Makes the hook decorator return an AgentMiddleware
- Changes the hook function input params to match Middlewares
- Removes the underlying code, that implements hooks, as these are now
  middlewares.
- Removes use of LC middlewares in debug logging

Author: mateusz834

.basedpyright/baseline.json

@@ -30354,24 +30354,6 @@
                 }
             }
         ],
-        "./tests/integration/ai/test_hooks.py": [
-            {
-                "code": "reportOptionalMemberAccess",
-                "range": {
-                    "startColumn": 26,
-                    "endColumn": 30,
-                    "lineCount": 1
-                }
-            },
-            {
-                "code": "reportOptionalMemberAccess",
-                "range": {
-                    "startColumn": 26,
-                    "endColumn": 30,
-                    "lineCount": 1
-                }
-            }
-        ],
         "./tests/integration/ai/test_registry.py": [
             {
                 "code": "reportUnknownArgumentType",

splunklib/ai/README.md

@@ -592,12 +592,15 @@ They differ by the point in the execution flow where they are invoked:
 - before_agent: once per agent invocation, before any model calls
 - after_agent: once per agent invocation, after all model calls
 
+Hooks implement the same interface as middlewares, which allows them to be supplied
+directly as middleware instances in the Agent constructor.
+
 Example hook that logs token usage after each model call:
 
 ```py
 from splunklib.ai import Agent, OpenAIModel
 from splunklib.ai.hooks import after_model
-from splunklib.ai.middleware import AgentState
+from splunklib.ai.middleware import ModelResponse
 from splunklib.client import connect
 
 import logging
@@ -608,42 +611,15 @@ model = OpenAIModel(...)
 service = connect(...)
 
 @after_model
-def log_token_usage(state: AgentState) -> None:
-    logger.debug(f"Model used {state.token_count} tokens up to this point")
-
-
-async with Agent(
-    model=model,
-    service=service,
-    system_prompt="..." ,
-    hooks=[log_token_usage],
-) as agent: ...
-```
+def log_model_response(req: ModelResponse) -> None:
+    logger.debug(f"Model response {req.message.content}")
 
-The same hook can be defined as a class. It needs to provide the type and name attributes, and implement the `__call__` method:
-
-```py
-from typing import final, override
-from splunklib.ai.hooks import AgentHook
-from splunklib.ai.middleware import AgentState
-import logging
-
-logger = logging.getLogger(__name__)
-
-@final
-class LoggingHook(AgentHook):
-    type = "before_model"
-    name = "test_hook"
-
-    @override
-    def __call__(self, state: AgentState) -> None:
-        logger.debug(f"Model used {state.token_count} tokens up to this point")
 
 async with Agent(
     model=model,
     service=service,
     system_prompt="..." ,
-    hooks=[LoggingHook()],
+    middleware=[log_model_response],
 ) as agent: ...
 ```
 
@@ -652,25 +628,25 @@ The logic of the hook can be more advanced and include multiple conditions, for
 
 ```py
 from splunklib.ai import Agent, OpenAIModel
-from splunklib.ai.hooks import before_model, AgentHook
-from splunklib.ai.middleware import AgentState
+from splunklib.ai.hooks import before_model
+from splunklib.ai.middleware import AgentMiddleware, ModelRequest
 from time import monotonic
 
-def timeout_or_token_limit(seconds_limit: float, token_limit: float) -> AgentHook:
+def timeout_or_token_limit(seconds_limit: float, token_limit: float) -> AgentMiddleware:
     now = monotonic()
     timeout = now + seconds_limit
 
     @before_model
-    def _limit_hook(state: AgentState) -> None:
-        if state.token_count > token_limit or monotonic() >= timeout:
+    def _limit_hook(req: ModelRequest) -> None:
+        if req.state.token_count > token_limit or monotonic() >= timeout:
             raise Exception("Stopping Agentic Loop")
 
     return _limit_hook
 
 
 async with Agent(
     ...,
-    hooks=[timeout_or_token_limit(seconds_limit=10.0, token_limit=10000)],
+    middleware=[timeout_or_token_limit(seconds_limit=10.0, token_limit=10000)],
 ) as agent: ...
 ```
 

splunklib/ai/agent.py

@@ -23,7 +23,6 @@
 from splunklib.ai.base_agent import BaseAgent
 from splunklib.ai.core.backend import AgentImpl
 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
@@ -94,11 +93,6 @@ class Agent(BaseAgent[OutputT]):
             used as a *subagent*. The supervisor agent uses this schema to
             understand how to call the subagent and how to format its inputs.
 
-        hooks:
-            Optional sequence of `AgentHook`. Hooks are user-defined callback
-            functions that can be registered to execute at specific points
-            during the agent's operation.
-
         name:
             Name of the agent when used as a subagent. This is
             surfaced to the supervisor and used to decide whether this agent
@@ -130,7 +124,6 @@ def __init__(
         agents: Sequence[BaseAgent[BaseModel | None]] | None = None,
         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
@@ -144,7 +137,6 @@ def __init__(
             agents=agents,
             input_schema=input_schema,
             output_schema=output_schema,
-            hooks=hooks,
             middleware=middleware,
             logger=logger,
         )

splunklib/ai/base_agent.py

@@ -20,7 +20,6 @@
 
 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
@@ -36,7 +35,6 @@ class BaseAgent(Generic[OutputT], ABC):
     _description: str = ""
     _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
@@ -51,7 +49,6 @@ def __init__(
         agents: Sequence["BaseAgent[BaseModel | None]"] | None = None,
         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:
@@ -63,7 +60,6 @@ def __init__(
         self._agents = tuple(agents) if agents else ()
         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
 
@@ -112,10 +108,6 @@ def input_schema(self) -> type[BaseModel] | None:
     def output_schema(self) -> type[OutputT] | None:
         return self._output_schema
 
-    @property
-    def hooks(self) -> Sequence[AgentHook] | None:
-        return self._hooks
-
     @property
     def middleware(self) -> Sequence[AgentMiddleware] | None:
         return self._middleware