Merge pull request #240 from UiPath/feat/add-escalations-support

feat: add escalation tool for ReAct Agent

Author: cotovanu-cristian

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-langchain"
-version = "0.1.20"
+version = "0.1.21"
 description = "UiPath Langchain"
 readme = { file = "README.md", content-type = "text/markdown" }
 requires-python = ">=3.11"

src/uipath_langchain/agent/react/terminate_node.py

@@ -1,5 +1,9 @@
 """Termination node for the Agent graph."""
 
+from __future__ import annotations
+
+from typing import Any, NoReturn
+
 from langchain_core.messages import AIMessage
 from pydantic import BaseModel
 from uipath.agent.react import END_EXECUTION_TOOL, RAISE_ERROR_TOOL
@@ -9,15 +13,53 @@
     AgentNodeRoutingException,
     AgentTerminationException,
 )
-from .types import AgentGraphState
+from .types import AgentGraphState, AgentTermination
+
+
+def _handle_end_execution(
+    args: dict[str, Any], response_schema: type[BaseModel] | None
+) -> dict[str, Any]:
+    """Handle LLM-initiated termination via END_EXECUTION_TOOL."""
+    output_schema = response_schema or END_EXECUTION_TOOL.args_schema
+    validated = output_schema.model_validate(args)
+    return validated.model_dump()
+
+
+def _handle_raise_error(args: dict[str, Any]) -> NoReturn:
+    """Handle LLM-initiated error via RAISE_ERROR_TOOL."""
+    error_message = args.get("message", "The LLM did not set the error message")
+    detail = args.get("details", "")
+    raise AgentTerminationException(
+        code=UiPathErrorCode.EXECUTION_ERROR,
+        title=error_message,
+        detail=detail,
+    )
+
+
+def _handle_agent_termination(termination: AgentTermination) -> NoReturn:
+    """Handle Command-based termination."""
+    raise AgentTerminationException(
+        code=UiPathErrorCode.EXECUTION_ERROR,
+        title=termination.title,
+        detail=termination.detail,
+    )
 
 
 def create_terminate_node(
     response_schema: type[BaseModel] | None = None,
 ):
-    """Validates and extracts end_execution args to state output field."""
+    """Handles Agent Graph termination for multiple sources and output or error propagation to Orchestrator.
+
+    Termination scenarios:
+    1. Command based termination with information in state (e.g: escalation)
+    2. LLM-initiated termination (END_EXECUTION_TOOL)
+    3. LLM-initiated error (RAISE_ERROR_TOOL)
+    """
 
     def terminate_node(state: AgentGraphState):
