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

uipath_langchain/agent/guardrails/actions/filter_action.py

@@ -0,0 +1,290 @@
+import re
+from typing import Any
+
+from langchain_core.messages import AIMessage, ToolMessage
+from langgraph.types import Command
+from uipath.core.guardrails.guardrails import FieldReference, FieldSource
+from uipath.platform.guardrails import BaseGuardrail, GuardrailScope
+from uipath.runtime.errors import UiPathErrorCategory, UiPathErrorCode
+
+from uipath_langchain.agent.guardrails.types import ExecutionStage
+
+from ...exceptions import AgentTerminationException
+from ...react.types import AgentGuardrailsGraphState
+from .base_action import GuardrailAction, GuardrailActionNode
+
+
+class FilterAction(GuardrailAction):
+    """Action that filters inputs/outputs on guardrail failure.
+
+    For Tool scope, this action removes specified fields from tool call arguments.
+    For AGENT and LLM scopes, this action raises an exception as it's not supported yet.
+    """
+
+    def __init__(self, fields: list[FieldReference] | None = None):
+        """Initialize FilterAction with fields to filter.
+
+        Args:
+            fields: List of FieldReference objects specifying which fields to filter.
+        """
+        self.fields = fields or []
+
+    def action_node(
+        self,
+        *,
+        guardrail: BaseGuardrail,
+        scope: GuardrailScope,
+        execution_stage: ExecutionStage,
+        guarded_component_name: str,
+    ) -> GuardrailActionNode:
+        """Create a guardrail action node that performs filtering.
+
+        Args:
+            guardrail: The guardrail responsible for the validation.
+            scope: The scope in which the guardrail applies.
+            execution_stage: Whether this runs before or after execution.
+            guarded_component_name: Name of the guarded component.
+
+        Returns:
+            A tuple containing the node name and the async node callable.
+        """
+        raw_node_name = f"{scope.name}_{execution_stage.name}_{guardrail.name}_filter"
+        node_name = re.sub(r"\W+", "_", raw_node_name.lower()).strip("_")
+
+        async def _node(
+            _state: AgentGuardrailsGraphState,
+        ) -> dict[str, Any] | Command[Any]:
+            if scope == GuardrailScope.TOOL:
+                return _filter_tool_fields(
+                    _state,
+                    self.fields,
+                    execution_stage,
+                    guarded_component_name,
+                    guardrail.name,
+                )
+
+            raise AgentTerminationException(
+                code=UiPathErrorCode.EXECUTION_ERROR,
+                title="Guardrail filter action not supported",
+                detail=f"FilterAction is not supported for scope [{scope.name}] at this time.",
+                category=UiPathErrorCategory.USER,
+            )
+
+        return node_name, _node
+
+
+def _filter_tool_fields(
+    state: AgentGuardrailsGraphState,
+    fields_to_filter: list[FieldReference],
+    execution_stage: ExecutionStage,
+    tool_name: str,
+    guardrail_name: str,
+) -> dict[str, Any] | Command[Any]:
+    """Filter specified fields from tool call arguments or tool output.
+
+    The filter action filters fields based on the execution stage:
+    - PRE_EXECUTION: Only input fields are filtered
+    - POST_EXECUTION: Only output fields are filtered
+
+    Args:
+        state: The current agent graph state.
+        fields_to_filter: List of FieldReference objects specifying which fields to filter.
+        execution_stage: The execution stage (PRE_EXECUTION or POST_EXECUTION).
+        tool_name: Name of the tool to filter.
+        guardrail_name: Name of the guardrail for logging purposes.
+
+    Returns:
+        Command to update messages with filtered tool call args or output.
+
+    Raises:
+        AgentTerminationException: If filtering fails.
+    """
+    try:
+        if not fields_to_filter:
+            return {}
+
+        if execution_stage == ExecutionStage.PRE_EXECUTION:
+            return _filter_tool_input_fields(state, fields_to_filter, tool_name)
+        else:
+            return _filter_tool_output_fields(state, fields_to_filter)
+
+    except Exception as e:
+        raise AgentTerminationException(
+            code=UiPathErrorCode.EXECUTION_ERROR,
+            title="Filter action failed",
+            detail=f"Failed to filter tool fields: {str(e)}",
+            category=UiPathErrorCategory.USER,
+        ) from e
+
+
+def _filter_tool_input_fields(
+    state: AgentGuardrailsGraphState,
+    fields_to_filter: list[FieldReference],
+    tool_name: str,
+) -> dict[str, Any] | Command[Any]:
+    """Filter specified input fields from tool call arguments (PRE_EXECUTION only).
+
+    This function is called at PRE_EXECUTION to filter input fields from tool call arguments
+    before the tool is executed.
+
+    Args:
+        state: The current agent graph state.
+        fields_to_filter: List of FieldReference objects specifying which fields to filter.
+        tool_name: Name of the tool to filter.
+
+    Returns:
+        Command to update messages with filtered tool call args, or empty dict if no input fields to filter.
+    """
+    # Check if there are any input fields to filter
+    has_input_fields = any(
+        field_ref.source == FieldSource.INPUT for field_ref in fields_to_filter
+    )
+
+    if not has_input_fields:
+        return {}
+
+    msgs = state.messages.copy()
+    if not msgs:
+        return {}
+
+    # Find the AIMessage with tool calls
+    # At PRE_EXECUTION, this is always the last message
+    ai_message = None
+    for i in range(len(msgs) - 1, -1, -1):
+        msg = msgs[i]
+        if isinstance(msg, AIMessage) and msg.tool_calls:
+            ai_message = msg
+            break
+
+    if ai_message is None:
+        return {}
+
+    # Find and filter the tool call with matching name
+    # Type assertion: we know ai_message is AIMessage from the check above
+    assert isinstance(ai_message, AIMessage)
+    tool_calls = list(ai_message.tool_calls)
+    modified = False
+
+    for tool_call in tool_calls:
+        call_name = (
+            tool_call.get("name")
+            if isinstance(tool_call, dict)
+            else getattr(tool_call, "name", None)
+        )
+
+        if call_name == tool_name:
+            # Get the current args
+            args = (
+                tool_call.get("args")
+                if isinstance(tool_call, dict)
+                else getattr(tool_call, "args", None)
+            )
+
+            if args and isinstance(args, dict):
+                # Filter out the specified input fields
+                filtered_args = args.copy()
+                for field_ref in fields_to_filter:
+                    # Only filter input fields
+                    if (
+                        field_ref.source == FieldSource.INPUT
+                        and field_ref.path in filtered_args
+                    ):
+                        del filtered_args[field_ref.path]
+                        modified = True
+
+                # Update the tool call with filtered args
+                if isinstance(tool_call, dict):
+                    tool_call["args"] = filtered_args
+                else:
+                    tool_call.args = filtered_args
+
+            break
+
+    if modified:
+        ai_message.tool_calls = tool_calls
+        return Command(update={"messages": msgs})
+
+    return {}
+
+
+def _filter_tool_output_fields(
+    state: AgentGuardrailsGraphState,
+    fields_to_filter: list[FieldReference],
+) -> dict[str, Any] | Command[Any]:
+    """Filter specified output fields from tool output (POST_EXECUTION only).
+
+    This function is called at POST_EXECUTION to filter output fields from tool results
+    after the tool has been executed.
+
+    Args:
+        state: The current agent graph state.
+        fields_to_filter: List of FieldReference objects specifying which fields to filter.
+
+    Returns:
+        Command to update messages with filtered tool output, or empty dict if no output fields to filter.
+    """
+    # Check if there are any output fields to filter
+    has_output_fields = any(
+        field_ref.source == FieldSource.OUTPUT for field_ref in fields_to_filter
+    )
+
+    if not has_output_fields:
+        return {}
+
+    msgs = state.messages.copy()
+    if not msgs:
+        return {}
+
+    last_message = msgs[-1]
+    if not isinstance(last_message, ToolMessage):
+        return {}
+
+    # Parse the tool output content
+    import json
+
+    content = last_message.content
+    if not content:
+        return {}
+
+    # Try to parse the content as JSON or dict
+    try:
+        if isinstance(content, dict):
+            output_data = content
+        elif isinstance(content, str):
+            try:
+                output_data = json.loads(content)
+            except json.JSONDecodeError:
+                # Try to parse as Python literal (dict representation)
+                import ast
+
+                try:
+                    output_data = ast.literal_eval(content)
+                    if not isinstance(output_data, dict):
+                        return {}
+                except (ValueError, SyntaxError):
+                    return {}
+        else:
+            # Content is not JSON-parseable, can't filter specific fields
+            return {}
+    except Exception:
+        return {}
+
+    if not isinstance(output_data, dict):
+        return {}
+
+    # Filter out the specified fields
+    filtered_output = output_data.copy()
+    modified = False
+
+    for field_ref in fields_to_filter:
+        # Only filter output fields
+        if field_ref.source == FieldSource.OUTPUT and field_ref.path in filtered_output:
+            del filtered_output[field_ref.path]
+            modified = True
+
+    if modified:
+        # Update the tool message content with filtered output
+        last_message.content = json.dumps(filtered_output)
+        return Command(update={"messages": msgs})
+
+    return {}

