Merge pull request #366 from UiPath/fix-escalation-tool-calls-extraction

fix: update escalation action to extract tool calls

Author: ctiliescuuipath

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-langchain"
-version = "0.1.37"
+version = "0.1.38"
 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/guardrails/actions/escalate_action.py

@@ -1,10 +1,11 @@
 from __future__ import annotations
 
+import ast
 import json
 import re
-from typing import Any, Dict, Literal
+from typing import Any, Dict, Literal, cast
 
-from langchain_core.messages import AIMessage, ToolMessage
+from langchain_core.messages import AIMessage, AnyMessage, BaseMessage, ToolMessage
 from langgraph.types import Command, interrupt
 from uipath.platform.common import CreateEscalation
 from uipath.platform.guardrails import (
@@ -72,20 +73,48 @@ def action_node(
         async def _node(
             state: AgentGuardrailsGraphState,
         ) -> Dict[str, Any] | Command[Any]:
-            input = _extract_escalation_content(
-                state, scope, execution_stage, guarded_component_name
-            )
-            escalation_field = _execution_stage_to_escalation_field(execution_stage)
+            # Validate message count based on execution stage
+            _validate_message_count(state, execution_stage)
 
-            data = {
+            # Build base data dictionary with common fields
+            data: Dict[str, Any] = {
                 "GuardrailName": guardrail.name,
                 "GuardrailDescription": guardrail.description,
                 "Component": scope.name.lower(),
                 "ExecutionStage": _execution_stage_to_string(execution_stage),
                 "GuardrailResult": state.guardrail_validation_result,
-                escalation_field: input,
             }
 
+            # Add stage-specific fields
+            if execution_stage == ExecutionStage.PRE_EXECUTION:
+                # PRE_EXECUTION: Only Inputs field from last message
+                input_content = _extract_escalation_content(
+                    state.messages[-1],
+                    scope,
+                    execution_stage,
+                    guarded_component_name,
+                )
+                data["Inputs"] = input_content
+            else:  # POST_EXECUTION
+                # Extract Inputs from second-to-last message using PRE_EXECUTION logic
+                input_content = _extract_escalation_content(
+                    state.messages[-2],
+                    scope,
+                    ExecutionStage.PRE_EXECUTION,
+                    guarded_component_name,
+                )
+
+                # Extract Outputs from last message using POST_EXECUTION logic
+                output_content = _extract_escalation_content(
+                    state.messages[-1],
+                    scope,
+                    execution_stage,
+                    guarded_component_name,
+                )
+
+                data["Inputs"] = input_content
+                data["Outputs"] = output_content
+
             escalation_result = interrupt(
                 CreateEscalation(
                     app_name=self.app_name,
@@ -114,6 +143,39 @@ async def _node(
         return node_name, _node
 
 
+def _validate_message_count(
+    state: AgentGuardrailsGraphState,
+    execution_stage: ExecutionStage,
+) -> None:
+    """Validate that state has the required number of messages for the execution stage.
+
+    Args:
+        state: The current agent graph state.
+        execution_stage: The execution stage (PRE_EXECUTION or POST_EXECUTION).
+
+    Raises:
+        AgentTerminationException: If the state doesn't have enough messages.
+    """
+    required_messages = 1 if execution_stage == ExecutionStage.PRE_EXECUTION else 2
+    actual_messages = len(state.messages)
+
+    if actual_messages < required_messages:
+        stage_name = (
+            "PRE_EXECUTION"
+            if execution_stage == ExecutionStage.PRE_EXECUTION
+            else "POST_EXECUTION"
+        )
+        detail = f"{stage_name} requires at least {required_messages} message{'s' if required_messages > 1 else ''} in state, but found {actual_messages}."
+        if execution_stage == ExecutionStage.POST_EXECUTION:
+            detail += " Cannot extract Inputs from previous message."
+
+        raise AgentTerminationException(
+            code=UiPathErrorCode.EXECUTION_ERROR,
+            title=f"Invalid state for {stage_name}",
+            detail=detail,
+        )
+
+
 def _get_node_name(
     execution_stage: ExecutionStage, guardrail: BaseGuardrail, scope: GuardrailScope
 ) -> str:
@@ -196,39 +258,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:
@@ -326,97 +403,80 @@ def _process_tool_escalation_response(
 
 
 def _extract_escalation_content(
-    state: AgentGuardrailsGraphState,
+    message: BaseMessage,
     scope: GuardrailScope,
     execution_stage: ExecutionStage,
     guarded_node_name: str,
 ) -> str | list[str | Dict[str, Any]]:
-    """Extract escalation content from state based on guardrail scope and execution stage.
+    """Extract escalation content from a message based on guardrail scope and execution stage.
 
     Args:
-        state: The current agent graph state.
+        message: The message to extract content from.
         scope: The guardrail scope (LLM/AGENT/TOOL).
         execution_stage: The execution stage (PRE_EXECUTION or POST_EXECUTION).
+        guarded_node_name: Name of the guarded component.
 
     Returns:
         str or list[str | Dict[str, Any]]: For LLM scope, returns JSON string or list with message/tool call content.
         For AGENT scope, returns empty string. For TOOL scope, returns JSON string or list with tool-specific content.
-
-    Raises:
-        AgentTerminationException: If no messages are found in state.
     """
-    if not state.messages:
-        raise AgentTerminationException(
-            code=UiPathErrorCode.EXECUTION_ERROR,
-            title="Invalid state message",
-            detail="No message found into agent state",
-        )
-
     match scope:
         case GuardrailScope.LLM:
-            return _extract_llm_escalation_content(state, execution_stage)
+            return _extract_llm_escalation_content(message, execution_stage)
         case GuardrailScope.AGENT:
-            return _extract_agent_escalation_content(state, execution_stage)
+            return _extract_agent_escalation_content(message, execution_stage)
         case GuardrailScope.TOOL:
             return _extract_tool_escalation_content(
-                state, execution_stage, guarded_node_name
+                message, execution_stage, guarded_node_name
             )
 
 
 def _extract_llm_escalation_content(
-    state: AgentGuardrailsGraphState, execution_stage: ExecutionStage
+    message: BaseMessage, execution_stage: ExecutionStage
 ) -> str | list[str | Dict[str, Any]]:
     """Extract escalation content for LLM scope guardrails.
 
     Args:
-        state: The current agent graph state.
+        message: The message to extract content from.
         execution_stage: The execution stage (PRE_EXECUTION or POST_EXECUTION).
 
     Returns:
         str or list[str | Dict[str, Any]]: For PreExecution, returns JSON string with message content or empty string.
         For PostExecution, returns JSON string (array) with tool call content and message content.
         Returns empty string if no content found.
     """
-    last_message = state.messages[-1]
     if execution_stage == ExecutionStage.PRE_EXECUTION:
-        if isinstance(last_message, ToolMessage):
-            return last_message.content
+        if isinstance(message, ToolMessage):
+            return message.content
 
-        content = get_message_content(last_message)
+        content = get_message_content(cast(AnyMessage, message))
         return json.dumps(content) if content else ""
 
     # For AI messages, process tool calls if present
-    if isinstance(last_message, AIMessage):
-        ai_message: AIMessage = last_message
-        content_list: list[str] = []
+    if isinstance(message, AIMessage):
+        ai_message: AIMessage = message
 
         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)
+    return get_message_content(cast(AnyMessage, message))
 
 
 def _extract_agent_escalation_content(
-    state: AgentGuardrailsGraphState, execution_stage: ExecutionStage
+    message: BaseMessage, execution_stage: ExecutionStage
 ) -> str | list[str | Dict[str, Any]]:
     """Extract escalation content for AGENT scope guardrails.
 
     Args:
-        state: The current agent graph state.
+        message: The message to extract content from.
         execution_stage: The execution stage (PRE_EXECUTION or POST_EXECUTION).
 
     Returns:
@@ -426,12 +486,12 @@ def _extract_agent_escalation_content(
 
 
 def _extract_tool_escalation_content(
-    state: AgentGuardrailsGraphState, execution_stage: ExecutionStage, tool_name: str
+    message: BaseMessage, execution_stage: ExecutionStage, tool_name: str
 ) -> str | list[str | Dict[str, Any]]:
     """Extract escalation content for TOOL scope guardrails.
 
     Args:
-        state: The current agent graph state.
+        message: The message to extract content from.
         execution_stage: The execution stage (PRE_EXECUTION or POST_EXECUTION).
         tool_name: Optional tool name to filter tool calls. If provided, only extracts args for matching tool.
 
@@ -440,16 +500,31 @@ def _extract_tool_escalation_content(
         for the specified tool name, or empty string if not found. For PostExecution, returns string with
         tool message content, or empty string if message type doesn't match.
     """
-    last_message = state.messages[-1]
     if execution_stage == ExecutionStage.PRE_EXECUTION:
-        args = _extract_tool_args_from_message(last_message, tool_name)
+        args = _extract_tool_args_from_message(cast(AnyMessage, message), tool_name)
         if args:
             return json.dumps(args)
         return ""
     else:
-        if not isinstance(last_message, ToolMessage):
+        if not isinstance(message, ToolMessage):
             return ""
-        return last_message.content
+        content = message.content
+
+        # If content is already dict/list, serialize to JSON
+        if isinstance(content, (dict, list)):
+            return json.dumps(content)
+
+        # If content is a string that looks like a Python literal, convert to JSON
+        if isinstance(content, str):
+            try:
+                # Try to parse as Python literal and convert to JSON
+                parsed_content = ast.literal_eval(content)
+                return json.dumps(parsed_content)
+            except (ValueError, SyntaxError):
+                # If parsing fails, return as-is
+                pass
+
+        return content
 
 
 def _execution_stage_to_escalation_field(

src/uipath_langchain/agent/react/guardrails/guardrails_subgraph.py

@@ -117,7 +117,7 @@ def _create_guardrails_subgraph(
             ExecutionStage.POST_EXECUTION,
             node_factory,
             END,
-            inner_node,
+            inner_name,
         )
         subgraph.add_edge(inner_name, first_post_exec_guardrail_node)
     else: