WIP Al 201 block guardrail subgraph

src/uipath_langchain/agent/exceptions/__init__.py

@@ -0,0 +1,6 @@
+from .exceptions import AgentNodeRoutingException, AgentTerminationException
+
+__all__ = [
+    "AgentNodeRoutingException",
+    "AgentTerminationException",
+]

src/uipath_langchain/agent/guardrails/__init__.py

@@ -0,0 +1,13 @@
+from .guardrail_nodes import create_llm_guardrail_node, create_tool_guardrail_node, create_agent_guardrail_node
+from .guardrails_subgraph import create_llm_guardrails_subgraph, create_agent_guardrails_subgraph, \
+    create_tool_guardrails_subgraph
+
+__all__ = [
+    "create_llm_guardrails_subgraph",
+    "create_agent_guardrails_subgraph",
+    "create_tool_guardrails_subgraph",
+    "create_llm_guardrail_node",
+    "create_agent_guardrail_node",
+    "create_tool_guardrail_node",
+]
+

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

@@ -0,0 +1,14 @@
+from .base_action import GuardrailAction
+from .block_action import BlockAction
+from .escalate_action import EscalateAction
+from .log_action import LogAction
+
+__all__ = [
+    "GuardrailAction",
+    "BlockAction",
+    "LogAction",
+    "EscalateAction",
+]
+
+
+

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

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Literal, Tuple, Callable, Any
+
+from uipath.platform.guardrails import CustomGuardrail, BuiltInValidatorGuardrail, GuardrailScope
+
+from ..types import AgentGuardrailsGraphState
+
+
+class GuardrailAction(ABC):
+    """Extensible action interface producing a node for validation failure."""
+
+    @abstractmethod
+    def action_node(
+        self,
+        *,
+        guardrail: CustomGuardrail | BuiltInValidatorGuardrail,
+        scope: GuardrailScope,
+        execution_stage: Literal["PreExecution", "PostExecution"],
+    ) -> GuardrailActionNode:
+        """Create and return the GraphNode to execute on validation failure."""
+        ...
+
+
+GuardrailActionNode = Tuple[str, Callable[[AgentGuardrailsGraphState], Any]]

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

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+import re
+from typing import Literal, Dict, Any
+
+from uipath.platform.guardrails import CustomGuardrail, BuiltInValidatorGuardrail, GuardrailScope
+from uipath.runtime.errors import UiPathErrorCode, UiPathErrorCategory
+
+from ..types import AgentGuardrailsGraphState
+from .base_action import GuardrailAction, GuardrailActionNode
+from ...exceptions import AgentTerminationException
+
+
+class BlockAction(GuardrailAction):
+    """Action that terminates execution when a guardrail fails.
+
+    Args:
+        reason: Optional reason string to include in the raised exception title.
+    """
+
+    def __init__(self, reason: str) -> None:
+        self.reason = reason
+
+    def action_node(
+        self,
+        *,
+        guardrail: CustomGuardrail | BuiltInValidatorGuardrail,
+        scope: GuardrailScope,
+        execution_stage: Literal["PreExecution", "PostExecution"],
+    ) -> GuardrailActionNode:
+        sanitized = re.sub(r"\W+", "_", getattr(guardrail, "name", "guardrail")).strip(
+            "_"
+        )
+        node_name = f"{sanitized}_{execution_stage.lower()}_{scope.lower()}_block"
+
+        async def _node(_state: AgentGuardrailsGraphState) -> Dict[str, Any]:
+            raise AgentTerminationException(
+                code=UiPathErrorCode.EXECUTION_ERROR,
+                title="Guardrail violation",
+                detail=self.reason,
+                category=UiPathErrorCategory.USER,
+            )
+
+        return node_name, _node

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

