Support async local tools (#48)

Author: mateusz834

splunklib/ai/registry.py

@@ -102,12 +102,14 @@ async def _() -> list[types.Tool]:
 
         @self._server.call_tool(validate_input=True)
         async def _(name: str, arguments: dict[str, Any]) -> types.CallToolResult:
-            return self._call_tool(name, arguments)
+            return await self._call_tool(name, arguments)
 
     def _list_tools(self) -> list[types.Tool]:
         return self._tools
 
-    def _call_tool(self, name: str, arguments: dict[str, Any]) -> types.CallToolResult:
+    async def _call_tool(
+        self, name: str, arguments: dict[str, Any]
+    ) -> types.CallToolResult:
         func = self._tools_func.get(name)
         if func is None:
             raise ValueError(f"Tool {name} does not exist")
@@ -129,6 +131,11 @@ def _call_tool(self, name: str, arguments: dict[str, Any]) -> types.CallToolResu
 
         res = func(**arguments)
 
+        # In case func was an async function, await the returned coroutine.
+        # If not then we already have the result.
+        if inspect.isawaitable(res):
+            res = await res
+
         if self._tools_wrapped_result.get(name):
             res = _WrappedResult(res)
 

tests/integration/ai/test_registry.py

@@ -112,6 +112,24 @@ async def test_missing_meta_params(self):
             self.assertEqual(res.structuredContent, None)
 
 
+class TestAsyncToolRegistry(TestRegistryTestCase):
+    async def test_tool_hello(self):
+        async with self.connect("async_tool.py") as session:
+            res = await session.call_tool(
+                "hello",
+                arguments={"name": "Stefan"},
+                meta={
+                    "splunk": {
+                        "management_token": self.get_splunk_token(),
+                        "management_url": self.splunk_url,
+                    }
+                },
+            )
+            self.assertEqual(res.isError, False)
+            self.assertEqual(res.content, [])
+            self.assertEqual(res.structuredContent, {"result": "Hello Stefan"})
+
+
 if __name__ == "__main__":
     import unittest
 

tests/integration/ai/testdata/async_tool.py

@@ -0,0 +1,11 @@
+from splunklib.ai.registry import ToolRegistry
+
+registry = ToolRegistry()
+
+
+@registry.tool()
+async def hello(name: str) -> str:
+    return f"Hello {name}"
+
+
+registry.run()

tests/unit/ai/test_registry_unit.py

@@ -331,6 +331,29 @@ def fancy_tool(foo: int | None, bar: Data, baz: int = -1) -> Data:
             },
         )
 
+    def test_async_tool(self) -> None:
+        r = ToolRegistry()
+
+        @r.tool()
+        async def str_tool() -> str:
+            return ""
+
+        tool = r._tools[0]
+        self.assertEqual(tool.name, "str_tool")
+        self.assertEqual(
+            tool.inputSchema,
+            {"properties": {}, "type": "object", "additionalProperties": False},
+        )
+        self.assertEqual(
+            tool.outputSchema,
+            {
+                "properties": {"result": {"title": "Result", "type": "string"}},
+                "required": ["result"],
+                "title": "_WrappedResult",
+                "type": "object",
+            },
+        )
+
 
 class TestParams(unittest.TestCase):
     def test_description_param(self) -> None: