Add hooks (#45)

Hooks help with inserting callbacks into different moments
of the Agent Loop execution.
This can be used to log/audit/debug and to stop the
Agentic Loop early on.

Author: szykol

splunklib/ai/README.md

@@ -365,12 +365,105 @@ async with Agent(
 
 **Note**: Currently input schemas can only be used by subagents, not by regular agents.
 
-## Loop Stop Conditions
+## Hooks
+
+Hooks are user-defined callback functions that can be registered to execute at specific points
+during the agent's operation. Hooks allow developers to add custom behavior, logging and monitoring
+or implement custom stopping conditions for the agent loop without modifying the core agent logic.
+
+There are several types of hooks available.
+They differ by the point in the execution flow where they are invoked:
+
+- before_model: before each model call
+- after_model: after each model call
+- before_agent: once per agent invocation, before any model calls
+- after_agent: once per agent invocation, after all model calls
+
+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.client import connect
+
+import logging
+
+logger = logging.getLogger(__name__)
+
+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: ...
+```
+
+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, 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()],
+) as agent: ...
+```
+
+The hooks can stop the Agentic Loop under custom conditions by raising exceptions.
+The logic of the hook can be more advanced and include multiple conditions, for example, based on both token usage and execution time:
+
+```py
+from splunklib.ai import Agent, OpenAIModel
+from splunklib.ai.hooks import before_model, AgentHook
+from time import monotonic
+
+def timeout_or_token_limit(seconds_limit: float, token_limit: float) -> AgentHook:
+    now = monotonic()
+    timeout = now + seconds_limit
+
+    @before_model
+    def _limit_hook(state: AgentState) -> None:
+        if 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)],
+) as agent: ...
+```
+
+### Predefined hooks for loop stopping conditions
 
 To prevent excessive token usage or runaway execution, an Agent can be constrained
-using loop stop conditions.
+using predefined hooks.
 
-Stop conditions allow you to automatically terminate the agent loop when one or more
+Those hooks allow you to automatically terminate the agent loop when one or more
 limits are reached, such as:
 
 - Maximum number of generated tokens
@@ -379,7 +472,7 @@ limits are reached, such as:
 
 ```py
 from splunklib.ai import Agent, OpenAIModel
-from splunklib.ai.stop_conditions import StopConditions
+from splunklib.ai.hooks import token_limit, step_limit, timeout_limit
 from splunklib.client import connect
 
 model = OpenAIModel(...)
@@ -389,11 +482,11 @@ async with Agent(
         model=model,
         service=service,
         system_prompt="..." ,
-        loop_stop_conditions=StopConditions(
-            token_limit = 10000,
-            steps_limit = 25,
-            timeout_seconds = 10.5,
-        ),
+        hooks=[
+            token_limit(10000),
+            step_limit(25),
+            timeout_limit(10.5),
+        ],
     ) as agent: ...
 ```
 

splunklib/ai/agent.py

@@ -22,9 +22,9 @@
 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.model import PredefinedModel
-from splunklib.ai.stop_conditions import StopConditions
 from splunklib.ai.tool_filtering import ToolFilters, filter_tools
 from splunklib.ai.tools import (
     Tool,
@@ -88,11 +88,10 @@ 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.
 
-        loop_stop_conditions:
-            Optional `StopConditions` instance defining automatic termination.
-            If any limit is exceeded, the corresponding exception
-            (`TokenLimitExceededException`, `StepsLimitExceededException`,
-            or `TimeoutExceededException`) is raised.
+        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
@@ -122,7 +121,7 @@ def __init__(
         agents: Sequence[BaseAgent[BaseModel | None]] | None = None,
         output_schema: type[OutputT] | None = None,
         input_schema: type[BaseModel] | None = None,  # Only used by Subgents
-        loop_stop_conditions: StopConditions | None = None,
+        hooks: Sequence[AgentHook] | None = None,
         name: str = "",  # Only used by Subgents
         description: str = "",  # Only used by Subagents
     ) -> None:
@@ -134,9 +133,12 @@ def __init__(
             agents=agents,
             input_schema=input_schema,
             output_schema=output_schema,
-            loop_stop_conditions=loop_stop_conditions,
+            hooks=hooks,
         )
 
+        if duplicate_hook_names := _find_duplicate_hook_names(self.hooks):
+            raise ValueError(f"Duplicate hook names found: {duplicate_hook_names!r}")
+
         self._use_mcp_tools = use_mcp_tools
         self._tool_filters = tool_filters
         self._service = service
@@ -181,3 +183,19 @@ async def _load_tools_from_mcp(
         return filter_tools(mcp_tools, filters)
 
     return mcp_tools
+
+
+def _find_duplicate_hook_names(hooks: Sequence[AgentHook] | None) -> set[str]:
+    seen: set[str] = set()
+    duplicates: set[str] = set()
+
+    if not hooks:
+        return set()
+
+    for hook in hooks:
+        if hook.name in seen:
+            duplicates.add(hook.name)
+        else:
+            seen.add(hook.name)
+
+    return duplicates

splunklib/ai/base_agent.py

@@ -19,9 +19,9 @@
 
 from pydantic import BaseModel
 
+from splunklib.ai.hooks import AgentHook
 from splunklib.ai.messages import AgentResponse, BaseMessage, OutputT
 from splunklib.ai.model import PredefinedModel
-from splunklib.ai.stop_conditions import StopConditions
 from splunklib.ai.tools import Tool
 
 
@@ -34,7 +34,7 @@ class BaseAgent(Generic[OutputT], ABC):
     _description: str = ""
     _input_schema: type[BaseModel] | None = None
     _output_schema: type[OutputT] | None = None
-    _loop_stop_conditions: StopConditions | None = None
+    _hooks: Sequence[AgentHook] | None = None
 
     def __init__(
         self,
@@ -46,7 +46,7 @@ def __init__(
         agents: Sequence["BaseAgent[BaseModel | None]"] | None = None,
         input_schema: type[BaseModel] | None = None,
         output_schema: type[OutputT] | None = None,
-        loop_stop_conditions: StopConditions | None = None,
+        hooks: Sequence[AgentHook] | None = None,
     ) -> None:
         self._system_prompt = system_prompt
         self._model = model
@@ -56,7 +56,7 @@ def __init__(
         self._agents = tuple(agents) if agents else ()
         self._input_schema = input_schema
         self._output_schema = output_schema
-        self._loop_stop_conditions = loop_stop_conditions
+        self._hooks = tuple(hooks) if hooks else ()
 
     @abstractmethod
     async def invoke(self, messages: list[BaseMessage]) -> AgentResponse[OutputT]: ...
@@ -94,5 +94,5 @@ def output_schema(self) -> type[OutputT] | None:
         return self._output_schema
 
     @property
-    def loop_stop_conditions(self) -> StopConditions | None:
-        return self._loop_stop_conditions
+    def hooks(self) -> Sequence[AgentHook] | None:
+        return self._hooks