@@ -0,0 +1,240 @@
+from __future__ import annotations
+
+import json
+import re
+from typing import Any, Dict, Literal
+
+from langchain_core.messages import AIMessage
+from langgraph.types import Command, interrupt
+from uipath.platform.common import CreateEscalation
+from uipath.platform.guardrails import (
+    BuiltInValidatorGuardrail,
+    CustomGuardrail,
+    GuardrailScope,
+)
+from uipath.runtime.errors import UiPathErrorCode
+
+from ...exceptions import AgentTerminationException
+from ..guardrail_nodes import _message_text
+from ..types import AgentGuardrailsGraphState
+from .base_action import GuardrailAction, GuardrailActionNode
+
+
+class EscalateAction(GuardrailAction):
+    """Node-producing action that inserts a HITL interruption node into the graph.
+
+    The returned node triggers a dynamic interrupt for HITL without re-evaluating.
+    The runtime will persist a resume trigger and suspend execution.
+    """
+
+    def __init__(
+        self,
+        app_name: str,
+        app_folder_path: str,
+        title: str,
+        version: int,
+        assignee: str,
+    ):
+        self.app_name = app_name
+        self.app_folder_path = app_folder_path
+        self.title = title
+        self.version = version
+        self.assignee = assignee
+
+    def action_node(
+        self,
+        *,
+        guardrail: CustomGuardrail | BuiltInValidatorGuardrail,
+        scope: GuardrailScope,
+        execution_stage: Literal["PreExecution", "PostExecution"],
+    ) -> GuardrailActionNode:
+        sanitized = re.sub(r"\W+", "_", guardrail.name).strip("_").lower()
+        node_name = f"{sanitized}_hitl_{execution_stage}_{scope.lower()}"
+
+        async def _node(state: AgentGuardrailsGraphState) -> Dict[str, Any]:
+            input = _extract_escalation_content(state, scope, execution_stage)
+            tool_field = _hook_type_to_tool_field(execution_stage)
+            data = {
+                "GuardrailName": guardrail.name,
+                "GuardrailDescription": guardrail.description,
+                "TenantName": "AgentsRuntime",
+                "AgentTrace": "https://alpha.uipath.com/f88fa028-ccdd-4b5f-bee4-01ef94d134d8/studio_/designer/48fff406-52e9-4a37-ba66-76c0212d9c6b",
+                "Tool": "Create_Issue",
+                "ExecutionStage": execution_stage,
+                "GuardrailResult": state.guardrail_validation_result,
+                tool_field: input,
+            }
+            escalation_result = interrupt(
+                CreateEscalation(
+                    app_name=self.app_name,
+                    app_folder_path=self.app_folder_path,
+                    title="Test",
+                    data=data,
+                    app_version=self.version,
+                    assignee=self.assignee,
+                )
+            )
+
+            if escalation_result.action == "Approve":
+                return _process_escalation_response(
+                    state, escalation_result.data, scope, execution_stage
+                )
+
+            raise AgentTerminationException(
+                code=UiPathErrorCode.CREATE_RESUME_TRIGGER_ERROR,
+                title="Escalation rejected",
+                detail="Escalation rejected",
+            )
+
+        return node_name, _node
+
+
+def _process_escalation_response(
+    state: AgentGuardrailsGraphState,
+    escalation_result: Dict[str, Any],
+    scope: GuardrailScope,
+    hook_type: Literal["PreExecution", "PostExecution"],
+) -> Dict[str, Any] | Command:
+    """Process escalation response and update state based on guardrail scope.
+
+    Args:
+        state: The current agent graph state.
+        escalation_result: The result from the escalation interrupt.
+        scope: The guardrail scope (LLM/AGENT/TOOL).
+        hook_type: The hook type ("PreExecution" or "PostExecution").
+
+    Returns:
+        For LLM scope: Command to update messages with reviewed inputs/outputs.
+        For non-LLM scope: Empty dict (no message alteration).
+
+    Raises:
+        AgentTerminationException: If escalation response processing fails.
+    """
+    if scope != GuardrailScope.LLM:
+        return {}
+
+    try:
+        reviewed_field = (
+            "ReviewedInputs" if hook_type == "PreExecution" else "ReviewedOutputs"
+        )
+
+        msgs = state.messages.copy()
+        if not msgs or reviewed_field not in escalation_result:
+            return {}
+
+        last_message = msgs[-1]
+
+        if hook_type == "PreExecution":
+            reviewed_content = escalation_result[reviewed_field]
+            if reviewed_content:
+                last_message.content = json.loads(reviewed_content)
+        else:
+            reviewed_outputs_json = escalation_result[reviewed_field]
+            if not reviewed_outputs_json:
+                return {}
+
+            content_list = json.loads(reviewed_outputs_json)
+            if not content_list:
+                return {}
+
+            ai_message: AIMessage = last_message  # type: ignore[assignment]
+            content_index = 0
+
+            if ai_message.tool_calls:
+                tool_calls = list(ai_message.tool_calls)
+                for tool_call in tool_calls:
+                    args = (
+                        tool_call["args"]
+                        if isinstance(tool_call, dict)
+                        else tool_call.args
+                    )
+                    if (
+                        isinstance(args, dict)
+                        and "content" in args
+                        and args["content"] is not None
+                    ):
+                        if content_index < len(content_list):
+                            updated_content = json.loads(content_list[content_index])
+                            args["content"] = updated_content
+                            if isinstance(tool_call, dict):
+                                tool_call["args"] = args
+                            else:
+                                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]
+
+        return Command(update={"messages": msgs})
+    except Exception as e:
+        raise AgentTerminationException(
+            code=UiPathErrorCode.EXECUTION_ERROR,
+            title="Escalation rejected",
+            detail=str(e)
+        ) from e
+
+
+def _extract_escalation_content(
+    state: AgentGuardrailsGraphState,
+    scope: GuardrailScope,
+    hook_type: Literal["PreExecution", "PostExecution"],
+) -> str:
+    """Extract escalation content from state based on guardrail scope and hook type.
+
+    Args:
+        state: The current agent graph state.
+        scope: The guardrail scope (LLM/AGENT/TOOL).
+        hook_type: The hook type ("PreExecution" or "PostExecution").
+
+    Returns:
+        For non-LLM scope: Empty string.
+        For LLM PreExecution: JSON string with message content.
+        For LLM PostExecution: JSON array with tool call content and message content.
+    """
+    if scope != GuardrailScope.LLM:
+        return ""
+
+    if not state.messages:
+        raise AgentTerminationException(
+            code=UiPathErrorCode.EXECUTION_ERROR,
+            title="Invalid state message",
+        )
+
+    last_message = state.messages[-1]
+    if hook_type == "PreExecution":
+        content = _message_text(last_message)
+        return json.dumps(content) if content else ""
+
+    ai_message: AIMessage = last_message  # type: ignore[assignment]
+    content_list: list[str] = []
+
+    if ai_message.tool_calls:
+        for tool_call in ai_message.tool_calls:
+            args = tool_call["args"] if isinstance(tool_call, dict) else 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)
+
+
+def _hook_type_to_tool_field(
+    hook_type: Literal["PreExecution", "PostExecution"],
+) -> str:
+    """Convert hook type to tool field name.
+
+    Args:
+        hook_type: The hook type ("PreExecution" or "PostExecution").
+
+    Returns:
+        "ToolInputs" for "PreExecution", "ToolOutputs" for "PostExecution".
+    """
+    return "ToolInputs" if hook_type == "PreExecution" else "ToolOutputs"

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

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+import logging
+import re
+from typing import Literal, Dict, Any
+
+from uipath.platform.guardrails import CustomGuardrail, BuiltInValidatorGuardrail, GuardrailScope
+
+from .base_action import GuardrailAction, GuardrailActionNode
+from ..types import AgentGuardrailsGraphState
+
+
+class LogAction(GuardrailAction):
+    """Action that logs guardrail violations and continues."""
+
+    def __init__(self, level: int = logging.WARNING) -> None:
+        self.level = level
+
+    def action_node(
+        self,
+        *,
+        guardrail: CustomGuardrail | BuiltInValidatorGuardrail,
+        scope: GuardrailScope,
+        execution_stage: Literal["PreExecution", "PostExecution"],
+    ) -> GuardrailActionNode:
+        sanitized = re.sub(r"\W+", "_", getattr(guardrail, "name", "guardrail")).strip(
+            "_"
+        )
+        node_name = f"{sanitized}_{execution_stage.lower()}_{scope.lower()}_log"
+
+        # TODO: add complete implementation for Log action
+        async def _node(_state: AgentGuardrailsGraphState) -> Dict[str, Any]:
+            print(
+                self.level,
+                "Guardrail '%s' failed at %s %s: %s",
+                guardrail.name,
+                execution_stage,
+                scope.value if hasattr(scope, "value") else str(scope),
+                _state.guardrail_validation_result,
+            )
+            return {}
+
+        return node_name, _node