fix: update escalation action to extract tool calls

Author: ctiliescuuipath

src/uipath_langchain/agent/guardrails/actions/escalate_action.py

@@ -196,39 +196,54 @@ def _process_llm_escalation_response(
             if not reviewed_outputs_json:
                 return {}
 
-            content_list = json.loads(reviewed_outputs_json)
-            if not content_list:
+            reviewed_tool_calls_list = json.loads(reviewed_outputs_json)
+            if not reviewed_tool_calls_list:
                 return {}
 
+            # Track if tool calls were successfully processed
+            tool_calls_processed = False
+
             # For AI messages, process tool calls if present
             if isinstance(last_message, AIMessage):
                 ai_message: AIMessage = last_message
-                content_index = 0
 
-                if ai_message.tool_calls:
+                if ai_message.tool_calls and isinstance(reviewed_tool_calls_list, list):
                     tool_calls = list(ai_message.tool_calls)
-                    for tool_call in tool_calls:
-                        args = tool_call["args"]
+
+                    # Create a name-to-args mapping from reviewed tool call data
+                    reviewed_tool_calls_map = {}
+                    for reviewed_data in reviewed_tool_calls_list:
                         if (
-                            isinstance(args, dict)
-                            and "content" in args
-                            and args["content"] is not None
+                            isinstance(reviewed_data, dict)
+                            and "name" in reviewed_data
+                            and "args" in reviewed_data
                         ):
-                            if content_index < len(content_list):
-                                updated_content = json.loads(
-                                    content_list[content_index]
-                                )
-                                args["content"] = updated_content
-                                tool_call["args"] = args
-                                content_index += 1
-                    ai_message.tool_calls = tool_calls
-
-                if len(content_list) > content_index:
-                    ai_message.content = content_list[-1]
-            else:
-                # Fallback for other message types
-                if content_list:
-                    last_message.content = content_list[-1]
+                            reviewed_tool_calls_map[reviewed_data["name"]] = (
+                                reviewed_data["args"]
+                            )
+
+                    # Update tool calls with reviewed args by matching name
+                    if reviewed_tool_calls_map:
+                        for tool_call in tool_calls:
+                            tool_name = (
+                                tool_call.get("name")
+                                if isinstance(tool_call, dict)
+                                else getattr(tool_call, "name", None)
+                            )
+                            if tool_name and tool_name in reviewed_tool_calls_map:
+                                if isinstance(tool_call, dict):
+                                    tool_call["args"] = reviewed_tool_calls_map[
+                                        tool_name
+                                    ]
+                                else:
+                                    tool_call.args = reviewed_tool_calls_map[tool_name]
+
+                        ai_message.tool_calls = tool_calls
+                        tool_calls_processed = True
+
+            # Fallback: update message content if tool_calls weren't processed
+            if not tool_calls_processed:
+                last_message.content = reviewed_outputs_json
 
         return Command(update={"messages": msgs})
     except Exception as e:
@@ -388,23 +403,16 @@ def _extract_llm_escalation_content(
     # For AI messages, process tool calls if present
     if isinstance(last_message, AIMessage):
         ai_message: AIMessage = last_message
-        content_list: list[str] = []
 
         if ai_message.tool_calls:
+            content_list: list[Dict[str, Any]] = []
             for tool_call in ai_message.tool_calls:
-                args = tool_call["args"]
-                if (
-                    isinstance(args, dict)
-                    and "content" in args
-                    and args["content"] is not None
-                ):
-                    content_list.append(json.dumps(args["content"]))
-
-        message_content = get_message_content(last_message)
-        if message_content:
-            content_list.append(message_content)
-
-        return json.dumps(content_list)
+                tool_call_data = {
+                    "name": tool_call.get("name"),
+                    "args": tool_call.get("args"),
+                }
+                content_list.append(tool_call_data)
+            return json.dumps(content_list)
 
     # Fallback for other message types
     return get_message_content(last_message)

tests/agent/guardrails/actions/test_escalate_action.py

@@ -242,7 +242,7 @@ async def test_node_post_execution_tool_field(self, mock_interrupt):
 
         # Verify ToolOutputs is used for PostExecution
         call_args = mock_interrupt.call_args[0][0]
-        assert call_args.data["Outputs"] == '["Test response"]'
+        assert call_args.data["Outputs"] == "Test response"
         assert "Inputs" not in call_args.data
 
     @pytest.mark.asyncio
@@ -287,12 +287,13 @@ async def test_post_execution_ai_message_with_tool_calls_extraction(
 
         await node(state)
 
-        # Verify interrupt was called with tool calls and content in ToolOutputs
+        # Verify interrupt was called with tool calls (name and args) in ToolOutputs
         call_args = mock_interrupt.call_args[0][0]
         tool_outputs = call_args.data["Outputs"]
         parsed = json.loads(tool_outputs)
-        assert len(parsed) == 2  # Tool call content + message content
-        assert parsed[1] == "AI response"
+        assert len(parsed) == 1  # Tool call data with name and args
+        assert parsed[0]["name"] == "test_tool"
+        assert parsed[0]["args"] == {"content": {"input": "test"}}
 
     @pytest.mark.asyncio
     @patch("uipath_langchain.agent.guardrails.actions.escalate_action.interrupt")
@@ -367,10 +368,10 @@ async def test_post_execution_human_message_with_reviewed_outputs(
 
         result = await node(state)
 
-        # Verify HumanMessage content was updated (ignores tool calls)
+        # Verify HumanMessage content was updated with fallback (raw JSON string)
         assert isinstance(result, Command)
         assert result.update is not None
-        assert result.update["messages"][0].content == "Updated content"
+        assert result.update["messages"][0].content == json.dumps(reviewed_content)
 
     @pytest.mark.asyncio
     @patch("uipath_langchain.agent.guardrails.actions.escalate_action.interrupt")
@@ -388,12 +389,8 @@ async def test_post_execution_ai_message_with_reviewed_outputs_and_tool_calls(
         guardrail.name = "Test Guardrail"
         guardrail.description = "Test description"
 
-        reviewed_tool_content = {"updated": "tool_content"}
-        reviewed_message_content = "Updated message"
-        reviewed_outputs = [
-            json.dumps(reviewed_tool_content),
-            reviewed_message_content,
-        ]
+        reviewed_tool_args = {"updated": "tool_content"}
+        reviewed_outputs = [{"name": "test_tool", "args": reviewed_tool_args}]
         mock_escalation_result = MagicMock()
         mock_escalation_result.action = "Approve"
         mock_escalation_result.data = {"ReviewedOutputs": json.dumps(reviewed_outputs)}
@@ -420,12 +417,11 @@ async def test_post_execution_ai_message_with_reviewed_outputs_and_tool_calls(
 
         result = await node(state)
 
-        # Verify tool calls and message content were updated
+        # Verify tool calls args were updated by matching name
         assert isinstance(result, Command)
         assert result.update is not None
         updated_message = result.update["messages"][0]
-        assert updated_message.tool_calls[0]["args"]["content"] == reviewed_tool_content
-        assert updated_message.content == reviewed_message_content
+        assert updated_message.tool_calls[0]["args"] == reviewed_tool_args
 
     @pytest.mark.asyncio
     @patch("uipath_langchain.agent.guardrails.actions.escalate_action.interrupt")
@@ -1206,7 +1202,7 @@ async def test_extract_llm_content_pre_execution_empty_content(self):
 
     @pytest.mark.asyncio
     async def test_extract_llm_content_post_execution_tool_calls_no_content_field(self):
-        """Extract LLM content PostExecution: tool calls without content field are skipped."""
+        """Extract LLM content PostExecution: extracts all tool calls with name and args."""
         from uipath_langchain.agent.guardrails.actions.escalate_action import (
             _extract_llm_escalation_content,
         )
@@ -1227,9 +1223,10 @@ async def test_extract_llm_content_post_execution_tool_calls_no_content_field(se
 
         assert isinstance(result, str)
         parsed = json.loads(result)
-        # Should only contain message content, not tool call content
+        # Should extract tool call data with name and args
         assert len(parsed) == 1
-        assert parsed[0] == "Response"
+        assert parsed[0]["name"] == "tool_without_content"
+        assert parsed[0]["args"] == {"param": "value"}
 
     @pytest.mark.asyncio
     async def test_extract_escalation_content_empty_messages_raises_exception(self):