Reflect unstructured inputs in SubagentCall.args (#110)

Author: mateusz834

splunklib/ai/engines/langchain.py

@@ -146,6 +146,7 @@ def __init__(self, agent: BaseAgent[OutputT]) -> None:
         tools = _prepare_langchain_tools(agent.tools)
 
         system_prompt = agent.system_prompt
+        structured_subagents: list[str] = []
         if agent.agents:
             seen_names: set[str] = set()
             for subagent in agent.agents:
@@ -161,6 +162,9 @@ def __init__(self, agent: BaseAgent[OutputT]) -> None:
                 seen_names.add(subagent.name)
                 tools.append(tool)
 
+                if subagent.input_schema is not None:
+                    structured_subagents.append(subagent.name)
+
                 system_prompt = AGENT_AS_TOOLS_PROMPT + "\n" + system_prompt
 
         before_user_middlewares, after_user_middlewares = _debugging_middleware(
@@ -211,7 +215,87 @@ async def awrap_tool_call(
 
                 return resp
 
+        class _SubagentArgumentPacker(LC_AgentMiddleware):
+            # For non-structured subagents, the SubagentCall.args field is an `str | dict[str, Any]`,
+            # to differentiate that we wrap the resulting args in an SubagentLCArgs.
+            #
+            # This middleware performs the corresponding pack/unpack at the two
+            # points in the LangChain call graph where raw args are needed/retreived.
+            #
+            # TODO: once we move middlewares into one LC middleware, we should move
+            # that piece of logic there (DVPL-12959).
+            @override
+            async def awrap_model_call(
+                self,
+                request: LC_ModelRequest,
+                handler: Callable[[LC_ModelRequest], Awaitable[LC_ModelCallResult]],
+            ) -> LC_ModelCallResult:
+                # Unpack existing messages.
+                messages: list[LC_AnyMessage] = []
+                for msg in request.messages:
+                    if isinstance(msg, LC_AIMessage):
+                        new_calls: list[LC_ToolCall] = []
+                        for call in msg.tool_calls:
+                            new_calls.append(self.unpack_tool_call(call))
+                        msg = msg.model_copy(update={"tool_calls": new_calls})
+                    messages.append(msg)
+
+                response = await handler(request.override(messages=messages))
+
+                ai_message = response
+                if isinstance(ai_message, LC_ExtendedModelResponse):
+                    ai_message = ai_message.model_response
+                if isinstance(ai_message, LC_ModelResponse):
+                    ai_message = next(
+                        (m for m in ai_message.result if isinstance(m, LC_AIMessage)),
+                        None,
+                    )
+                    assert ai_message, "AIMessage not found found in response"
+
+                # Pack new message.
+                for call in ai_message.tool_calls:
+                    if call["name"].startswith(AGENT_PREFIX):
+                        if (
+                            _denormalize_agent_name(call["name"])
+                            in structured_subagents
+                        ):
+                            args = SubagentLCArgs(call["args"])
+                        else:
+                            content: str = call["args"].get("content", "")
+                            args = SubagentLCArgs(content)
+                        call["args"] = asdict(args)
+
+                return response
+
+            # Unpack args, just before tool call.
+            @override
+            async def awrap_tool_call(
+                self,
+                request: LC_ToolCallRequest,
+                handler: Callable[
+                    [LC_ToolCallRequest], Awaitable[LC_ToolMessage | LC_Command[None]]
+                ],
+            ) -> LC_ToolMessage | LC_Command[None]:
+                return await handler(
+                    request.override(
+                        tool_call=self.unpack_tool_call(request.tool_call),
+                    )
+                )
+
+            def unpack_tool_call(self, call: LC_ToolCall) -> LC_ToolCall:
+                if call["name"].startswith(AGENT_PREFIX):
+                    unpacked_args = SubagentLCArgs(**call["args"]).args
+                    if isinstance(unpacked_args, str):
+                        unpacked_args = {"content": unpacked_args}
+                    return LC_ToolCall(
+                        id=call["id"],
+                        name=call["name"],
+                        args=unpacked_args,
+                    )
+                return call
+
         lc_middleware.append(_ToolFailureArtifact())
+        lc_middleware.append(_SubagentArgumentPacker())
 
         self._agent = create_agent(
             model=model_impl,
@@ -933,12 +1017,17 @@ async def _run(
     )
 
 
+@dataclass(frozen=True)
+class SubagentLCArgs:
+    args: str | dict[str, Any]
+
+
 def _map_tool_call_from_langchain(tool_call: LC_ToolCall) -> ToolCall | SubagentCall:
     name = tool_call["name"]
     if name.startswith(AGENT_PREFIX):
         return SubagentCall(
             name=_denormalize_agent_name(name),
-            args=tool_call["args"],
+            args=SubagentLCArgs(**tool_call["args"]).args,
             id=tool_call["id"],
         )
 
@@ -957,10 +1046,12 @@ def _map_tool_call_to_langchain(call: ToolCall | SubagentCall) -> LC_ToolCall:
     match call:
         case SubagentCall():
             name = _normalize_agent_name(call.name)
+            args = asdict(SubagentLCArgs(call.args))
         case ToolCall():
             name = _normalize_tool_name(call.name, call.type)
+            args = call.args
 
-    return LC_ToolCall(id=call.id, name=name, args=call.args)
+    return LC_ToolCall(id=call.id, name=name, args=args)
 
 
 def _map_message_from_langchain(message: LC_BaseMessage) -> BaseMessage:

splunklib/ai/messages.py

@@ -32,8 +32,7 @@ class ToolCall:
 @dataclass(frozen=True)
 class SubagentCall:
     name: str
-    # TODO: should be a str | dict[str, Any] for subagents without structured inputs
-    args: dict[str, Any]
+    args: str | dict[str, Any]
     id: str | None  # TODO: can be None?
 
 

tests/integration/ai/test_agent.py

@@ -16,7 +16,7 @@
 from pydantic import BaseModel, Field
 
 from splunklib.ai import Agent
-from splunklib.ai.messages import HumanMessage, SubagentMessage
+from splunklib.ai.messages import AIMessage, HumanMessage, SubagentCall, SubagentMessage
 from tests.ai_testlib import AITestCase
 
 OPENAI_BASE_URL = "http://localhost:11434/v1"
@@ -175,7 +175,17 @@ class NicknameGeneratorInput(BaseModel):
                 ]
             )
 
-            response = result.final_message.content
+            first_ai_message = next(
+                m for m in result.messages if isinstance(m, AIMessage)
+            )
+            assert first_ai_message
+            assert len(first_ai_message.calls) == 1
+            assert isinstance(first_ai_message.calls[0], SubagentCall)
+            args = first_ai_message.calls[0].args
+            assert isinstance(args, dict)
+
+            # asserts that can create NicknameGeneratorInput from args
+            NicknameGeneratorInput(**args)
 
             subagent_message = next(
                 filter(lambda m: m.role == "subagent", result.messages), None
@@ -184,6 +194,8 @@ class NicknameGeneratorInput(BaseModel):
                 "Invalid subagent message"
             )
             assert subagent_message, "No subagent message found in response"
+
+            response = result.final_message.content
             assert "Chris-zilla" in response, "Agent did generate valid nickname"
 
     @pytest.mark.asyncio
@@ -217,6 +229,15 @@ async def test_subagent_without_input_schema(self):
                 ]
             )
 
