fix: invoke the tool using modified ToolCall object (#542)

Author: cristian-groza

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-langchain"
-version = "0.5.36"
+version = "0.5.37"
 description = "Python SDK that enables developers to build and deploy LangGraph agents to the UiPath Cloud Platform"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"

src/uipath_langchain/agent/tools/integration_tool.py

@@ -179,8 +179,8 @@ async def integration_tool_wrapper(
         call: ToolCall,
         state: AgentGraphState,
     ) -> ToolWrapperReturnType:
-        modified_args = handle_static_args(resource, state, call["args"])
-        return await tool.ainvoke(modified_args)
+        call["args"] = handle_static_args(resource, state, call["args"])
+        return await tool.ainvoke(call)
 
     tool = StructuredToolWithWrapper(
         name=tool_name,

src/uipath_langchain/agent/tools/tool_node.py

@@ -75,7 +75,7 @@ def _func(self, state: AgentGraphState) -> OutputType:
             inputs = self._prepare_wrapper_inputs(self.wrapper, self.tool, call, state)
             result = self.wrapper(*inputs)
         else:
-            result = self.tool.invoke(call["args"])
+            result = self.tool.invoke(call)
         return self._process_result(call, result)
 
     async def _afunc(self, state: AgentGraphState) -> OutputType:
@@ -86,7 +86,7 @@ async def _afunc(self, state: AgentGraphState) -> OutputType:
             inputs = self._prepare_wrapper_inputs(self.awrapper, self.tool, call, state)
             result = await self.awrapper(*inputs)
         else:
-            result = await self.tool.ainvoke(call["args"])
+            result = await self.tool.ainvoke(call)
         return self._process_result(call, result)
 
     def _extract_tool_call(self, state: AgentGraphState) -> ToolCall | None:
@@ -110,11 +110,13 @@ def _extract_tool_call(self, state: AgentGraphState) -> ToolCall | None:
         return latest_ai_message.tool_calls[current_tool_call_index]
 
     def _process_result(
-        self, call: ToolCall, result: dict[str, Any] | Command[Any] | None
+        self, call: ToolCall, result: dict[str, Any] | Command[Any] | ToolMessage | None
     ) -> OutputType:
         """Process the tool result into a message format or return a Command."""
         if isinstance(result, Command):
             return result
+        elif isinstance(result, ToolMessage):
+            return {"messages": [result]}
         else:
             message = ToolMessage(
                 content=str(result), name=call["name"], tool_call_id=call["id"]

src/uipath_langchain/agent/wrappers/job_attachment_wrapper.py

@@ -1,6 +1,7 @@
+import json
 from typing import Any
 
-from langchain_core.messages.tool import ToolCall, ToolMessage
+from langchain_core.messages.tool import ToolCall
 from langchain_core.tools import BaseTool
 from langgraph.types import Command
 from pydantic import BaseModel
@@ -14,6 +15,18 @@
 from uipath_langchain.agent.tools.tool_node import AsyncToolWrapperWithState
 
 
+def _parse(content: str) -> Any:
+    if not content:
+        return content
+
+    try:
+        return json.loads(content)
+    except (json.JSONDecodeError, TypeError):
+        pass
+
+    return content
+
+
 def get_job_attachment_wrapper(
     output_type: Any | None = None,
 ) -> AsyncToolWrapperWithState:
@@ -71,24 +84,20 @@ async def job_attachment_wrapper(
 
             if errors:
                 return {"error": "\n".join(errors)}
-
-        tool_result = await tool.ainvoke(modified_input_args)
+        call["args"] = modified_input_args
+        tool_result = await tool.ainvoke(call)
         job_attachments_dict = {}
         if output_type is not None:
-            job_attachments = get_job_attachments(output_type, tool_result)
+            job_attachments = get_job_attachments(
+                output_type, _parse(tool_result.content)
+            )
             job_attachments_dict = {
                 str(att.id): att for att in job_attachments if att.id is not None
             }
 
         return Command(
             update={
-                "messages": [
-                    ToolMessage(
-                        content=str(tool_result),
-                        name=call["name"],
-                        tool_call_id=call["id"],
-                    )
-                ],
+                "messages": [tool_result],
                 "inner_state": {"job_attachments": job_attachments_dict},
             }
         )

tests/agent/wrappers/test_job_attachment_wrapper.py

@@ -1,10 +1,12 @@
 """Tests for job_attachment_wrapper module."""
 
+import json
 import uuid
 from typing import Any, cast
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
+from langchain_core.messages import ToolMessage
 from langchain_core.messages.tool import ToolCall
 from langchain_core.tools import BaseTool
 from langgraph.types import Command
@@ -85,7 +87,14 @@ def create_mock_attachment(self, attachment_id: uuid.UUID) -> MagicMock:
     def mock_tool(self):
         """Create a mock tool."""
         tool = MagicMock(spec=BaseTool)
-        tool.ainvoke = AsyncMock(return_value={"result": "success"})
+
+        async def ainvoke_side_effect(call):
+            tool_message = ToolMessage(
+                content="{'result': 'success'}", tool_call_id=call.get("id", "call_123")
+            )
+            return tool_message
+
+        tool.ainvoke = AsyncMock(side_effect=ainvoke_side_effect)
         return tool
 
     @pytest.fixture
@@ -129,7 +138,7 @@ async def test_tool_without_args_schema(
 
         assert isinstance(result, Command)
         self.assert_command_success(result)
-        mock_tool.ainvoke.assert_awaited_once_with(mock_tool_call["args"])
+        mock_tool.ainvoke.assert_awaited_once_with(mock_tool_call)
 
     @pytest.mark.asyncio
     async def test_tool_with_dict_args_schema(
@@ -143,7 +152,7 @@ async def test_tool_with_dict_args_schema(
 
         assert isinstance(result, Command)
         self.assert_command_success(result)
-        mock_tool.ainvoke.assert_awaited_once_with(mock_tool_call["args"])
+        mock_tool.ainvoke.assert_awaited_once_with(mock_tool_call)
 
     @pytest.mark.asyncio
     async def test_tool_with_non_basemodel_schema(
@@ -157,7 +166,7 @@ async def test_tool_with_non_basemodel_schema(
 
         assert isinstance(result, Command)
         self.assert_command_success(result)
-        mock_tool.ainvoke.assert_awaited_once_with(mock_tool_call["args"])
+        mock_tool.ainvoke.assert_awaited_once_with(mock_tool_call)
 
     @pytest.mark.asyncio
     @patch(
@@ -180,7 +189,7 @@ async def test_tool_with_no_attachment_paths(
         assert isinstance(result, Command)
         self.assert_command_success(result)
         mock_get_paths.assert_called_once_with(MockAttachmentSchema)
-        mock_tool.ainvoke.assert_awaited_once_with(mock_tool_call["args"])
+        mock_tool.ainvoke.assert_awaited_once_with(mock_tool_call)
 
     @pytest.mark.asyncio
     @patch(
@@ -219,9 +228,9 @@ async def test_tool_with_valid_attachments(
         self.assert_command_success(result)
         # Verify that tool.ainvoke was called (with replaced attachment)
         mock_tool.ainvoke.assert_awaited_once()
-        called_args = mock_tool.ainvoke.call_args[0][0]
-        assert called_args["name"] == "test"
-        assert "attachment" in called_args
+        called_tool_call = mock_tool.ainvoke.call_args[0][0]
+        assert called_tool_call["args"]["name"] == "test"
+        assert "attachment" in called_tool_call["args"]
 
     @pytest.mark.asyncio
     @patch(
@@ -557,7 +566,8 @@ async def test_tool_with_complex_nested_structure_all_valid(
         self.assert_command_success(result, tool_call_id="call_complex_456")
         # Tool should be invoked with replaced attachments
         mock_tool.ainvoke.assert_awaited_once()
-        called_args = mock_tool.ainvoke.call_args[0][0]
+        called_tool_call = mock_tool.ainvoke.call_args[0][0]
+        called_args = called_tool_call["args"]
 
         # Verify structure is preserved
         assert "request" in called_args
@@ -588,12 +598,14 @@ async def test_structured_tool_with_output_attachments(
         # Create a tool with output type
         mock_tool = MagicMock(spec=BaseTool)
         mock_tool.args_schema = None
-        mock_tool.ainvoke = AsyncMock(
-            return_value={
-                "result_attachment": output_attachment.model_dump(),
-                "additional_attachments": [],
-            }
+        tool_output = {
+            "result_attachment": output_attachment.model_dump(),
+            "additional_attachments": [],
+        }
+        tool_message = ToolMessage(
+            content=json.dumps(tool_output), tool_call_id=mock_tool_call["id"]
         )
+        mock_tool.ainvoke = AsyncMock(return_value=tool_message)
 
         wrapper = get_job_attachment_wrapper(output_type=MockOutputSchema)
         result = await wrapper(mock_tool, mock_tool_call, mock_state)
@@ -606,9 +618,7 @@ async def test_structured_tool_with_output_attachments(
             expected_content=None,
         )
         # Verify get_job_attachments was called with correct parameters
-        mock_get_job_attachments.assert_called_once_with(
-            MockOutputSchema, mock_tool.ainvoke.return_value
-        )
+        mock_get_job_attachments.assert_called_once_with(MockOutputSchema, tool_output)
 
     @pytest.mark.asyncio
     @patch("uipath_langchain.agent.wrappers.job_attachment_wrapper.get_job_attachments")
@@ -631,15 +641,17 @@ async def test_structured_tool_with_multiple_output_attachments(
         # Create a tool with output type
         mock_tool = MagicMock(spec=BaseTool)
         mock_tool.args_schema = None
-        mock_tool.ainvoke = AsyncMock(
-            return_value={
-                "result_attachment": attachment1.model_dump(),
-                "additional_attachments": [
-                    attachment2.model_dump(),
-                    attachment3.model_dump(),
-                ],
-            }
+        tool_output = {
+            "result_attachment": attachment1.model_dump(),
+            "additional_attachments": [
+                attachment2.model_dump(),
+                attachment3.model_dump(),
+            ],
+        }
+        tool_message = ToolMessage(
+            content=json.dumps(tool_output), tool_call_id=mock_tool_call["id"]
         )
+        mock_tool.ainvoke = AsyncMock(return_value=tool_message)
 
         wrapper = get_job_attachment_wrapper(output_type=MockOutputSchema)
         result = await wrapper(mock_tool, mock_tool_call, mock_state)
@@ -668,9 +680,14 @@ async def test_structured_tool_with_no_output_attachments(
         # Create a tool with output type
         mock_tool = MagicMock(spec=BaseTool)
         mock_tool.args_schema = None
-        mock_tool.ainvoke = AsyncMock(
-            return_value={"result_attachment": None, "additional_attachments": []}
+        tool_output: dict[str, Any] = {
+            "result_attachment": None,
+            "additional_attachments": [],
+        }
+        tool_message = ToolMessage(
+            content=json.dumps(tool_output), tool_call_id=mock_tool_call["id"]
         )
+        mock_tool.ainvoke = AsyncMock(return_value=tool_message)
 
         wrapper = get_job_attachment_wrapper(output_type=MockOutputSchema)
         result = await wrapper(mock_tool, mock_tool_call, mock_state)
@@ -703,12 +720,14 @@ async def test_structured_tool_filters_attachments_with_none_id(
         # Create a tool with output type
         mock_tool = MagicMock(spec=BaseTool)
         mock_tool.args_schema = None
-        mock_tool.ainvoke = AsyncMock(
-            return_value={
-                "result_attachment": attachment_with_id.model_dump(),
-                "additional_attachments": [attachment_without_id.model_dump()],
-            }
+        tool_output = {
+            "result_attachment": attachment_with_id.model_dump(),
+            "additional_attachments": [attachment_without_id.model_dump()],
+        }
+        tool_message = ToolMessage(
+            content=json.dumps(tool_output), tool_call_id=mock_tool_call["id"]
         )
+        mock_tool.ainvoke = AsyncMock(return_value=tool_message)
 
         wrapper = get_job_attachment_wrapper(output_type=MockOutputSchema)
         result = await wrapper(mock_tool, mock_tool_call, mock_state)