Move AgentState to splunklib.ai.middleware (#79)

Author: mateusz834

.basedpyright/baseline.json

@@ -82,16 +82,6 @@
                 }
             }
         ],
-        "./splunklib/ai/hooks.py": [
-            {
-                "code": "reportDeprecated",
-                "range": {
-                    "startColumn": 24,
-                    "endColumn": 33,
-                    "lineCount": 1
-                }
-            }
-        ],
         "./splunklib/ai/model.py": [
             {
                 "code": "reportDeprecated",

splunklib/ai/README.md

@@ -562,6 +562,7 @@ 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.client import connect
 
 import logging
@@ -588,7 +589,8 @@ The same hook can be defined as a class. It needs to provide the type and name a
 
 ```py
 from typing import final, override
-from splunklib.ai.hooks import AgentHook, AgentState
+from splunklib.ai.hooks import AgentHook
+from splunklib.ai.middleware import AgentState
 import logging
 
 logger = logging.getLogger(__name__)
@@ -616,6 +618,7 @@ 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 time import monotonic
 
 def timeout_or_token_limit(seconds_limit: float, token_limit: float) -> AgentHook:

splunklib/ai/engines/langchain.py

@@ -62,7 +62,6 @@
 )
 from splunklib.ai.hooks import (
     AgentHook,
-    AgentState,
     FunctionHook,
     after_model as hook_after_model,
     before_model as hook_before_model,
@@ -80,6 +79,7 @@
     ToolMessage,
 )
 from splunklib.ai.middleware import (
+    AgentState,
     AgentMiddleware,
     ModelMiddlewareHandler,
     ModelRequest,

splunklib/ai/hooks.py

@@ -1,8 +1,8 @@
-from dataclasses import dataclass
+from collections.abc import Awaitable, Callable
 from time import monotonic
-from typing import Any, Awaitable, Callable, Literal, Protocol, final, override
+from typing import Literal, Protocol, final, override
 
-from splunklib.ai.messages import AgentResponse
+from splunklib.ai.middleware import AgentState
 
 # Hook type decides when the hook is called during agent execution.
 # before_model: before each model call
@@ -12,18 +12,6 @@
 HookType = Literal["before_model", "after_model", "before_agent", "after_agent"]
 
 
-@dataclass(frozen=True)
-class AgentState:
-    """AgentState is passed to each hook and contains information about the current state of the agent execution."""
-
-    # holds messages exchanged so far in the conversation
-    response: AgentResponse[Any | None]
-    # steps taken so far in the conversation
-    total_steps: int
-    # tokens used so far in the conversation
-    token_count: float
-
-
 class AgentHook(Protocol):
     """AgentHook is a callable that can be registered to be called at specific points during the agent execution.
 

splunklib/ai/middleware.py

@@ -14,16 +14,28 @@
 
 from collections.abc import Awaitable, Callable
 from dataclasses import dataclass
-from typing import Literal, override
+from typing import Any, Literal, override
 
-from splunklib.ai.hooks import AgentState
 from splunklib.ai.messages import (
     AIMessage,
+    AgentResponse,
     SubagentCall,
     ToolCall,
 )
 
 
+@dataclass(frozen=True)
+class AgentState:
+    """AgentState is passed to middleware and contains information about the current state of the agent execution."""
+
+    # holds messages exchanged so far in the conversation
+    response: AgentResponse[Any | None]
+    # steps taken so far in the conversation
+    total_steps: int
+    # tokens used so far in the conversation
+    token_count: float
+
+
 @dataclass
 class ToolRequest:
     call: ToolCall

tests/integration/ai/test_hooks.py

@@ -21,7 +21,6 @@
 from splunklib.ai import Agent
 from splunklib.ai.hooks import (
     AgentHook,
-    AgentState,
     StepsLimitExceededException,
     TimeoutExceededException,
     TokenLimitExceededException,
@@ -34,6 +33,7 @@
     token_limit,
 )
 from splunklib.ai.messages import HumanMessage
+from splunklib.ai.middleware import AgentState
 from tests.ai_testlib import AITestCase