+        if state.termination:
+            _handle_agent_termination(state.termination)
+
         last_message = state.messages[-1]
         if not isinstance(last_message, AIMessage):
             raise AgentNodeRoutingException(
@@ -28,21 +70,10 @@ def terminate_node(state: AgentGraphState):
             tool_name = tool_call["name"]
 
             if tool_name == END_EXECUTION_TOOL.name:
-                args = tool_call["args"]
-                output_schema = response_schema or END_EXECUTION_TOOL.args_schema
-                validated = output_schema.model_validate(args)
-                return validated.model_dump()
+                return _handle_end_execution(tool_call["args"], response_schema)
 
             if tool_name == RAISE_ERROR_TOOL.name:
-                error_message = tool_call["args"].get(
-                    "message", "The LLM did not set the error message"
-                )
-                detail = tool_call["args"].get("details", "")
-                raise AgentTerminationException(
-                    code=UiPathErrorCode.EXECUTION_ERROR,
-                    title=error_message,
-                    detail=detail,
-                )
+                _handle_raise_error(tool_call["args"])
 
         raise AgentNodeRoutingException(
             "No control flow tool call found in terminate node. Unexpected state."

src/uipath_langchain/agent/react/types.py

@@ -6,10 +6,23 @@
 from pydantic import BaseModel, Field
 
 
+class AgentTerminationSource(StrEnum):
+    ESCALATION = "escalation"
+
+
+class AgentTermination(BaseModel):
+    """Agent Graph Termination model."""
+
+    source: AgentTerminationSource
+    title: str
+    detail: str = ""
+
+
 class AgentGraphState(BaseModel):
     """Agent Graph state for standard loop execution."""
 
     messages: Annotated[list[AnyMessage], add_messages] = []
+    termination: AgentTermination | None = None
 
 
 class AgentGraphNode(StrEnum):

src/uipath_langchain/agent/react/utils.py

@@ -2,18 +2,18 @@
 
 from typing import Any, Sequence
 
-from jsonschema_pydantic_converter import transform as create_model
 from langchain_core.messages import AIMessage, BaseMessage
 from pydantic import BaseModel
 from uipath.agent.react import END_EXECUTION_TOOL
+from uipath.utils.dynamic_schema import jsonschema_to_pydantic
 
 
 def resolve_input_model(
     input_schema: dict[str, Any] | None,
 ) -> type[BaseModel]:
     """Resolve the input model from the input schema."""
     if input_schema:
-        return create_model(input_schema)
+        return jsonschema_to_pydantic(input_schema)
 
     return BaseModel
 
@@ -23,7 +23,7 @@ def resolve_output_model(
 ) -> type[BaseModel]:
     """Fallback to default end_execution tool schema when no agent output schema is provided."""
     if output_schema:
-        return create_model(output_schema)
+        return jsonschema_to_pydantic(output_schema)
 
     return END_EXECUTION_TOOL.args_schema
 

src/uipath_langchain/agent/tools/escalation_tool.py

@@ -0,0 +1,111 @@
+"""Escalation tool creation for Action Center integration."""
+
+from enum import Enum
+from typing import Any
+
+from jsonschema_pydantic_converter import transform as create_model
+from langchain.tools import ToolRuntime
+from langchain_core.messages import ToolMessage
+from langchain_core.tools import StructuredTool
+from langgraph.types import Command, interrupt
+from uipath.agent.models.agent import (
+    AgentEscalationChannel,
+    AgentEscalationRecipientType,
+    AgentEscalationResourceConfig,
+)
+from uipath.eval.mocks import mockable
+from uipath.platform.common import CreateEscalation
+
+from ..react.types import AgentGraphNode, AgentTerminationSource
+from .utils import sanitize_tool_name
+
+
+class EscalationAction(str, Enum):
+    """Actions that can be taken after an escalation completes."""
+
+    CONTINUE = "continue"
+    END = "end"
+
+
+def create_escalation_tool(resource: AgentEscalationResourceConfig) -> StructuredTool:
+    """Uses interrupt() for Action Center human-in-the-loop."""
+
+    tool_name: str = f"escalate_{sanitize_tool_name(resource.name)}"
+    channel: AgentEscalationChannel = resource.channels[0]
+
+    input_model: Any = create_model(channel.input_schema)
+    output_model: Any = create_model(channel.output_schema)
+
+    assignee: str | None = (
+        channel.recipients[0].value
+        if channel.recipients
+        and channel.recipients[0].type == AgentEscalationRecipientType.USER_EMAIL
+        else None
+    )
+
+    @mockable(
+        name=resource.name,
+        description=resource.description,
+        input_schema=input_model.model_json_schema(),
+        output_schema=output_model.model_json_schema(),
+    )
+    async def escalation_tool_fn(
+        runtime: ToolRuntime, **kwargs: Any
+    ) -> Command[Any] | Any:
+        task_title = channel.task_title or "Escalation Task"
+
+        result = interrupt(
+            CreateEscalation(
+                title=task_title,
+                data=kwargs,
+                assignee=assignee,
+                app_name=channel.properties.app_name,
+                app_folder_path=channel.properties.folder_name,
+                app_version=channel.properties.app_version,
+                priority=channel.priority,
+                labels=channel.labels,
+                is_actionable_message_enabled=channel.properties.is_actionable_message_enabled,
+                actionable_message_metadata=channel.properties.actionable_message_meta_data,
+            )
+        )
+
+        escalation_action = getattr(result, "action", None)
+        escalation_output = getattr(result, "data", {})
+
+        outcome = (
+            channel.outcome_mapping.get(escalation_action)
+            if channel.outcome_mapping and escalation_action
+            else None
+        )
+
+        if outcome == EscalationAction.END:
+            output_detail = f"Escalation output: {escalation_output}"
+            termination_title = f"Agent run ended based on escalation outcome {outcome} with directive {escalation_action}"
+
+            return Command(
+                update={
+                    "messages": [
+                        ToolMessage(
+                            content=f"{termination_title}. {output_detail}",
+                            tool_call_id=runtime.tool_call_id,
+                        )
+                    ],
+                    "termination": {
+                        "source": AgentTerminationSource.ESCALATION,
+                        "title": termination_title,
+                        "detail": output_detail,
+                    },
+                },
+                goto=AgentGraphNode.TERMINATE,
+            )
+
+        return escalation_output
+
+    tool = StructuredTool(
+        name=tool_name,
+        description=resource.description,
+        args_schema=input_model,
+        coroutine=escalation_tool_fn,
+    )
+
+    return tool

src/uipath_langchain/agent/tools/tool_factory.py

@@ -3,13 +3,15 @@
 from langchain_core.tools import BaseTool, StructuredTool
 from uipath.agent.models.agent import (
     AgentContextResourceConfig,
+    AgentEscalationResourceConfig,
     AgentIntegrationToolResourceConfig,
     AgentProcessToolResourceConfig,
     BaseAgentResourceConfig,
     LowCodeAgentDefinition,
 )
 
 from .context_tool import create_context_tool
+from .escalation_tool import create_escalation_tool
 from .integration_tool import create_integration_tool
 from .process_tool import create_process_tool
 
@@ -34,6 +36,9 @@ async def _build_tool_for_resource(
     elif isinstance(resource, AgentContextResourceConfig):
         return create_context_tool(resource)
 
+    elif isinstance(resource, AgentEscalationResourceConfig):
+        return create_escalation_tool(resource)
+
     elif isinstance(resource, AgentIntegrationToolResourceConfig):
         return create_integration_tool(resource)