Support async hooks (#65)

Author: mateusz834

splunklib/ai/engines/langchain.py

@@ -12,6 +12,7 @@
 # License for the specific language governing permissions and limitations
 # under the License.
 
+from inspect import isawaitable
 import logging
 import uuid
 from collections.abc import Sequence
@@ -500,7 +501,9 @@ def _convert_hook_to_middleware(
         case _:
             raise AssertionError(f"Unsupported middleware type: {hook.type}")
 
-    def _middleware(state: LC_AgentState, runtime: Runtime) -> dict[str, Any] | None:
+    async def _middleware(
+        state: LC_AgentState, runtime: Runtime
+    ) -> dict[str, Any] | None:
         # NOTE: We're converting the langchain AgentState into the SDK AgentState
         # on each middleware call.
         # We're converting all the messages back to the SDK format and counting the
@@ -516,7 +519,10 @@ def _middleware(state: LC_AgentState, runtime: Runtime) -> dict[str, Any] | None
         if logger:
             logger.debug(f"Executing {hook.type} hook {hook.name}")
 
-        hook(sdk_state)
+        res = hook(sdk_state)
+        if isawaitable(res):
+            await res
+        return None
 
     return wrapper(_middleware)
 

splunklib/ai/hooks.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 from time import monotonic
-from typing import Any, Callable, Literal, Protocol, final, override
+from typing import Any, Awaitable, Callable, Literal, Protocol, final, override
 
 from splunklib.ai.messages import AgentResponse
 
@@ -34,7 +34,7 @@ class AgentHook(Protocol):
     # Name of the middleware must be unique
     name: str
 
-    def __call__(self, state: AgentState) -> None:
+    def __call__(self, state: AgentState) -> None | Awaitable[None]:
         """Called at specific points during the agent execution, depending on the hook type."""
 
 
@@ -65,7 +65,7 @@ def __init__(self, timeout_seconds: float) -> None:
 
 def _create_hook(
     type: HookType,
-    func: Callable[[AgentState], None],
+    func: Callable[[AgentState], None | Awaitable[None]],
     name: str | None = None,
 ) -> AgentHook:
     mw_name = name or func.__name__
@@ -77,31 +77,31 @@ class CustomHook(AgentHook):
         name = mw_name
 
         @override
-        def __call__(self, state: AgentState) -> None:
+        def __call__(self, state: AgentState) -> None | Awaitable[None]:
             return func(state)
 
     return CustomHook()
 
 
-def before_model(func: Callable[[AgentState], None]) -> AgentHook:
+def before_model(func: Callable[[AgentState], None | Awaitable[None]]) -> AgentHook:
     """This hook is called before each model call."""
 
     return _create_hook("before_model", func)
 
 
-def after_model(func: Callable[[AgentState], None]) -> AgentHook:
+def after_model(func: Callable[[AgentState], None | Awaitable[None]]) -> AgentHook:
     """This hook is called after each model call."""
 
     return _create_hook("after_model", func)
 
 
-def before_agent(func: Callable[[AgentState], None]) -> AgentHook:
+def before_agent(func: Callable[[AgentState], None | Awaitable[None]]) -> AgentHook:
     """This hook is called once per agent invocation. Before any model calls."""
 
     return _create_hook("before_agent", func)
 
 
-def after_agent(func: Callable[[AgentState], None]) -> AgentHook:
+def after_agent(func: Callable[[AgentState], None | Awaitable[None]]) -> AgentHook:
     """This hook is called once per agent invocation. After all model calls."""
 
     return _create_hook("after_agent", func)

tests/integration/ai/test_hooks.py

@@ -57,20 +57,35 @@ async def test_agent_hooks_duplicated(self):
     async def test_agent_hook(self):
         pytest.importorskip("langchain_openai")
 
+        hook_calls = 0
+
         @final
         class TestHook(AgentHook):
             type = "before_model"
-            name = "test_hook"
+            name = "test_async_hook"
 
             @override
             def __call__(self, state: AgentState) -> None:
+                nonlocal hook_calls
+                hook_calls += 1
+                assert len(state.response.messages) == 1
+
+        @final
+        class TestAsyncHook(AgentHook):
+            type = "before_model"
+            name = "test_hook"
+
+            @override
+            async def __call__(self, state: AgentState) -> None:
+                nonlocal hook_calls
+                hook_calls += 1
                 assert len(state.response.messages) == 1
 
         async with Agent(
             model=(await self.model()),
             system_prompt="Your name is stefan",
             service=self.service,
-            hooks=[TestHook()],
+            hooks=[TestHook(), TestAsyncHook()],
         ) as agent:
             result = await agent.invoke(
                 [
@@ -82,6 +97,7 @@ def __call__(self, state: AgentState) -> None:
 
             response = result.messages[-1].content.strip().lower().replace(".", "")
             assert "stefan" == response
+            assert hook_calls == 2
 
     @pytest.mark.asyncio
     async def test_agent_hook_decorator(self):
@@ -96,18 +112,37 @@ def test_hook_before(state: AgentState) -> None:
 
             assert len(state.response.messages) == 1
 
+        @before_model
+        async def test_async_hook_before(state: AgentState) -> None:
+            nonlocal hook_calls
+            hook_calls += 1
+
+            assert len(state.response.messages) == 1
+
         @after_model
         def test_hook_after(state: AgentState) -> None:
             nonlocal hook_calls
             hook_calls += 1
 
             assert len(state.response.messages) == 2
 
+        @after_model
+        async def test_async_hook_after(state: AgentState) -> None:
+            nonlocal hook_calls
+            hook_calls += 1
+
+            assert len(state.response.messages) == 2
+
         async with Agent(
             model=(await self.model()),
             system_prompt="Your name is stefan",
             service=self.service,
-            hooks=[test_hook_before, test_hook_after],
+            hooks=[
+                test_hook_before,
+                test_async_hook_before,
+                test_hook_after,
+                test_async_hook_after,
+            ],
         ) as agent:
             result = await agent.invoke(
                 [
@@ -119,7 +154,7 @@ def test_hook_after(state: AgentState) -> None:
 
             response = result.messages[-1].content.strip().lower().replace(".", "")
             assert "stefan" == response
-            assert hook_calls == 2
+            assert hook_calls == 4
 
     @pytest.mark.asyncio
     async def test_agent_hook_agent(self):
@@ -137,8 +172,24 @@ def before_agent_hook(state: AgentState) -> None:
 
             assert len(state.response.messages) == 1
 
+        @before_agent
+        async def before_async_agent_hook(state: AgentState) -> None:
+            nonlocal hook_calls
+            hook_calls += 1
+
+            assert len(state.response.messages) == 1
+
         @after_agent
-        def after_agent_hook(state: AgentState) -> None:
+        async def after_agent_hook(state: AgentState) -> None:
+            nonlocal hook_calls
+            hook_calls += 1
+
+            person = state.response.structured_output
+            assert person.name.lower() == "stefan"
+            assert len(state.response.messages) == 2
+
+        @after_agent
+        async def after_async_agent_hook(state: AgentState) -> None:
             nonlocal hook_calls
             hook_calls += 1
 
@@ -150,7 +201,12 @@ def after_agent_hook(state: AgentState) -> None:
             model=(await self.model()),
             system_prompt="Your name is stefan",
             service=self.service,
-            hooks=[before_agent_hook, after_agent_hook],
+            hooks=[
+                before_agent_hook,
+                before_async_agent_hook,
+                after_agent_hook,
+                after_async_agent_hook,
+            ],
             output_schema=Person,
         ) as agent:
             result = await agent.invoke(
@@ -163,7 +219,7 @@ def after_agent_hook(state: AgentState) -> None:
 
             response = result.messages[-1].content.strip().lower().replace(".", "")
             assert '{"name":"stefan"}' == response
-            assert hook_calls == 2
+            assert hook_calls == 4
 
     @pytest.mark.asyncio
     async def test_agent_loop_stop_conditions_token_limit(self):