uipath_langchain/agent/guardrails/actions/log_action.py

@@ -6,7 +6,7 @@ from uipath.platform.guardrails import BaseGuardrail, GuardrailScope
 
 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__)

uipath_langchain/agent/guardrails/guardrail_nodes.py

@@ -3,25 +3,108 @@ import logging
 import re
 from typing import Any, Callable
 
-from langchain_core.messages import AIMessage, AnyMessage, HumanMessage, SystemMessage
 from langgraph.types import Command
+from uipath.core.guardrails import (
+    DeterministicGuardrail,
+    DeterministicGuardrailsService,
+)
 from uipath.platform import UiPath
 from uipath.platform.guardrails import (
     BaseGuardrail,
+    BuiltInValidatorGuardrail,
     GuardrailScope,
 )
+from uipath.runtime.errors import UiPathErrorCode
 
 from uipath_langchain.agent.guardrails.types import ExecutionStage
+from uipath_langchain.agent.guardrails.utils import (
+    _extract_tool_args_from_message,
+    _extract_tool_output_data,
+    _extract_tools_args_from_message,
+    get_message_content,
+)
+from uipath_langchain.agent.react.types import AgentGuardrailsGraphState
 
-from .types import AgentGuardrailsGraphState
+from ..exceptions import AgentTerminationException
 
 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 _evaluate_deterministic_guardrail(
+    state: AgentGuardrailsGraphState,
+    guardrail: DeterministicGuardrail,
+    execution_stage: ExecutionStage,
+    input_data_extractor: Callable[[AgentGuardrailsGraphState], dict[str, Any]],
+    output_data_extractor: Callable[[AgentGuardrailsGraphState], dict[str, Any]] | None,
+):
+    """Evaluate deterministic guardrail.
+
+    Args:
+        state: The current agent graph state.
+        guardrail: The deterministic guardrail to evaluate.
+        execution_stage: The execution stage (PRE_EXECUTION or POST_EXECUTION).
+        input_data_extractor: Function to extract input data from state.
+        output_data_extractor: Function to extract output data from state (optional).
+
+    Returns:
+        The guardrail evaluation result.
+    """
+    service = DeterministicGuardrailsService()
+    input_data = input_data_extractor(state)
+
+    if execution_stage == ExecutionStage.PRE_EXECUTION:
+        return service.evaluate_pre_deterministic_guardrail(
+            input_data=input_data, guardrail=guardrail
+        )
+    else:  # POST_EXECUTION
+        output_data = output_data_extractor(state) if output_data_extractor else {}
+        return service.evaluate_post_deterministic_guardrail(
+            input_data=input_data,
+            output_data=output_data,
+            guardrail=guardrail,
+        )
+
+
+def _evaluate_builtin_guardrail(
+    state: AgentGuardrailsGraphState,
+    guardrail: BuiltInValidatorGuardrail,
+    payload_generator: Callable[[AgentGuardrailsGraphState], str],
+):
+    """Evaluate built-in validator guardrail.
+
+    Args:
+        state: The current agent graph state.
+        guardrail: The built-in validator guardrail to evaluate.
+        payload_generator: Function to generate payload text from state.
+
+    Returns:
+        The guardrail evaluation result.
+    """
+    text = payload_generator(state)
+    uipath = UiPath()
+    return uipath.guardrails.evaluate_guardrail(text, guardrail)
+
+
+def _create_validation_command(
+    result,
+    success_node: str,
+    failure_node: str,
+) -> Command[Any]:
+    """Create command based on validation result.
+
+    Args:
+        result: The guardrail evaluation result.
+        success_node: Node to route to on validation pass.
+        failure_node: Node to route to on validation fail.
+
+    Returns:
+        Command to update state and route to appropriate node.
+    """
+    if not result.validation_passed:
+        return Command(
+            goto=failure_node, update={"guardrail_validation_result": result.reason}
+        )
+    return Command(goto=success_node, update={"guardrail_validation_result": None})
 
 
 def _create_guardrail_node(
@@ -31,10 +114,15 @@ def _create_guardrail_node(
     payload_generator: Callable[[AgentGuardrailsGraphState], str],
     success_node: str,
     failure_node: str,
+    input_data_extractor: Callable[[AgentGuardrailsGraphState], dict[str, Any]]
+    | None = None,
+    output_data_extractor: Callable[[AgentGuardrailsGraphState], dict[str, Any]]
+    | None = None,
+    tool_name: str | None = None,
 ) -> tuple[str, Callable[[AgentGuardrailsGraphState], Any]]:
     """Private factory for guardrail evaluation nodes.
 
-    Returns a node that evaluates the guardrail and routes via Command:
+    Returns a node with observability metadata attached as __metadata__ attribute:
     - goto success_node on validation pass
     - goto failure_node on validation fail
     """
@@ -44,19 +132,60 @@ def _create_guardrail_node(
     async def node(
         state: AgentGuardrailsGraphState,
     ):
-        text = payload_generator(state)
         try:
-            uipath = UiPath()
-            result = uipath.guardrails.evaluate_guardrail(text, guardrail)
+            # Route to appropriate evaluation service based on guardrail type and scope
+            if (
+                isinstance(guardrail, DeterministicGuardrail)
+                and scope == GuardrailScope.TOOL
+                and input_data_extractor is not None
+            ):
+                result = _evaluate_deterministic_guardrail(
+                    state,
+                    guardrail,
+                    execution_stage,
+                    input_data_extractor,
+                    output_data_extractor,
+                )
+            elif isinstance(guardrail, BuiltInValidatorGuardrail):
+                result = _evaluate_builtin_guardrail(
+                    state, guardrail, payload_generator
+                )
+            else:
+                # Provide specific error message for DeterministicGuardrails with wrong scope
+                if isinstance(guardrail, DeterministicGuardrail):
+                    raise AgentTerminationException(
+                        code=UiPathErrorCode.EXECUTION_ERROR,
+                        title="Invalid guardrail scope",
+                        detail=f"DeterministicGuardrail '{guardrail.name}' can only be used with TOOL scope. "
+                        f"Current scope: {scope.name}. "
+                        f"Please configure this guardrail to use only TOOL scope.",
+                    )
+                else:
+                    raise AgentTerminationException(
+                        code=UiPathErrorCode.EXECUTION_ERROR,
+                        title="Unsupported guardrail type",
+                        detail=f"Guardrail type '{type(guardrail).__name__}' is not supported. "
+                        f"Expected DeterministicGuardrail (TOOL scope only) or BuiltInValidatorGuardrail.",
+                    )
+
+            return _create_validation_command(result, success_node, failure_node)
+
         except Exception as exc:
-            logger.error("Failed to evaluate guardrail: %s", exc)
+            logger.error(
+                "Failed to evaluate guardrail '%s': %s",
+                guardrail.name,
+                exc,
+            )
             raise
 
-        if not result.validation_passed:
-            return Command(
-                goto=failure_node, update={"guardrail_validation_result": result.reason}
-            )
-        return Command(goto=success_node, update={"guardrail_validation_result": None})
+    # Attach observability metadata as function attribute
+    node.__metadata__ = {  # type: ignore[attr-defined]
+        "guardrail_name": guardrail.name,
+        "guardrail_description": getattr(guardrail, "description", None),
+        "guardrail_scope": scope.value,
+        "guardrail_stage": execution_stage.value,
+        "tool_name": tool_name,
+    }
 
     return node_name, node
 
@@ -70,7 +199,11 @@ def create_llm_guardrail_node(
     def _payload_generator(state: AgentGuardrailsGraphState) -> str:
         if not state.messages:
             return ""
-        return _message_text(state.messages[-1])
+        match execution_stage:
+            case ExecutionStage.PRE_EXECUTION:
+                return get_message_content(state.messages[-1])
+            case ExecutionStage.POST_EXECUTION:
+                return json.dumps(_extract_tools_args_from_message(state.messages[-1]))
 
     return _create_guardrail_node(
         guardrail,
@@ -82,17 +215,35 @@ def create_llm_guardrail_node(
     )
 
 
-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,
@@ -137,31 +288,28 @@ def create_tool_guardrail_node(
             return ""
 
         if execution_stage == ExecutionStage.PRE_EXECUTION:
-            if not isinstance(state.messages[-1], AIMessage):
-                return ""
-            message = state.messages[-1]
+            last_message = state.messages[-1]
+            args_dict = _extract_tool_args_from_message(last_message, tool_name)
+            if args_dict:
+                return json.dumps(args_dict)
 
-            if not message.tool_calls:
-                return ""
+        return get_message_content(state.messages[-1])
 
-            # 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:
-                        return json.dumps(args)
+    # Create closures for input/output data extraction (for deterministic guardrails)
+    def _input_data_extractor(state: AgentGuardrailsGraphState) -> dict[str, Any]:
+        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)
 
-        return _message_text(state.messages[-1])
+    def _output_data_extractor(state: AgentGuardrailsGraphState) -> dict[str, Any]:
+        return _extract_tool_output_data(state)
 
     return _create_guardrail_node(
         guardrail,
@@ -170,4 +318,7 @@ def create_tool_guardrail_node(
         _payload_generator,
         success_node,
         failure_node,
+        _input_data_extractor,
+        _output_data_extractor,
+        tool_name,
     )