Merge pull request #312 from UiPath/AL-228-log-block-guardrail-agent

feat(api): add support for OOTB guardrails at Agent level [AL-228, AL-241]

Author: valentinabojan

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-langchain"
-version = "0.1.29"
+version = "0.1.30"
 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/__init__.py

@@ -1,21 +1,5 @@
-from .guardrail_nodes import (
-    create_agent_guardrail_node,
-    create_llm_guardrail_node,
-    create_tool_guardrail_node,
-)
 from .guardrails_factory import build_guardrails_with_actions
-from .guardrails_subgraph import (
-    create_agent_guardrails_subgraph,
-    create_llm_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",
     "build_guardrails_with_actions",
 ]

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

@@ -6,7 +6,7 @@
 from uipath_langchain.agent.guardrails.types import ExecutionStage
 
 from ...exceptions import AgentTerminationException
-from ..types import AgentGuardrailsGraphState
+from ...react.types import AgentGuardrailsGraphState
 from .base_action import GuardrailAction, GuardrailActionNode
 
 

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

@@ -14,8 +14,9 @@
 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 get_message_content
 from .base_action import GuardrailAction, GuardrailActionNode
 
 
@@ -229,7 +230,7 @@ def _process_llm_escalation_response(
                 if content_list:
                     last_message.content = content_list[-1]
 
-        return Command[Any](update={"messages": msgs})
+        return Command(update={"messages": msgs})
     except Exception as e:
         raise AgentTerminationException(
             code=UiPathErrorCode.EXECUTION_ERROR,
@@ -311,7 +312,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,
@@ -377,7 +378,7 @@ def _extract_llm_escalation_content(
         if isinstance(last_message, ToolMessage):
             return last_message.content
 
-        content = _message_text(last_message)
+        content = get_message_content(last_message)
         return json.dumps(content) if content else ""
 
     # For AI messages, process tool calls if present
@@ -395,14 +396,14 @@ def _extract_llm_escalation_content(
                 ):
                     content_list.append(json.dumps(args["content"]))
 
-        message_content = _message_text(last_message)
+        message_content = get_message_content(last_message)
         if message_content:
             content_list.append(message_content)
 
         return json.dumps(content_list)
 
     # Fallback for other message types
-    return _message_text(last_message)
+    return get_message_content(last_message)
 
 
 def _extract_agent_escalation_content(

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

@@ -6,7 +6,7 @@
 
 from uipath_langchain.agent.guardrails.types import ExecutionStage
 
-from ..types import AgentGuardrailsGraphState
+from ...react.types import AgentGuardrailsGraphState
 from .base_action import GuardrailAction, GuardrailActionNode
 
 logger = logging.getLogger(__name__)

src/uipath_langchain/agent/guardrails/guardrail_nodes.py

@@ -3,7 +3,7 @@
 import re
 from typing import Any, Callable
 
-from langchain_core.messages import AIMessage, AnyMessage, HumanMessage, SystemMessage
+from langchain_core.messages import AIMessage
 from langgraph.types import Command
 from uipath.platform import UiPath
 from uipath.platform.guardrails import (
@@ -12,18 +12,12 @@
 )
 
 from uipath_langchain.agent.guardrails.types import ExecutionStage
-
-from .types import AgentGuardrailsGraphState
+from uipath_langchain.agent.guardrails.utils import get_message_content
+from uipath_langchain.agent.react.types import AgentGuardrailsGraphState
 
 logger = logging.getLogger(__name__)
 
 
-def _message_text(msg: AnyMessage) -> str:
-    if isinstance(msg, (HumanMessage, SystemMessage)):
-        return msg.content if isinstance(msg.content, str) else str(msg.content)
-    return str(getattr(msg, "content", "")) if hasattr(msg, "content") else ""
-
-
 def _create_guardrail_node(
     guardrail: BaseGuardrail,
     scope: GuardrailScope,
@@ -70,7 +64,7 @@ def create_llm_guardrail_node(
     def _payload_generator(state: AgentGuardrailsGraphState) -> str:
         if not state.messages:
             return ""
-        return _message_text(state.messages[-1])
+        return get_message_content(state.messages[-1])
 
     return _create_guardrail_node(
         guardrail,
@@ -82,17 +76,35 @@ def _payload_generator(state: AgentGuardrailsGraphState) -> str:
     )
 
 
-def create_agent_guardrail_node(
+def create_agent_init_guardrail_node(
     guardrail: BaseGuardrail,
     execution_stage: ExecutionStage,
     success_node: str,
     failure_node: str,
 ) -> tuple[str, Callable[[AgentGuardrailsGraphState], Any]]:
-    # To be implemented in future PR
     def _payload_generator(state: AgentGuardrailsGraphState) -> str:
         if not state.messages:
             return ""
-        return _message_text(state.messages[-1])
+        return get_message_content(state.messages[-1])
+
+    return _create_guardrail_node(
+        guardrail,
+        GuardrailScope.AGENT,
+        execution_stage,
+        _payload_generator,
+        success_node,
+        failure_node,
+    )
+
+
+def create_agent_terminate_guardrail_node(
+    guardrail: BaseGuardrail,
+    execution_stage: ExecutionStage,
+    success_node: str,
+    failure_node: str,
+) -> tuple[str, Callable[[AgentGuardrailsGraphState], Any]]:
+    def _payload_generator(state: AgentGuardrailsGraphState) -> str:
+        return str(state.agent_result)
 
     return _create_guardrail_node(
         guardrail,
@@ -161,7 +173,7 @@ def _payload_generator(state: AgentGuardrailsGraphState) -> str:
                     if args is not None:
                         return json.dumps(args)
 
-        return _message_text(state.messages[-1])
+        return get_message_content(state.messages[-1])
 
     return _create_guardrail_node(
         guardrail,

src/uipath_langchain/agent/guardrails/types.py

@@ -1,16 +1,4 @@
 from enum import Enum
-from typing import Annotated, Optional
-
-from langchain_core.messages import AnyMessage
-from langgraph.graph.message import add_messages
-from pydantic import BaseModel
-
-
-class AgentGuardrailsGraphState(BaseModel):
-    """Agent Guardrails Graph state for guardrail subgraph."""
-
-    messages: Annotated[list[AnyMessage], add_messages] = []
-    guardrail_validation_result: Optional[str] = None
 
 
 class ExecutionStage(str, Enum):

src/uipath_langchain/agent/guardrails/utils.py

@@ -0,0 +1,7 @@
+from langchain_core.messages import AnyMessage, HumanMessage, SystemMessage
+
+
+def get_message_content(msg: AnyMessage) -> str:
+    if isinstance(msg, (HumanMessage, SystemMessage)):
+        return msg.content if isinstance(msg.content, str) else str(msg.content)
+    return str(getattr(msg, "content", "")) if hasattr(msg, "content") else ""

src/uipath_langchain/agent/react/agent.py

@@ -9,10 +9,14 @@
 from pydantic import BaseModel
 from uipath.platform.guardrails import BaseGuardrail
 
-from ..guardrails import create_llm_guardrails_subgraph
 from ..guardrails.actions import GuardrailAction
-from ..guardrails.guardrails_subgraph import create_tools_guardrails_subgraph
 from ..tools import create_tool_node
+from .guardrails.guardrails_subgraph import (
+    create_agent_init_guardrails_subgraph,
+    create_agent_terminate_guardrails_subgraph,
+    create_llm_guardrails_subgraph,
+    create_tools_guardrails_subgraph,
+)
 from .init_node import (
     create_init_node,
 )
@@ -86,12 +90,20 @@ def create_agent(
     builder: StateGraph[AgentGraphState, None, InputT, OutputT] = StateGraph(
         InnerAgentGraphState, input_schema=input_schema, output_schema=output_schema
     )
-    builder.add_node(AgentGraphNode.INIT, init_node)
+    init_with_guardrails_subgraph = create_agent_init_guardrails_subgraph(
+        (AgentGraphNode.GUARDED_INIT, init_node),
+        guardrails,
+    )
+    builder.add_node(AgentGraphNode.INIT, init_with_guardrails_subgraph)
 
     for tool_name, tool_node in tool_nodes_with_guardrails.items():
         builder.add_node(tool_name, tool_node)
 
-    builder.add_node(AgentGraphNode.TERMINATE, terminate_node)
+    terminate_with_guardrails_subgraph = create_agent_terminate_guardrails_subgraph(
+        (AgentGraphNode.GUARDED_TERMINATE, terminate_node),
+        guardrails,
+    )
+    builder.add_node(AgentGraphNode.TERMINATE, terminate_with_guardrails_subgraph)
 
     builder.add_edge(START, AgentGraphNode.INIT)
 

…/agent/guardrails/guardrails_subgraph.py …/react/guardrails/guardrails_subgraph.pysrc/uipath_langchain/agent/guardrails/guardrails_subgraph.py renamed to src/uipath_langchain/agent/react/guardrails/guardrails_subgraph.py src/uipath_langchain/agent/guardrails/guardrails_subgraph.py renamed to src/uipath_langchain/agent/react/guardrails/guardrails_subgraph.py

@@ -10,15 +10,21 @@
     GuardrailScope,
 )
 
-from uipath_langchain.agent.guardrails.types import ExecutionStage
-
-from .actions.base_action import GuardrailAction, GuardrailActionNode
-from .guardrail_nodes import (
-    create_agent_guardrail_node,
+from uipath_langchain.agent.guardrails.actions.base_action import (
+    GuardrailAction,
+    GuardrailActionNode,
+)
+from uipath_langchain.agent.guardrails.guardrail_nodes import (
+    create_agent_init_guardrail_node,
+    create_agent_terminate_guardrail_node,
     create_llm_guardrail_node,
     create_tool_guardrail_node,
 )
-from .types import AgentGuardrailsGraphState
+from uipath_langchain.agent.guardrails.types import ExecutionStage
+from uipath_langchain.agent.react.types import (
+    AgentGraphState,
+    AgentGuardrailsGraphState,
+)
 
 _VALIDATOR_ALLOWED_STAGES = {
     "prompt_injection": {ExecutionStage.PRE_EXECUTION},
@@ -232,32 +238,65 @@ def create_tools_guardrails_subgraph(
     return result
 
 
-def create_agent_guardrails_subgraph(
-    agent_node: tuple[str, Any],
+def create_agent_init_guardrails_subgraph(
+    init_node: tuple[str, Any],
     guardrails: Sequence[tuple[BaseGuardrail, GuardrailAction]] | None,
-    execution_stage: ExecutionStage,
 ):
-    """Create a subgraph for AGENT-scoped guardrails that applies checks at the specified stage.
-
-    This is intended for wrapping nodes like INIT or TERMINATE, where guardrails should run
-    either before (pre-execution) or after (post-execution) the node logic.
-    """
+    """Create a subgraph for INIT node that applies guardrails on the state messages."""
     applicable_guardrails = [
         (guardrail, _)
         for (guardrail, _) in (guardrails or [])
         if GuardrailScope.AGENT in guardrail.selector.scopes
     ]
     if applicable_guardrails is None or len(applicable_guardrails) == 0:
-        return agent_node[1]
+        return init_node[1]
 
     return _create_guardrails_subgraph(
-        main_inner_node=agent_node,
+        main_inner_node=init_node,
+        guardrails=applicable_guardrails,
+        scope=GuardrailScope.AGENT,
+        execution_stages=[ExecutionStage.POST_EXECUTION],
+        node_factory=create_agent_init_guardrail_node,
+    )
+
+
+def create_agent_terminate_guardrails_subgraph(
+    terminate_node: tuple[str, Any],
+    guardrails: Sequence[tuple[BaseGuardrail, GuardrailAction]] | None,
+):
+    """Create a subgraph for TERMINATE node that applies guardrails on the agent result."""
+    node_name, node_func = terminate_node
+
+    def terminate_wrapper(state: Any) -> dict[str, Any]:
+        # Call original terminate node
+        result = node_func(state)
+        # Store result in state
+        return {"agent_result": result, "messages": state.messages}
+
+    applicable_guardrails = [
+        (guardrail, _)
+        for (guardrail, _) in (guardrails or [])
+        if GuardrailScope.AGENT in guardrail.selector.scopes
+    ]
+    if applicable_guardrails is None or len(applicable_guardrails) == 0:
+        return terminate_node[1]
+
+    subgraph = _create_guardrails_subgraph(
+        main_inner_node=(node_name, terminate_wrapper),
         guardrails=applicable_guardrails,
         scope=GuardrailScope.AGENT,
-        execution_stages=[execution_stage],
-        node_factory=create_agent_guardrail_node,
+        execution_stages=[ExecutionStage.POST_EXECUTION],
+        node_factory=create_agent_terminate_guardrail_node,
     )
 
+    async def run_terminate_subgraph(
+        state: AgentGraphState,
+    ) -> dict[str, Any]:
+        result_state = await subgraph.ainvoke(state)
+        return result_state["agent_result"]
+
+    return run_terminate_subgraph
+
 
 def create_tool_guardrails_subgraph(
     tool_node: tuple[str, Any],