uipath-langchain 0.1.28__py3-none-any.whl → 0.3.1__py3-none-any.whl

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, ToolCall, 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 (
@@ -14,8 +15,9 @@ from uipath.platform.guardrails import (
 from uipath.runtime.errors import UiPathErrorCode
 
 from ...exceptions import AgentTerminationException
-from ..guardrail_nodes import _message_text
-from ..types import AgentGuardrailsGraphState, ExecutionStage
+from ...react.types import AgentGuardrailsGraphState
+from ..types import ExecutionStage
+from ..utils import _extract_tool_args_from_message, get_message_content
 from .base_action import GuardrailAction, GuardrailActionNode
 
 
@@ -71,25 +73,59 @@ class EscalateAction(GuardrailAction):
         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],
+                    state,
+                    scope,
+                    execution_stage,
+                    guarded_component_name,
+                )
+                data["Inputs"] = input_content
+            else:  # POST_EXECUTION
+                if scope == GuardrailScope.AGENT:
+                    input_message = state.messages[1]
+                else:
+                    input_message = state.messages[-2]
+                input_content = _extract_escalation_content(
+                    input_message,
+                    state,
+                    scope,
+                    ExecutionStage.PRE_EXECUTION,
+                    guarded_component_name,
+                )
+
+                # Extract Outputs from last message using POST_EXECUTION logic
+                output_content = _extract_escalation_content(
+                    state.messages[-1],
+                    state,
+                    scope,
+                    execution_stage,
+                    guarded_component_name,
+                )
+
+                data["Inputs"] = input_content
+                data["Outputs"] = output_content
+
             escalation_result = interrupt(
                 CreateEscalation(
                     app_name=self.app_name,
                     app_folder_path=self.app_folder_path,
-                    title=self.app_name,
+                    title="Agents Guardrail Task",
                     data=data,
                     assignee=self.assignee,
                 )
@@ -107,17 +143,50 @@ class EscalateAction(GuardrailAction):
             raise AgentTerminationException(
                 code=UiPathErrorCode.EXECUTION_ERROR,
                 title="Escalation rejected",
-                detail=f"Action was rejected after reviewing the task created by guardrail [{guardrail.name}]. Please contact your administrator.",
+                detail=f"Please contact your administrator. Action was rejected after reviewing the task created by guardrail [{guardrail.name}], with reason: {escalation_result.data['Reason']}",
             )
 
         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:
-    sanitized = re.sub(r"\W+", "_", guardrail.name).strip("_").lower()
-    node_name = f"{sanitized}_hitl_{execution_stage.name.lower()}_{scope.lower()}"
+    raw_node_name = f"{scope.name}_{execution_stage.name}_{guardrail.name}_hitl"
+    node_name = re.sub(r"\W+", "_", raw_node_name.lower()).strip("_")
     return node_name
 
 
@@ -137,8 +206,8 @@ def _process_escalation_response(
         execution_stage: The execution stage (PRE_EXECUTION or POST_EXECUTION).
 
     Returns:
-        For LLM/TOOL scope: Command to update messages with reviewed inputs/outputs, or empty dict.
-        For AGENT scope: Empty dict (no message alteration).
+        Command updates for the state (e.g., updating messages / tool calls / agent_result),
+        or an empty dict if no update is needed.
     """
     match scope:
         case GuardrailScope.LLM:
@@ -150,8 +219,71 @@ def _process_escalation_response(
                 state, escalation_result, execution_stage, guarded_node_name
             )
         case GuardrailScope.AGENT:
+            return _process_agent_escalation_response(
+                state, escalation_result, execution_stage
+            )
+
+
+def _process_agent_escalation_response(
+    state: AgentGuardrailsGraphState,
+    escalation_result: Dict[str, Any],
+    execution_stage: ExecutionStage,
+) -> Dict[str, Any] | Command[Any]:
+    """Process escalation response for AGENT scope guardrails.
+
+    For AGENT scope:
+    - PRE_EXECUTION: updates the last message content using `ReviewedInputs`
+    - POST_EXECUTION: updates `agent_result` using `ReviewedOutputs`
+
+    Args:
+        state: The current agent graph state.
+        escalation_result: The result from the escalation interrupt containing reviewed inputs/outputs.
+        execution_stage: The execution stage (PRE_EXECUTION or POST_EXECUTION).
+
+    Returns:
+        Command to update state, or empty dict if no updates are needed.
+
+    Raises:
+        AgentTerminationException: If escalation response processing fails.
+    """
+    try:
+        reviewed_field = get_reviewed_field_name(execution_stage)
+        if reviewed_field not in escalation_result:
+            return {}
+
+        reviewed_value = escalation_result.get(reviewed_field)
+        if not reviewed_value:
             return {}
 
+        try:
+            parsed = json.loads(reviewed_value)
+        except json.JSONDecodeError:
+            parsed = reviewed_value
+
+        if execution_stage == ExecutionStage.PRE_EXECUTION:
+            msgs = state.messages.copy()
+            if not msgs:
+                return {}
+            msgs[-1].content = parsed
+            return Command(update={"messages": msgs})
+
+        # POST_EXECUTION: update agent_result
+        return Command(update={"agent_result": parsed})
+    except Exception as e:
+        raise AgentTerminationException(
+            code=UiPathErrorCode.EXECUTION_ERROR,
+            title="Escalation rejected",
+            detail=str(e),
+        ) from e
+
+
+def get_reviewed_field_name(execution_stage):
+    return (
+        "ReviewedInputs"
+        if execution_stage == ExecutionStage.PRE_EXECUTION
+        else "ReviewedOutputs"
+    )
+
 
 def _process_llm_escalation_response(
     state: AgentGuardrailsGraphState,
@@ -174,11 +306,7 @@ def _process_llm_escalation_response(
         AgentTerminationException: If escalation response processing fails.
     """
     try:
-        reviewed_field = (
-            "ReviewedInputs"
-            if execution_stage == ExecutionStage.PRE_EXECUTION
-            else "ReviewedOutputs"
-        )
+        reviewed_field = get_reviewed_field_name(execution_stage)
 
         msgs = state.messages.copy()
         if not msgs or reviewed_field not in escalation_result:
@@ -195,41 +323,56 @@ 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]
-
-        return Command[Any](update={"messages": msgs})
+                            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:
         raise AgentTerminationException(
             code=UiPathErrorCode.EXECUTION_ERROR,
@@ -264,11 +407,7 @@ def _process_tool_escalation_response(
         AgentTerminationException: If escalation response processing fails.
     """
     try:
-        reviewed_field = (
-            "ReviewedInputs"
-            if execution_stage == ExecutionStage.PRE_EXECUTION
-            else "ReviewedOutputs"
-        )
+        reviewed_field = get_reviewed_field_name(execution_stage)
 
         msgs = state.messages.copy()
         if not msgs or reviewed_field not in escalation_result:
@@ -292,7 +431,11 @@ def _process_tool_escalation_response(
             if last_message.tool_calls:
                 tool_calls = list(last_message.tool_calls)
                 for tool_call in tool_calls:
-                    call_name = extract_tool_name(tool_call)
+                    call_name = (
+                        tool_call.get("name")
+                        if isinstance(tool_call, dict)
+                        else getattr(tool_call, "name", None)
+                    )
                     if call_name == tool_name:
                         # Update args for the matching tool call
                         if isinstance(reviewed_tool_calls_args, dict):
@@ -311,7 +454,7 @@ def _process_tool_escalation_response(
             if reviewed_outputs_json:
                 last_message.content = reviewed_outputs_json
 
-        return Command[Any](update={"messages": msgs})
+        return Command(update={"messages": msgs})
     except Exception as e:
         raise AgentTerminationException(
             code=UiPathErrorCode.EXECUTION_ERROR,
@@ -321,50 +464,65 @@ def _process_tool_escalation_response(
 
 
 def _extract_escalation_content(
+    message: BaseMessage,
     state: AgentGuardrailsGraphState,
     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, state, execution_stage)
         case GuardrailScope.TOOL:
             return _extract_tool_escalation_content(
-                state, execution_stage, guarded_node_name
+                message, execution_stage, guarded_node_name
             )
 
 
+def _extract_agent_escalation_content(
+    message: BaseMessage,
+    state: AgentGuardrailsGraphState,
+    execution_stage: ExecutionStage,
+) -> str | list[str | Dict[str, Any]]:
+    """Extract escalation content for AGENT scope guardrails.
+
+    Args:
+        message: The message used to extract the agent input content.
+        state: The current agent guardrails graph state. Used to read `agent_result` for POST_EXECUTION.
+        execution_stage: PRE_EXECUTION or POST_EXECUTION.
+
+    Returns:
+        - PRE_EXECUTION: the agent input string (from message content).
+        - POST_EXECUTION: a JSON-serialized representation of `state.agent_result`.
+    """
+    if execution_stage == ExecutionStage.PRE_EXECUTION:
+        return get_message_content(cast(AnyMessage, message))
+
+    output_content = state.agent_result or ""
+    return json.dumps(output_content)
+
+
 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:
@@ -372,61 +530,37 @@ def _extract_llm_escalation_content(
         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 = _message_text(last_message)
-        return json.dumps(content) if content else ""
+        return get_message_content(cast(AnyMessage, message))
 
     # 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 = _message_text(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 _message_text(last_message)
-
-
-def _extract_agent_escalation_content(
-    state: AgentGuardrailsGraphState, execution_stage: ExecutionStage
-) -> str | list[str | Dict[str, Any]]:
-    """Extract escalation content for AGENT scope guardrails.
-
-    Args:
-        state: The current agent graph state.
-        execution_stage: The execution stage (PRE_EXECUTION or POST_EXECUTION).
-
-    Returns:
-        str: Empty string (AGENT scope guardrails do not extract escalation content).
-    """
-    return ""
+    return get_message_content(cast(AnyMessage, message))
 
 
 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.
 
@@ -435,38 +569,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:
-        if not isinstance(last_message, AIMessage):
-            return ""
-        if not last_message.tool_calls:
-            return ""
-
-        # Find the tool call with matching name
-        for tool_call in last_message.tool_calls:
-            call_name = extract_tool_name(tool_call)
-            if call_name == tool_name:
-                # Extract args from the matching tool call
-                args = (
-                    tool_call.get("args")
-                    if isinstance(tool_call, dict)
-                    else getattr(tool_call, "args", None)
-                )
-                if args is not None:
-                    return json.dumps(args)
+        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
-
-
-def extract_tool_name(tool_call: ToolCall) -> Any | None:
-    return (
-        tool_call.get("name")
-        if isinstance(tool_call, dict)
-        else getattr(tool_call, "name", None)
-    )
+        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(