fix(tool_registry): Support plain dict as return types of local tools (#141)

- Add reproducer tool that returns a dict
- Fix the bug

Author: szykol

splunklib/ai/registry.py

@@ -17,7 +17,7 @@
 import logging
 import string
 from collections.abc import Callable, Sequence
-from dataclasses import asdict, dataclass
+from dataclasses import asdict, dataclass, is_dataclass
 from logging import Logger
 from typing import (
     Any,
@@ -317,8 +317,14 @@ async def _call_tool(
                 if self._tools_wrapped_result.get(name):
                     res = _WrappedResult(res)
 
+                if is_dataclass(res) and not isinstance(res, type):
+                    res = asdict(res)
+
+                if not isinstance(res, dict):
+                    raise AssertionError("invalid type of tool response")
+
                 return types.CallToolResult(
-                    structuredContent=asdict(res),
+                    structuredContent=res,  # pyright: ignore[reportUnknownArgumentType]
                     content=[],
                 )
         except BaseExceptionGroup as e:
@@ -354,6 +360,7 @@ def _input_schema(self, func: Callable[_P, _R]) -> dict[str, Any]:
 
         return input_schema
 
+    # TODO: figure out how to handle custom classes as output type
     def _output_schema(self, func: Callable[_P, _R]) -> tuple[dict[str, Any], bool]:
         """
         Generates a output schema for the provided func, if necessary wraps the

tests/integration/ai/test_agent_mcp_tools.py

@@ -25,18 +25,27 @@
 from splunklib.ai import Agent
 from splunklib.ai.engines.langchain import LOCAL_TOOL_PREFIX
 from splunklib.ai.messages import (
+    AIMessage,
     HumanMessage,
+    ToolCall,
     ToolFailureResult,
     ToolMessage,
     ToolResult,
 )
+from splunklib.ai.middleware import (
+    ModelMiddlewareHandler,
+    ModelRequest,
+    ModelResponse,
+    model_middleware,
+)
 from splunklib.ai.tool_settings import (
     LocalToolSettings,
     RemoteToolSettings,
     ToolAllowlist,
     ToolSettings,
 )
 from splunklib.ai.tools import (
+    ToolType,
     _get_splunk_username,  # pyright: ignore[reportPrivateUsage]
     locate_app,
 )
@@ -589,6 +598,67 @@ async def lifespan(app: Starlette) -> AsyncGenerator[None, Any]:
                 response = result.final_message.content
                 assert "31.5" in response, "Invalid LLM response"
 
+    @patch(
+        "splunklib.ai.agent._testing_local_tools_path",
+        os.path.join(os.path.dirname(__file__), "testdata", "temperature_as_dict.py"),
+    )
+    @patch("splunklib.ai.agent._testing_app_id", "app_id")
+    @pytest.mark.asyncio
+    async def test_supports_plain_dicts_as_tool_outputs(self) -> None:
+        """Regression test for DVPL-13022"""
+        pytest.importorskip("langchain_openai")
+
+        messages: list[AIMessage] = [
+            AIMessage(
+                content="",
+                calls=[
+                    ToolCall(
+                        name="temperature",
+                        args={"city": "Krakow"},
+                        id="call_hSdIJSuUZOh2IiBsqfrzhA7d",
+                        type=ToolType.LOCAL,
+                    )
+                ],
+            ),
+            AIMessage(content="The temperature in Krakow is 22°C.", calls=[]),
+        ]
+
+        responses = (m for m in messages)
+
+        @model_middleware
+        async def middleware(
+            req: ModelRequest, handler: ModelMiddlewareHandler
+        ) -> ModelResponse:
+            return ModelResponse(message=next(responses))
+
+        async with Agent(
+            model=(await self.model()),
+            system_prompt="You must use the available tools to perform requested operations",
+            service=self.service,
+            tool_settings=ToolSettings(local=True, remote=None),
+            middleware=[middleware],
+        ) as agent:
+            result = await agent.invoke(
+                [
+                    HumanMessage(
+                        content=(
+                            "What is the weather like today in Krakow? Use the provided tools to check the temperature."
+                            + "Return a short response, containing the tool response."
+                        ),
+                    )
+                ]
+            )
+
+            tool_message = next(
+                filter(lambda m: m.role == "tool", result.messages), None
+            )
+            assert isinstance(tool_message, ToolMessage), "Invalid tool message"
+            assert tool_message, "No tool message found in response"
+            assert tool_message.name == "temperature", "Invalid tool name"
+
+            response = result.final_message.content
+            assert "22" in response, "Invalid LLM response"
+
 
 class TestHandlingToolNameCollision(AITestCase):
     @patch(

tests/integration/ai/test_registry.py

@@ -118,6 +118,21 @@ async def test_tool_hello(self):
             self.assertEqual(res.structuredContent, {"result": "Hello Stefan"})
 
 
+class TestTemperatureAsDictRegistry(TestRegistryTestCase):
+    async def test_tool_temperature_returning_dict(self):
+        async with self.connect("temperature_as_dict.py") as session:
+            res = await session.call_tool(
+                "temperature",
+                arguments={"city": "Krakow"},
+                meta={"splunk": {"service": self.serialized_service.model_dump()}},
+            )
+            self.assertEqual(res.isError, False)
+            self.assertEqual(res.content, [])
+            self.assertEqual(
+                res.structuredContent, {"city": "Krakow", "temperature": 22}
+            )
+
+
 @dataclass
 class Log:
     level: LoggingLevel

tests/integration/ai/testdata/temperature_as_dict.py

@@ -0,0 +1,15 @@
+from typing import Any
+
+from splunklib.ai.registry import ToolContext, ToolRegistry
+
+registry = ToolRegistry()
+
+
+@registry.tool(name="temperature", tags=["read"])
+def temperature(city: str, _ctx: ToolContext) -> dict[str, Any]:
+    """A simple tool that returns a temperature for the city."""
+
+    return {"city": city, "temperature": 22}
+
+
+registry.run()

tests/unit/ai/test_registry_unit.py

@@ -66,6 +66,25 @@ def structured_tool() -> Output:
             "type": "object",
         }
 
+    def test_output_non_wrapped_dict(self) -> None:
+        r = ToolRegistry()
+
+        @r.tool()
+        def structured_tool() -> dict[str, Any]:
+            return {"some": "info"}
+
+        tool = r._tools[0]
+        assert tool.name == "structured_tool"
+        assert tool.inputSchema == {
+            "properties": {},
+            "type": "object",
+            "additionalProperties": False,
+        }
+        assert tool.outputSchema == {
+            "additionalProperties": True,
+            "type": "object",
+        }
+
     def test_output_wrapped(self) -> None:
         r = ToolRegistry()