+            first_ai_message = next(
+                m for m in result.messages if isinstance(m, AIMessage)
+            )
+            assert first_ai_message
+            assert len(first_ai_message.calls) == 1
+            assert isinstance(first_ai_message.calls[0], SubagentCall)
+            assert isinstance(first_ai_message.calls[0].args, str)
+            assert first_ai_message.calls[0].args.lower() == "chris"
+
             response = result.final_message.content
             assert "Chris-zilla" in response, "Agent did generate valid nickname"
 

tests/unit/ai/engine/test_langchain_backend.py

@@ -58,7 +58,7 @@ def test_map_message_from_langchain_ai_with_tool_calls(self) -> None:
 
     def test_map_message_from_langchain_ai_with_agent_call(self) -> None:
         tool_call = LC_ToolCall(
-            name=f"{lc.AGENT_PREFIX}assistant", args={"q": "test"}, id="tc-2"
+            name=f"{lc.AGENT_PREFIX}assistant", args={"args": {"q": "test"}}, id="tc-2"
         )
         message = LC_AIMessage(content="done", tool_calls=[tool_call])
         mapped = lc._map_message_from_langchain(message)
@@ -75,7 +75,7 @@ def test_map_message_from_langchain_ai_with_agent_call(self) -> None:
     def test_map_message_from_langchain_ai_with_mixed_calls(self) -> None:
         tool_call = LC_ToolCall(name="lookup", args={"q": "test"}, id="tc-1")
         agent_call = LC_ToolCall(
-            name=f"{lc.AGENT_PREFIX}assistant", args={"q": "test"}, id="tc-2"
+            name=f"{lc.AGENT_PREFIX}assistant", args={"args": {"q": "test"}}, id="tc-2"
         )
         message = LC_AIMessage(content="done", tool_calls=[tool_call, agent_call])
 
@@ -155,14 +155,22 @@ def test_map_message_to_langchain_ai(self) -> None:
     def test_map_message_to_langchain_ai_with_agent_call(self) -> None:
         message = AIMessage(
             content="hi",
-            calls=[SubagentCall(name="assistant", args={"q": "test"}, id="tc-2")],
+            calls=[
+                SubagentCall(
+                    name="assistant",
+                    args={"q": "test"},
+                    id="tc-2",
+                )
+            ],
         )
         mapped = lc._map_message_to_langchain(message)
 
         assert isinstance(mapped, LC_AIMessage)
         assert mapped.tool_calls == [
             LC_ToolCall(
-                name=f"{lc.AGENT_PREFIX}assistant", args={"q": "test"}, id="tc-2"
+                name=f"{lc.AGENT_PREFIX}assistant",
+                args={"args": {"q": "test"}},
+                id="tc-2",
             )
         ]
 
@@ -268,7 +276,7 @@ def test_map_message_to_langchain_agent_call_with_agent_prefix_raises(
         # Fine, but in practice a unnecessary prefix.
         assert isinstance(message, LC_AIMessage)
         assert message.tool_calls == [
-            LC_ToolCall(name="__agent-__agent-bad-agent", args={}, id="tc-1")
+            LC_ToolCall(name="__agent-__agent-bad-agent", args={"args": {}}, id="tc-1")
         ]
 
     def test_map_message_to_langchain_system(self) -> None: