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

uipath_langchain/agent/guardrails/guardrails_factory.py

@@ -1,13 +1,28 @@
 import logging
-from typing import Sequence
+import re
+from typing import Callable, Sequence
 
 from uipath.agent.models.agent import (
+    AgentBooleanOperator,
+    AgentBooleanRule,
+    AgentCustomGuardrail,
     AgentGuardrail,
     AgentGuardrailBlockAction,
     AgentGuardrailEscalateAction,
     AgentGuardrailLogAction,
     AgentGuardrailSeverityLevel,
+    AgentNumberOperator,
+    AgentNumberRule,
     AgentUnknownGuardrail,
+    AgentWordOperator,
+    AgentWordRule,
+)
+from uipath.core.guardrails import (
+    BooleanRule,
+    DeterministicGuardrail,
+    NumberRule,
+    UniversalRule,
+    WordRule,
 )
 from uipath.platform.guardrails import BaseGuardrail
 
@@ -19,6 +34,180 @@ from uipath_langchain.agent.guardrails.actions import (
 )
 
 
+def _assert_value_not_none(value: str | None, operator: AgentWordOperator) -> str:
+    """Assert value is not None and return as string."""
+    assert value is not None, f"value cannot be None for {operator.name} operator"
+    return value
+
+
+def _create_word_rule_func(
+    operator: AgentWordOperator, value: str | None
+) -> Callable[[str], bool]:
+    """Create a callable function from AgentWordOperator and value.
+
+    Args:
+        operator: The word operator to convert.
+        value: The value to compare against (may be None for isEmpty/isNotEmpty).
+
+    Returns:
+        A callable that takes a string and returns a boolean.
+    """
+    match operator:
+        case AgentWordOperator.CONTAINS:
+            val = _assert_value_not_none(value, operator)
+            return lambda s: val.lower() in s.lower()
+        case AgentWordOperator.DOES_NOT_CONTAIN:
+            val = _assert_value_not_none(value, operator)
+            return lambda s: val.lower() not in s.lower()
+        case AgentWordOperator.EQUALS:
+            val = _assert_value_not_none(value, operator)
+            return lambda s: s == val
+        case AgentWordOperator.DOES_NOT_EQUAL:
+            val = _assert_value_not_none(value, operator)
+            return lambda s: s != val
+        case AgentWordOperator.STARTS_WITH:
+            val = _assert_value_not_none(value, operator)
+            return lambda s: s.startswith(val)
+        case AgentWordOperator.DOES_NOT_START_WITH:
+            val = _assert_value_not_none(value, operator)
+            return lambda s: not s.startswith(val)
+        case AgentWordOperator.ENDS_WITH:
+            val = _assert_value_not_none(value, operator)
+            return lambda s: s.endswith(val)
+        case AgentWordOperator.DOES_NOT_END_WITH:
+            val = _assert_value_not_none(value, operator)
+            return lambda s: not s.endswith(val)
+        case AgentWordOperator.IS_EMPTY:
+            return lambda s: len(s) == 0
+        case AgentWordOperator.IS_NOT_EMPTY:
+            return lambda s: len(s) > 0
+        case AgentWordOperator.MATCHES_REGEX:
+            val = _assert_value_not_none(value, operator)
+            pattern = re.compile(val)
+            return lambda s: bool(pattern.match(s))
+        case _:
+            raise ValueError(f"Unsupported word operator: {operator}")
+
+
+def _create_number_rule_func(
+    operator: AgentNumberOperator, value: float
+) -> Callable[[float], bool]:
+    """Create a callable function from AgentNumberOperator and value.
+
+    Args:
+        operator: The number operator to convert.
+        value: The value to compare against.
+
+    Returns:
+        A callable that takes a float and returns a boolean.
+    """
+    match operator:
+        case AgentNumberOperator.EQUALS:
+            return lambda n: n == value
+        case AgentNumberOperator.DOES_NOT_EQUAL:
+            return lambda n: n != value
+        case AgentNumberOperator.GREATER_THAN:
+            return lambda n: n > value
+        case AgentNumberOperator.GREATER_THAN_OR_EQUAL:
+            return lambda n: n >= value
+        case AgentNumberOperator.LESS_THAN:
+            return lambda n: n < value
+        case AgentNumberOperator.LESS_THAN_OR_EQUAL:
+            return lambda n: n <= value
+        case _:
+            raise ValueError(f"Unsupported number operator: {operator}")
+
+
+def _create_boolean_rule_func(
+    operator: AgentBooleanOperator, value: bool
+) -> Callable[[bool], bool]:
+    """Create a callable function from AgentBooleanOperator and value.
+
+    Args:
+        operator: The boolean operator to convert.
+        value: The value to compare against.
+
+    Returns:
+        A callable that takes a boolean and returns a boolean.
+    """
+    match operator:
+        case AgentBooleanOperator.EQUALS:
+            return lambda b: b == value
+        case _:
+            raise ValueError(f"Unsupported boolean operator: {operator}")
+
+
+def _convert_agent_rule_to_deterministic(
+    agent_rule: AgentWordRule | AgentNumberRule | AgentBooleanRule | UniversalRule,
+) -> WordRule | NumberRule | BooleanRule | UniversalRule:
+    """Convert an Agent rule to its Deterministic equivalent.
+
+    Args:
+        agent_rule: The agent rule to convert.
+
+    Returns:
+        The corresponding deterministic rule with a callable function.
+    """
+    if isinstance(agent_rule, UniversalRule):
+        # UniversalRule is already compatible
+        return agent_rule
+
+    if isinstance(agent_rule, AgentWordRule):
+        return WordRule(
+            rule_type="word",
+            field_selector=agent_rule.field_selector,
+            detects_violation=_create_word_rule_func(
+                agent_rule.operator, agent_rule.value
+            ),
+        )
+
+    if isinstance(agent_rule, AgentNumberRule):
+        return NumberRule(
+            rule_type="number",
+            field_selector=agent_rule.field_selector,
+            detects_violation=_create_number_rule_func(
+                agent_rule.operator, agent_rule.value
+            ),
+        )
+
+    if isinstance(agent_rule, AgentBooleanRule):
+        return BooleanRule(
+            rule_type="boolean",
+            field_selector=agent_rule.field_selector,
+            detects_violation=_create_boolean_rule_func(
+                agent_rule.operator, agent_rule.value
+            ),
+        )
+
+    raise ValueError(f"Unsupported agent rule type: {type(agent_rule)}")
+
+
+def _convert_agent_custom_guardrail_to_deterministic(
+    guardrail: AgentCustomGuardrail,
+) -> DeterministicGuardrail:
+    """Convert AgentCustomGuardrail to DeterministicGuardrail.
+
+    Args:
+        guardrail: The agent custom guardrail to convert.
+
+    Returns:
+        A DeterministicGuardrail with converted rules.
+    """
+    converted_rules = [
+        _convert_agent_rule_to_deterministic(rule) for rule in guardrail.rules
+    ]
+
+    return DeterministicGuardrail(
+        id=guardrail.id,
+        name=guardrail.name,
+        description=guardrail.description,
+        enabled_for_evals=guardrail.enabled_for_evals,
+        selector=guardrail.selector,
+        guardrail_type="custom",
+        rules=converted_rules,
+    )
+
+
 def build_guardrails_with_actions(
     guardrails: Sequence[AgentGuardrail] | None,
 ) -> list[tuple[BaseGuardrail, GuardrailAction]]:
@@ -38,10 +227,17 @@ def build_guardrails_with_actions(
         if isinstance(guardrail, AgentUnknownGuardrail):
             continue
 
+        # Convert AgentCustomGuardrail to DeterministicGuardrail
+        converted_guardrail: BaseGuardrail = guardrail
+        if isinstance(guardrail, AgentCustomGuardrail):
+            converted_guardrail = _convert_agent_custom_guardrail_to_deterministic(
+                guardrail
+            )
+
         action = guardrail.action
 
         if isinstance(action, AgentGuardrailBlockAction):
-            result.append((guardrail, BlockAction(action.reason)))
+            result.append((converted_guardrail, BlockAction(action.reason)))
         elif isinstance(action, AgentGuardrailLogAction):
             severity_level_map = {
                 AgentGuardrailSeverityLevel.ERROR: logging.ERROR,
@@ -51,14 +247,14 @@ def build_guardrails_with_actions(
             level = severity_level_map.get(action.severity_level, logging.INFO)
             result.append(
                 (
-                    guardrail,
+                    converted_guardrail,
                     LogAction(message=action.message, level=level),
                 )
             )
         elif isinstance(action, AgentGuardrailEscalateAction):
             result.append(
                 (
-                    guardrail,
+                    converted_guardrail,
                     EscalateAction(
                         app_name=action.app.name,
                         app_folder_path=action.app.folder_name,

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):

uipath_langchain/agent/guardrails/utils.py

@@ -0,0 +1,146 @@
+import json
+import logging
+from typing import Any
+
+from langchain_core.messages import (
+    AIMessage,
+    AnyMessage,
+    HumanMessage,
+    SystemMessage,
+    ToolMessage,
+)
+
+from uipath_langchain.agent.guardrails.types import ExecutionStage
+from uipath_langchain.agent.react.types import AgentGuardrailsGraphState
+
+logger = logging.getLogger(__name__)
+
+
+def _extract_tool_args_from_message(
+    message: AnyMessage, tool_name: str
+) -> dict[str, Any]:
+    """Extract tool call arguments from an AIMessage.
+
+    Args:
+        message: The message to extract from.
+        tool_name: Name of the tool to extract arguments from.
+
+    Returns:
+        Dict containing tool call arguments, or empty dict if not found.
+    """
+    if not isinstance(message, AIMessage):
+        return {}
+
+    if not message.tool_calls:
+        return {}
+
+    # Find the first tool call with matching name
+    for tool_call in message.tool_calls:
+        call_name = (
+            tool_call.get("name")
+            if isinstance(tool_call, dict)
+            else getattr(tool_call, "name", None)
+        )
+        if call_name == tool_name:
+            # Extract args from the tool call
+            args = (
+                tool_call.get("args")
+                if isinstance(tool_call, dict)
+                else getattr(tool_call, "args", None)
+            )
+            if args is not None:
+                # Args should already be a dict
+                if isinstance(args, dict):
+                    return args
+                # If it's a JSON string, parse it
+                if isinstance(args, str):
+                    try:
+                        parsed = json.loads(args)
+                        if isinstance(parsed, dict):
+                            return parsed
+                    except json.JSONDecodeError:
+                        logger.warning(
+                            "Failed to parse tool args as JSON for tool '%s': %s",
+                            tool_name,
+                            args[:100] if len(args) > 100 else args,
+                        )
+                        return {}
+
+    return {}
+
+
+def _extract_tool_input_data(
+    state: AgentGuardrailsGraphState, tool_name: str, execution_stage: ExecutionStage
+) -> dict[str, Any]:
+    """Extract tool call arguments as dict for deterministic guardrails.
+
+    Args:
+        state: The current agent graph state.
+        tool_name: Name of the tool to extract arguments from.
+        execution_stage: PRE_EXECUTION or POST_EXECUTION.
+
+    Returns:
+        Dict containing tool call arguments, or empty dict if not found.
+        - For PRE_EXECUTION: extracts from last message
+        - For POST_EXECUTION: extracts from second-to-last message
+    """
+    if not state.messages:
+        return {}
+
+    # For PRE_EXECUTION, look at last message
+    # For POST_EXECUTION, look at second-to-last message (before the ToolMessage)
+    if execution_stage == ExecutionStage.PRE_EXECUTION:
+        if len(state.messages) < 1:
+            return {}
+        message = state.messages[-1]
+    else:  # POST_EXECUTION
+        if len(state.messages) < 2:
+            return {}
+        message = state.messages[-2]
+
+    return _extract_tool_args_from_message(message, tool_name)
+
+
+def _extract_tool_output_data(state: AgentGuardrailsGraphState) -> dict[str, Any]:
+    """Extract tool execution output as dict for POST_EXECUTION deterministic guardrails.
+
+    Args:
+        state: The current agent graph state.
+
+    Returns:
+        Dict containing tool output. If output is not valid JSON, wraps it in {"output": content}.
+    """
+    if not state.messages:
+        return {}
+
+    last_message = state.messages[-1]
+    if not isinstance(last_message, ToolMessage):
+        return {}
+
+    content = last_message.content
+    if not content:
+        return {}
+
+    # Try to parse as JSON first
+    if isinstance(content, str):
+        try:
+            parsed = json.loads(content)
+            if isinstance(parsed, dict):
+                return parsed
+            else:
+                # JSON array or primitive - wrap it
+                return {"output": parsed}
+        except json.JSONDecodeError:
+            logger.warning("Tool output is not valid JSON")
+            return {"output": content}
+    elif isinstance(content, dict):
+        return content
+    else:
+        # List or other type
+        return {"output": content}
+
+
+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 ""

uipath_langchain/agent/react/agent.py

@@ -9,10 +9,14 @@ from langgraph.graph import StateGraph
 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,
 )
@@ -20,7 +24,7 @@ from .llm_node import (
     create_llm_node,
 )
 from .router import (
-    route_agent,
+    create_route_agent,
 )
 from .terminate_node import (
     create_terminate_node,
@@ -54,7 +58,7 @@ def create_agent(
     config: AgentGraphConfig | None = None,
     guardrails: Sequence[tuple[BaseGuardrail, GuardrailAction]] | None = None,
 ) -> StateGraph[AgentGraphState, None, InputT, OutputT]:
-    """Build agent graph with INIT -> AGENT(subgraph) <-> TOOLS loop, terminated by control flow tools.
+    """Build agent graph with INIT -> AGENT (subgraph) <-> TOOLS loop, terminated by control flow tools.
 
     The AGENT node is a subgraph that runs:
     - before-agent guardrail middlewares
@@ -86,16 +90,24 @@ 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)
 
-    llm_node = create_llm_node(model, llm_tools)
+    llm_node = create_llm_node(model, llm_tools, config.thinking_messages_limit)
     llm_with_guardrails_subgraph = create_llm_guardrails_subgraph(
         (AgentGraphNode.LLM, llm_node), guardrails
     )
@@ -103,6 +115,7 @@ def create_agent(
     builder.add_edge(AgentGraphNode.INIT, AgentGraphNode.AGENT)
 
     tool_node_names = list(tool_nodes_with_guardrails.keys())
+    route_agent = create_route_agent(config.thinking_messages_limit)
     builder.add_conditional_edges(
         AgentGraphNode.AGENT,
         route_agent,

uipath_langchain/agent/react/constants.py

@@ -1,2 +1 @@
-# Agent routing configuration
-MAX_SUCCESSIVE_COMPLETIONS = 1
+MAX_CONSECUTIVE_THINKING_MESSAGES = 0

uipath_langchain/agent/{guardrails → react/guardrails}/guardrails_subgraph.py

@@ -10,15 +10,21 @@ from uipath.platform.guardrails import (
     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],

uipath_langchain/agent/react/llm_node.py

@@ -1,34 +1,65 @@
-"""LLM node implementation for LangGraph."""
+"""LLM node for ReAct Agent graph."""
 
-from typing import Sequence
+from typing import Literal, Sequence
 
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import AIMessage, AnyMessage
 from langchain_core.tools import BaseTool
 
-from .constants import MAX_SUCCESSIVE_COMPLETIONS
+from .constants import MAX_CONSECUTIVE_THINKING_MESSAGES
 from .types import AgentGraphState
-from .utils import count_successive_completions
+from .utils import count_consecutive_thinking_messages
+
+OPENAI_COMPATIBLE_CHAT_MODELS = (
+    "UiPathChatOpenAI",
+    "AzureChatOpenAI",
+    "ChatOpenAI",
+    "UiPathChat",
+    "UiPathAzureChatOpenAI",
+)
+
+
+def _get_required_tool_choice_by_model(
+    model: BaseChatModel,
+) -> Literal["required", "any"]:
+    """Get the appropriate tool_choice value to enforce tool usage based on model type.
+
+    "required" - OpenAI compatible required tool_choice value
+    "any" - Vertex and Bedrock parameter for required tool_choice value
+    """
+    model_class_name = model.__class__.__name__
+    if model_class_name in OPENAI_COMPATIBLE_CHAT_MODELS:
+        return "required"
+    return "any"
 
 
 def create_llm_node(
     model: BaseChatModel,
     tools: Sequence[BaseTool] | None = None,
+    thinking_messages_limit: int = MAX_CONSECUTIVE_THINKING_MESSAGES,
 ):
-    """Invoke LLM with tools and dynamically control tool_choice based on successive completions.
+    """Create LLM node with dynamic tool_choice enforcement.
 
-    When successive completions reach the limit, tool_choice is set to "required" to force
-    the LLM to use a tool and prevent infinite reasoning loops.
+    Controls when to force tool usage based on consecutive thinking steps
+    to prevent infinite loops and ensure progress.
+
+    Args:
+        model: The chat model to use
+        tools: Available tools to bind
+        thinking_messages_limit: Max consecutive LLM responses without tool calls
+            before enforcing tool usage. 0 = force tools every time.
     """
     bindable_tools = list(tools) if tools else []
     base_llm = model.bind_tools(bindable_tools) if bindable_tools else model
+    tool_choice_required_value = _get_required_tool_choice_by_model(model)
 
     async def llm_node(state: AgentGraphState):
         messages: list[AnyMessage] = state.messages
 
-        successive_completions = count_successive_completions(messages)
-        if successive_completions >= MAX_SUCCESSIVE_COMPLETIONS:
-            llm = base_llm.bind(tool_choice="required")
+        consecutive_thinking_messages = count_consecutive_thinking_messages(messages)
+
+        if bindable_tools and consecutive_thinking_messages >= thinking_messages_limit:
+            llm = base_llm.bind(tool_choice=tool_choice_required_value)
         else:
             llm = base_llm