feat: move guardrail state to inner_state; add inputs to guardrail state (#420)

Author: andreiancuta-uipath

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

@@ -83,7 +83,7 @@ async def _node(
                 "GuardrailDescription": guardrail.description,
                 "Component": scope.name.lower(),
                 "ExecutionStage": _execution_stage_to_string(execution_stage),
-                "GuardrailResult": state.guardrail_validation_result,
+                "GuardrailResult": state.inner_state.guardrail_validation_result,
             }
 
             # Add tenant and trace URL if base_url is configured
@@ -275,7 +275,7 @@ def _process_agent_escalation_response(
             return Command(update={"messages": msgs})
 
         # POST_EXECUTION: update agent_result
-        return Command(update={"agent_result": parsed})
+        return Command(update={"inner_state": {"agent_result": parsed}})
     except Exception as e:
         raise AgentTerminationException(
             code=UiPathErrorCode.EXECUTION_ERROR,
@@ -519,7 +519,7 @@ def _extract_agent_escalation_content(
     if execution_stage == ExecutionStage.PRE_EXECUTION:
         return get_message_content(cast(AnyMessage, message))
 
-    output_content = state.agent_result or ""
+    output_content = state.inner_state.agent_result or ""
     return json.dumps(output_content)
 
 

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

@@ -49,7 +49,7 @@ def action_node(
         async def _node(_state: AgentGuardrailsGraphState) -> dict[str, Any]:
             message = (
                 self.message
-                or f"Guardrail [{guardrail.name}] validation failed for [{scope.name}] [{execution_stage.name}] with the following reason: {_state.guardrail_validation_result}"
+                or f"Guardrail [{guardrail.name}] validation failed for [{scope.name}] [{execution_stage.name}] with the following reason: {_state.inner_state.guardrail_validation_result}"
             )
 
             logger.log(self.level, message)

src/uipath_langchain/agent/guardrails/guardrail_nodes.py

@@ -106,12 +106,17 @@ def _create_validation_command(
         AgentTerminationException: If the result is neither PASSED nor VALIDATION_FAILED.
     """
     if guardrail_result.result == GuardrailValidationResultType.PASSED:
-        return Command(goto=success_node, update={"guardrail_validation_result": None})
+        return Command(
+            goto=success_node,
+            update={"inner_state": {"guardrail_validation_result": None}},
+        )
 
     if guardrail_result.result == GuardrailValidationResultType.VALIDATION_FAILED:
         return Command(
             goto=failure_node,
-            update={"guardrail_validation_result": guardrail_result.reason},
+            update={
+                "inner_state": {"guardrail_validation_result": guardrail_result.reason}
+            },
         )
 
     # For other results (FEATURE_DISABLED, ENTITLEMENTS_MISSING, etc.), interrupt execution
@@ -260,7 +265,7 @@ def create_agent_terminate_guardrail_node(
     failure_node: str,
 ) -> tuple[str, Callable[[AgentGuardrailsGraphState], Any]]:
     def _payload_generator(state: AgentGuardrailsGraphState) -> str:
-        return str(state.agent_result)
+        return str(state.inner_state.agent_result)
 
     return _create_guardrail_node(
         guardrail,

src/uipath_langchain/agent/react/agent.py

@@ -1,5 +1,5 @@
 import os
-from typing import Callable, Sequence, Type, TypeVar, cast
+from typing import Callable, Sequence, Type, TypeVar
 
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import HumanMessage, SystemMessage
@@ -30,23 +30,17 @@
     create_terminate_node,
 )
 from .tools import create_flow_control_tools
-from .types import AgentGraphConfig, AgentGraphNode, AgentGraphState
+from .types import (
+    AgentGraphConfig,
+    AgentGraphNode,
+    AgentGraphState,
+)
+from .utils import create_state_with_input
 
 InputT = TypeVar("InputT", bound=BaseModel)
 OutputT = TypeVar("OutputT", bound=BaseModel)
 
 
-def create_state_with_input(input_schema: Type[InputT]):
-    CompleteAgentGraphState = type(
-        "CompleteAgentGraphState",
-        (AgentGraphState, input_schema),
-        {},
-    )
-
-    cast(type[BaseModel], CompleteAgentGraphState).model_rebuild()
-    return CompleteAgentGraphState
-
-
 def create_agent(
     model: BaseChatModel,
     tools: Sequence[BaseTool],
@@ -84,7 +78,7 @@ def create_agent(
 
     tool_nodes = create_tool_node(agent_tools)
     tool_nodes_with_guardrails = create_tools_guardrails_subgraph(
-        tool_nodes, guardrails
+        tool_nodes, guardrails, input_schema=input_schema
     )
     terminate_node = create_terminate_node(output_schema, config.is_conversational)
 
@@ -98,6 +92,7 @@ def create_agent(
     init_with_guardrails_subgraph = create_agent_init_guardrails_subgraph(
         (AgentGraphNode.GUARDED_INIT, init_node),
         guardrails,
+        input_schema=input_schema,
     )
     builder.add_node(AgentGraphNode.INIT, init_with_guardrails_subgraph)
 
@@ -107,6 +102,7 @@ def create_agent(
     terminate_with_guardrails_subgraph = create_agent_terminate_guardrails_subgraph(
         (AgentGraphNode.GUARDED_TERMINATE, terminate_node),
         guardrails,
+        input_schema=input_schema,
     )
     builder.add_node(AgentGraphNode.TERMINATE, terminate_with_guardrails_subgraph)
 
@@ -116,7 +112,7 @@ def create_agent(
         model, llm_tools, config.thinking_messages_limit, config.is_conversational
     )
     llm_with_guardrails_subgraph = create_llm_guardrails_subgraph(
-        (AgentGraphNode.LLM, llm_node), guardrails
+        (AgentGraphNode.LLM, llm_node), guardrails, input_schema=input_schema
     )
     builder.add_node(AgentGraphNode.AGENT, llm_with_guardrails_subgraph)
     builder.add_edge(AgentGraphNode.INIT, AgentGraphNode.AGENT)

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

@@ -1,9 +1,10 @@
 from functools import partial
-from typing import Any, Callable, Mapping, Sequence
+from typing import Any, Callable, Mapping, Sequence, TypeVar
 
 from langgraph._internal._runnable import RunnableCallable
 from langgraph.constants import END, START
 from langgraph.graph import StateGraph
+from pydantic import BaseModel
 from uipath.core.guardrails import DeterministicGuardrail
 from uipath.platform.guardrails import (
     BaseGuardrail,
@@ -26,6 +27,7 @@
     AgentGraphState,
     AgentGuardrailsGraphState,
 )
+from uipath_langchain.agent.react.utils import create_guardrails_state_with_input
 
 _VALIDATOR_ALLOWED_STAGES = {
     "prompt_injection": {ExecutionStage.PRE_EXECUTION},
@@ -65,6 +67,7 @@ def _create_guardrails_subgraph(
         ],
         GuardrailActionNode,
     ] = create_llm_guardrail_node,
+    input_schema: type[BaseModel] | None = None,
 ):
     """Build a subgraph that enforces guardrails around an inner node.
 
@@ -83,7 +86,8 @@ def _create_guardrails_subgraph(
     """
     inner_name, inner_node = main_inner_node
 
-    subgraph = StateGraph(AgentGuardrailsGraphState)
+    CompleteAgentGuardrailsGraphState = create_guardrails_state_with_input(input_schema)
+    subgraph = StateGraph(CompleteAgentGuardrailsGraphState)
 
     subgraph.add_node(inner_name, inner_node)
 
@@ -203,12 +207,14 @@ def _build_guardrail_node_chain(
 def create_llm_guardrails_subgraph(
     llm_node: tuple[str, Any],
     guardrails: Sequence[tuple[BaseGuardrail, GuardrailAction]] | None,
+    input_schema: type[BaseModel] | None = None,
 ):
     """Create a guarded LLM node.
 
     Args:
         llm_node: Tuple of (node_name, node_callable) for the LLM node.
         guardrails: Optional sequence of (guardrail, action) tuples.
+        input_schema: Optional input schema to include in state.
 
     Returns:
         Either the original node callable (if no applicable guardrails) or a compiled
@@ -229,17 +235,20 @@ def create_llm_guardrails_subgraph(
         scope=GuardrailScope.LLM,
         execution_stages=[ExecutionStage.PRE_EXECUTION, ExecutionStage.POST_EXECUTION],
         node_factory=create_llm_guardrail_node,
+        input_schema=input_schema,
     )
 
 
 def create_tools_guardrails_subgraph(
     tool_nodes: Mapping[str, RunnableCallable],
     guardrails: Sequence[tuple[BaseGuardrail, GuardrailAction]] | None,
+    input_schema: type[BaseModel] | None = None,
 ) -> dict[str, RunnableCallable]:
     """Create tool nodes with guardrails applied.
     Args:
         tool_nodes: Mapping of tool name to a LangGraph `ToolNode`.
         guardrails: Optional sequence of (guardrail, action) tuples.
+        input_schema: Optional input schema to include in state.
 
     Returns:
         A mapping of tool name to either the original `ToolNode` or a compiled subgraph
@@ -250,6 +259,7 @@ def create_tools_guardrails_subgraph(
         subgraph = create_tool_guardrails_subgraph(
             (tool_name, tool_node),
             guardrails,
+            input_schema=input_schema,
         )
         result[tool_name] = subgraph
 
@@ -259,6 +269,7 @@ def create_tools_guardrails_subgraph(
 def create_agent_init_guardrails_subgraph(
     init_node: tuple[str, Any],
     guardrails: Sequence[tuple[BaseGuardrail, GuardrailAction]] | None,
+    input_schema: type[BaseModel] | None = None,
 ) -> Any:
     """Create a subgraph for the INIT node and apply AGENT guardrails after INIT.
 
@@ -269,6 +280,7 @@ def create_agent_init_guardrails_subgraph(
     Args:
         init_node: Tuple of (node_name, node_callable) for the INIT node.
         guardrails: Optional sequence of (guardrail, action) tuples.
+        input_schema: Optional input schema to include in state.
 
     Returns:
         Either the original node callable (if no applicable guardrails) or a compiled
@@ -287,7 +299,8 @@ def create_agent_init_guardrails_subgraph(
         return init_node[1]
 
     inner_name, inner_node = init_node
-    subgraph = StateGraph(AgentGuardrailsGraphState)
+    CompleteAgentGuardrailsGraphState = create_guardrails_state_with_input(input_schema)
+    subgraph = StateGraph(CompleteAgentGuardrailsGraphState)
     subgraph.add_node(inner_name, inner_node)
     subgraph.add_edge(START, inner_name)
 
@@ -307,15 +320,15 @@ def create_agent_init_guardrails_subgraph(
 def create_agent_terminate_guardrails_subgraph(
     terminate_node: tuple[str, Any],
     guardrails: Sequence[tuple[BaseGuardrail, GuardrailAction]] | None,
+    input_schema: type[BaseModel] | None = 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}
+        return {"inner_state": {"agent_result": result}}
 
     applicable_guardrails = [
         (guardrail, _)
@@ -332,26 +345,31 @@ def terminate_wrapper(state: Any) -> dict[str, Any]:
         scope=GuardrailScope.AGENT,
         execution_stages=[ExecutionStage.POST_EXECUTION],
         node_factory=create_agent_terminate_guardrail_node,
+        input_schema=input_schema,
     )
 
+    StateT = TypeVar("StateT", bound=AgentGraphState)
+
     async def run_terminate_subgraph(
-        state: AgentGraphState,
+        state: StateT,
     ) -> dict[str, Any]:
         result_state = await subgraph.ainvoke(state)
-        return result_state["agent_result"]
+        return result_state["inner_state"].agent_result
 
     return run_terminate_subgraph
 
 
 def create_tool_guardrails_subgraph(
     tool_node: tuple[str, Any],
     guardrails: Sequence[tuple[BaseGuardrail, GuardrailAction]] | None,
+    input_schema: type[BaseModel] | None = None,
 ):
     """Create a guarded tool node.
 
     Args:
         tool_node: Tuple of (tool_name, tool_node_callable).
         guardrails: Optional sequence of (guardrail, action) tuples.
+        input_schema: Optional input schema to include in state.
 
     Returns:
         Either the original tool node callable (if no matching guardrails) or a compiled
@@ -374,4 +392,5 @@ def create_tool_guardrails_subgraph(
         scope=GuardrailScope.TOOL,
         execution_stages=[ExecutionStage.PRE_EXECUTION, ExecutionStage.POST_EXECUTION],
         node_factory=partial(create_tool_guardrail_node, tool_name=tool_name),
+        input_schema=input_schema,
     )

src/uipath_langchain/agent/react/types.py

@@ -26,6 +26,13 @@ class InnerAgentGraphState(BaseModel):
     termination: AgentTermination | None = None
 
 
+class InnerAgentGuardrailsGraphState(InnerAgentGraphState):
+    """Extended inner state for guardrails subgraph."""
+
+    guardrail_validation_result: Optional[str] = None
+    agent_result: Optional[dict[str, Any]] = None
+
+
 class AgentGraphState(BaseModel):
     """Agent Graph state for standard loop execution."""
 
@@ -38,8 +45,9 @@ class AgentGraphState(BaseModel):
 class AgentGuardrailsGraphState(AgentGraphState):
     """Agent Guardrails Graph state for guardrail subgraph."""
 
-    guardrail_validation_result: Optional[str] = None
-    agent_result: Optional[dict[str, Any]] = None
+    inner_state: Annotated[InnerAgentGuardrailsGraphState, merge_objects] = Field(
+        default_factory=InnerAgentGuardrailsGraphState
+    )
 
 
 class AgentGraphNode(StrEnum):

src/uipath_langchain/agent/react/utils.py

@@ -1,12 +1,16 @@
 """ReAct Agent loop utilities."""
 
-from typing import Any, Sequence
+from typing import Any, Sequence, TypeVar, cast
 
 from langchain_core.messages import AIMessage, BaseMessage
 from pydantic import BaseModel
 from uipath.agent.react import END_EXECUTION_TOOL
 
 from uipath_langchain.agent.react.jsonschema_pydantic_converter import create_model
+from uipath_langchain.agent.react.types import (
+    AgentGraphState,
+    AgentGuardrailsGraphState,
+)
 
 
 def resolve_input_model(
@@ -48,3 +52,41 @@ def count_consecutive_thinking_messages(messages: Sequence[BaseMessage]) -> int:
         count += 1
 
     return count
+
+
+InputT = TypeVar("InputT", bound=BaseModel)
+GraphStateT = TypeVar("GraphStateT", bound=BaseModel)
+
+
+def _create_state_model_with_input(
+    state_model: type[GraphStateT],
+    input_schema: type[InputT] | None,
+    model_name: str = "CompleteStateModel",
+) -> type[GraphStateT]:
+    if input_schema is None:
+        return state_model
+
+    CompleteStateModel = type(
+        model_name,
+        (state_model, input_schema),
+        {},
+    )
+
+    cast(type[GraphStateT], CompleteStateModel).model_rebuild()
+    return CompleteStateModel
+
+
+def create_state_with_input(input_schema: type[InputT] | None) -> type[AgentGraphState]:
+    return _create_state_model_with_input(
+        AgentGraphState, input_schema, model_name="CompleteAgentGraphState"
+    )
+
+
+def create_guardrails_state_with_input(
+    input_schema: type[InputT] | None,
+) -> type[AgentGuardrailsGraphState]:
+    return _create_state_model_with_input(
+        AgentGuardrailsGraphState,
+        input_schema,
+        model_name="CompleteAgentGuardrailsGraphState",
+    )