Merge pull request #351 from UiPath/filter-agent-llm

feat: don't allow filter action for agent and llm scopes [AL-249]

Author: valentinabojan

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

@@ -1,11 +1,13 @@
 from .base_action import GuardrailAction
 from .block_action import BlockAction
 from .escalate_action import EscalateAction
+from .filter_action import FilterAction
 from .log_action import LogAction
 
 __all__ = [
     "GuardrailAction",
     "BlockAction",
     "LogAction",
     "EscalateAction",
+    "FilterAction",
 ]

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

@@ -0,0 +1,55 @@
+import re
+from typing import Any
+
+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 now, filtering is only supported for non-AGENT and non-LLM scopes.
+    If invoked for ``GuardrailScope.AGENT`` or ``GuardrailScope.LLM``, this action
+    raises an exception to indicate the operation is not supported yet.
+    """
+
+    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]:
+            if scope in (GuardrailScope.AGENT, GuardrailScope.LLM):
+                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,
+                )
+            # No-op for other scopes for now.
+            return {}
+
+        return node_name, _node

tests/agent/guardrails/actions/test_filter_action.py

@@ -0,0 +1,102 @@
+"""Tests for FilterAction guardrail failure behavior."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+from unittest.mock import MagicMock
+
+import pytest
+from uipath.platform.guardrails import GuardrailScope
+from uipath.runtime.errors import UiPathErrorCode
+
+from uipath_langchain.agent.exceptions import AgentTerminationException
+from uipath_langchain.agent.guardrails.actions.filter_action import FilterAction
+from uipath_langchain.agent.guardrails.types import ExecutionStage
+from uipath_langchain.agent.react.types import AgentGuardrailsGraphState
+
+if TYPE_CHECKING:
+    from _pytest.capture import CaptureFixture  # noqa: F401
+    from _pytest.fixtures import FixtureRequest  # noqa: F401
+    from _pytest.logging import LogCaptureFixture  # noqa: F401
+    from _pytest.monkeypatch import MonkeyPatch  # noqa: F401
+    from pytest_mock.plugin import MockerFixture  # noqa: F401
+
+
+class TestFilterAction:
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        ("scope", "stage", "expected_node_name"),
+        [
+            (
+                GuardrailScope.LLM,
+                ExecutionStage.PRE_EXECUTION,
+                "llm_pre_execution_my_guardrail_v1_filter",
+            ),
+            (
+                GuardrailScope.LLM,
+                ExecutionStage.POST_EXECUTION,
+                "llm_post_execution_my_guardrail_v1_filter",
+            ),
+            (
+                GuardrailScope.AGENT,
+                ExecutionStage.PRE_EXECUTION,
+                "agent_pre_execution_my_guardrail_v1_filter",
+            ),
+            (
+                GuardrailScope.AGENT,
+                ExecutionStage.POST_EXECUTION,
+                "agent_post_execution_my_guardrail_v1_filter",
+            ),
+        ],
+    )
+    async def test_node_name_and_exception_for_unsupported_scopes(
+        self, scope: GuardrailScope, stage: ExecutionStage, expected_node_name: str
+    ) -> None:
+        """AGENT/LLM scopes raise AgentTerminationException and node name is sanitized."""
+        action = FilterAction()
+        guardrail = MagicMock()
+        guardrail.name = "My Guardrail v1"
+
+        node_name, node = action.action_node(
+            guardrail=guardrail,
+            scope=scope,
+            execution_stage=stage,
+            guarded_component_name="guarded_node_name",
+        )
+
+        assert node_name == expected_node_name
+
+        with pytest.raises(AgentTerminationException) as excinfo:
+            await node(AgentGuardrailsGraphState(messages=[]))
+
+        # Validate rich error info
+        assert (
+            excinfo.value.error_info.code
+            == f"Python.{UiPathErrorCode.EXECUTION_ERROR.value}"
+        )
+        assert excinfo.value.error_info.title == "Guardrail filter action not supported"
+        assert (
+            excinfo.value.error_info.detail
+            == f"FilterAction is not supported for scope [{scope.name}] at this time."
+        )
+
+    @pytest.mark.asyncio
+    @pytest.mark.parametrize(
+        "stage",
+        [ExecutionStage.PRE_EXECUTION, ExecutionStage.POST_EXECUTION],
+    )
+    async def test_tool_scope_returns_empty_dict(self, stage: ExecutionStage) -> None:
+        """TOOL scope currently performs no-op and returns empty dict."""
+        action = FilterAction()
+        guardrail = MagicMock()
+        guardrail.name = "My Guardrail v1"
+
+        _, node = action.action_node(
+            guardrail=guardrail,
+            scope=GuardrailScope.TOOL,
+            execution_stage=stage,
+            guarded_component_name="test_tool",
+        )
+
+        result = await node(AgentGuardrailsGraphState(messages=[]))
+        assert result == {}