diff --git a/pyproject.toml b/pyproject.toml index 569f8541..ba4a05d5 100644 --- a/pyproject.toml +++ b/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" diff --git a/src/uipath_langchain/agent/guardrails/__init__.py b/src/uipath_langchain/agent/guardrails/__init__.py index d5091acb..f5a220af 100644 --- a/src/uipath_langchain/agent/guardrails/__init__.py +++ b/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", ] diff --git a/src/uipath_langchain/agent/guardrails/actions/block_action.py b/src/uipath_langchain/agent/guardrails/actions/block_action.py index 70e4ff1e..70681d54 100644 --- a/src/uipath_langchain/agent/guardrails/actions/block_action.py +++ b/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 diff --git a/src/uipath_langchain/agent/guardrails/actions/escalate_action.py b/src/uipath_langchain/agent/guardrails/actions/escalate_action.py index a2a316e6..6d9dd99a 100644 --- a/src/uipath_langchain/agent/guardrails/actions/escalate_action.py +++ b/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( diff --git a/src/uipath_langchain/agent/guardrails/actions/log_action.py b/src/uipath_langchain/agent/guardrails/actions/log_action.py index 2d5b52ff..c6ef93eb 100644 --- a/src/uipath_langchain/agent/guardrails/actions/log_action.py +++ b/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__) diff --git a/src/uipath_langchain/agent/guardrails/guardrail_nodes.py b/src/uipath_langchain/agent/guardrails/guardrail_nodes.py index 522f591d..66288fa4 100644 --- a/src/uipath_langchain/agent/guardrails/guardrail_nodes.py +++ b/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, diff --git a/src/uipath_langchain/agent/guardrails/types.py b/src/uipath_langchain/agent/guardrails/types.py index 2f13f795..0e4b6bab 100644 --- a/src/uipath_langchain/agent/guardrails/types.py +++ b/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): diff --git a/src/uipath_langchain/agent/guardrails/utils.py b/src/uipath_langchain/agent/guardrails/utils.py new file mode 100644 index 00000000..bd02ec44 --- /dev/null +++ b/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 "" diff --git a/src/uipath_langchain/agent/react/agent.py b/src/uipath_langchain/agent/react/agent.py index cae80fee..af7de6be 100644 --- a/src/uipath_langchain/agent/react/agent.py +++ b/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) diff --git a/src/uipath_langchain/agent/guardrails/guardrails_subgraph.py b/src/uipath_langchain/agent/react/guardrails/guardrails_subgraph.py similarity index 81% rename from src/uipath_langchain/agent/guardrails/guardrails_subgraph.py rename to src/uipath_langchain/agent/react/guardrails/guardrails_subgraph.py index ed6a97ad..51221b9a 100644 --- a/src/uipath_langchain/agent/guardrails/guardrails_subgraph.py +++ b/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], diff --git a/src/uipath_langchain/agent/react/types.py b/src/uipath_langchain/agent/react/types.py index 3b85e10f..b7d1da90 100644 --- a/src/uipath_langchain/agent/react/types.py +++ b/src/uipath_langchain/agent/react/types.py @@ -1,5 +1,5 @@ from enum import StrEnum -from typing import Annotated +from typing import Annotated, Any, Optional from langchain_core.messages import AnyMessage from langgraph.graph.message import add_messages @@ -25,12 +25,21 @@ class AgentGraphState(BaseModel): termination: AgentTermination | None = None +class AgentGuardrailsGraphState(AgentGraphState): + """Agent Guardrails Graph state for guardrail subgraph.""" + + guardrail_validation_result: Optional[str] = None + agent_result: Optional[dict[str, Any]] = None + + class AgentGraphNode(StrEnum): INIT = "init" + GUARDED_INIT = "guarded-init" AGENT = "agent" LLM = "llm" TOOLS = "tools" TERMINATE = "terminate" + GUARDED_TERMINATE = "guarded-terminate" class AgentGraphConfig(BaseModel): diff --git a/tests/agent/guardrails/actions/test_block_action.py b/tests/agent/guardrails/actions/test_block_action.py index d6b3f1e2..a5b69828 100644 --- a/tests/agent/guardrails/actions/test_block_action.py +++ b/tests/agent/guardrails/actions/test_block_action.py @@ -1,5 +1,7 @@ """Tests for BlockAction guardrail failure behavior.""" +from __future__ import annotations + from unittest.mock import MagicMock import pytest @@ -8,27 +10,64 @@ from uipath_langchain.agent.exceptions import AgentTerminationException from uipath_langchain.agent.guardrails.actions.block_action import BlockAction from uipath_langchain.agent.guardrails.types import ( - AgentGuardrailsGraphState, ExecutionStage, ) +from uipath_langchain.agent.react.types import AgentGuardrailsGraphState class TestBlockAction: @pytest.mark.asyncio - async def test_node_name_and_exception_pre_llm(self): - """PreExecution + LLM: name is sanitized and node raises correct exception.""" + @pytest.mark.parametrize( + ("scope", "stage", "expected_node_name"), + [ + ( + GuardrailScope.LLM, + ExecutionStage.PRE_EXECUTION, + "llm_pre_execution_my_guardrail_v1_block", + ), + ( + GuardrailScope.AGENT, + ExecutionStage.PRE_EXECUTION, + "agent_pre_execution_my_guardrail_v1_block", + ), + ( + GuardrailScope.TOOL, + ExecutionStage.PRE_EXECUTION, + "tool_pre_execution_my_guardrail_v1_block", + ), + ( + GuardrailScope.LLM, + ExecutionStage.POST_EXECUTION, + "llm_post_execution_my_guardrail_v1_block", + ), + ( + GuardrailScope.AGENT, + ExecutionStage.POST_EXECUTION, + "agent_post_execution_my_guardrail_v1_block", + ), + ( + GuardrailScope.TOOL, + ExecutionStage.POST_EXECUTION, + "tool_post_execution_my_guardrail_v1_block", + ), + ], + ) + async def test_node_name_and_exception( + self, scope: GuardrailScope, stage: ExecutionStage, expected_node_name: str + ) -> None: + """Name is sanitized and node raises correct exception for each scope/stage.""" action = BlockAction(reason="Sensitive data detected") guardrail = MagicMock() guardrail.name = "My Guardrail v1" node_name, node = action.action_node( guardrail=guardrail, - scope=GuardrailScope.LLM, - execution_stage=ExecutionStage.PRE_EXECUTION, + scope=scope, + execution_stage=stage, guarded_component_name="guarded_node_name", ) - assert node_name == "llm_pre_execution_my_guardrail_v1_block" + assert node_name == expected_node_name with pytest.raises(AgentTerminationException) as excinfo: await node(AgentGuardrailsGraphState(messages=[])) diff --git a/tests/agent/guardrails/actions/test_escalate_action.py b/tests/agent/guardrails/actions/test_escalate_action.py index db5dcc96..313f39ba 100644 --- a/tests/agent/guardrails/actions/test_escalate_action.py +++ b/tests/agent/guardrails/actions/test_escalate_action.py @@ -1,5 +1,7 @@ """Tests for EscalateAction guardrail failure behavior.""" +from __future__ import annotations + import json from unittest.mock import MagicMock, patch @@ -12,15 +14,52 @@ from uipath_langchain.agent.exceptions import AgentTerminationException from uipath_langchain.agent.guardrails.actions.escalate_action import EscalateAction from uipath_langchain.agent.guardrails.types import ( - AgentGuardrailsGraphState, ExecutionStage, ) +from uipath_langchain.agent.react.types import AgentGuardrailsGraphState class TestEscalateAction: @pytest.mark.asyncio - async def test_node_name_pre_llm(self): - """PreExecution + LLM: name is sanitized correctly.""" + @pytest.mark.parametrize( + ("scope", "stage", "expected_node_name"), + [ + ( + GuardrailScope.LLM, + ExecutionStage.PRE_EXECUTION, + "my_guardrail_1_hitl_pre_execution_llm", + ), + ( + GuardrailScope.LLM, + ExecutionStage.POST_EXECUTION, + "my_guardrail_1_hitl_post_execution_llm", + ), + ( + GuardrailScope.AGENT, + ExecutionStage.PRE_EXECUTION, + "my_guardrail_1_hitl_pre_execution_agent", + ), + ( + GuardrailScope.AGENT, + ExecutionStage.POST_EXECUTION, + "my_guardrail_1_hitl_post_execution_agent", + ), + ( + GuardrailScope.TOOL, + ExecutionStage.PRE_EXECUTION, + "my_guardrail_1_hitl_pre_execution_tool", + ), + ( + GuardrailScope.TOOL, + ExecutionStage.POST_EXECUTION, + "my_guardrail_1_hitl_post_execution_tool", + ), + ], + ) + async def test_node_name( + self, scope: GuardrailScope, stage: ExecutionStage, expected_node_name: str + ) -> None: + """Node name is sanitized correctly for each scope/stage.""" action = EscalateAction( app_name="TestApp", app_folder_path="TestFolder", @@ -28,39 +67,17 @@ async def test_node_name_pre_llm(self): assignee="test@example.com", ) guardrail = MagicMock() - guardrail.name = "My Guardrail v1" + guardrail.name = "My Guardrail 1" guardrail.description = "Test description" node_name, _ = action.action_node( guardrail=guardrail, - scope=GuardrailScope.LLM, - execution_stage=ExecutionStage.PRE_EXECUTION, - guarded_component_name="test_node", - ) - - assert node_name == "my_guardrail_v1_hitl_pre_execution_llm" - - @pytest.mark.asyncio - async def test_node_name_post_agent(self): - """PostExecution + AGENT: name is sanitized correctly.""" - action = EscalateAction( - app_name="TestApp", - app_folder_path="TestFolder", - version=1, - assignee="test@example.com", - ) - guardrail = MagicMock() - guardrail.name = "Special-Guardrail@2024" - guardrail.description = "Test description" - - node_name, _ = action.action_node( - guardrail=guardrail, - scope=GuardrailScope.AGENT, - execution_stage=ExecutionStage.POST_EXECUTION, + scope=scope, + execution_stage=stage, guarded_component_name="test_node", ) - assert node_name == "special_guardrail_2024_hitl_post_execution_agent" + assert node_name == expected_node_name @pytest.mark.asyncio @patch("uipath_langchain.agent.guardrails.actions.escalate_action.interrupt") diff --git a/tests/agent/guardrails/actions/test_log_action.py b/tests/agent/guardrails/actions/test_log_action.py index cf172ff3..bc6c70df 100644 --- a/tests/agent/guardrails/actions/test_log_action.py +++ b/tests/agent/guardrails/actions/test_log_action.py @@ -1,6 +1,7 @@ """Tests for LogAction guardrail failure behavior.""" import logging +from itertools import product from unittest.mock import MagicMock import pytest @@ -8,9 +9,9 @@ from uipath_langchain.agent.guardrails.actions.log_action import LogAction from uipath_langchain.agent.guardrails.types import ( - AgentGuardrailsGraphState, ExecutionStage, ) +from uipath_langchain.agent.react.types import AgentGuardrailsGraphState class TestLogAction: @@ -47,23 +48,39 @@ async def test_node_name_and_logs_custom_message( ) @pytest.mark.asyncio + @pytest.mark.parametrize( + ("scope", "stage", "level"), + list( + product( + (GuardrailScope.LLM, GuardrailScope.AGENT, GuardrailScope.TOOL), + (ExecutionStage.PRE_EXECUTION, ExecutionStage.POST_EXECUTION), + (logging.INFO, logging.WARNING, logging.ERROR), + ) + ), + ) async def test_default_message_includes_context( - self, caplog: pytest.LogCaptureFixture + self, + caplog: pytest.LogCaptureFixture, + scope: GuardrailScope, + stage: ExecutionStage, + level: int, ) -> None: - """PostExecution + TOOL: default message includes guardrail name, scope, stage, and reason.""" - action = LogAction(message=None, level=logging.WARNING) + """Default message includes guardrail name, scope, stage, and reason.""" + action = LogAction(message=None, level=level) guardrail = MagicMock() guardrail.name = "My Guardrail" node_name, node = action.action_node( guardrail=guardrail, - scope=GuardrailScope.TOOL, - execution_stage=ExecutionStage.POST_EXECUTION, + scope=scope, + execution_stage=stage, guarded_component_name="guarded_node_name", ) - assert node_name == "tool_post_execution_my_guardrail_log" + assert ( + node_name == f"{scope.name.lower()}_{stage.name.lower()}_my_guardrail_log" + ) - with caplog.at_level(logging.WARNING): + with caplog.at_level(level): result = await node( AgentGuardrailsGraphState( messages=[], guardrail_validation_result="bad input" @@ -72,9 +89,10 @@ async def test_default_message_includes_context( assert result == {} # Confirm default formatted message content + expected = ( + "Guardrail [My Guardrail] validation failed for " + f"[{scope.name}] [{stage.name}] with the following reason: bad input" + ) assert any( - rec.levelno == logging.WARNING - and rec.message - == "Guardrail [My Guardrail] validation failed for [TOOL] [POST_EXECUTION] with the following reason: bad input" - for rec in caplog.records + rec.levelno == level and rec.message == expected for rec in caplog.records ) diff --git a/tests/agent/guardrails/test_agent_init_guardrails_subgraph.py b/tests/agent/guardrails/test_agent_init_guardrails_subgraph.py new file mode 100644 index 00000000..5ff749f1 --- /dev/null +++ b/tests/agent/guardrails/test_agent_init_guardrails_subgraph.py @@ -0,0 +1,105 @@ +"""Tests for init-node (agent-scope) guardrails subgraph construction.""" + +from __future__ import annotations + +from typing import Sequence +from unittest.mock import MagicMock + +from _pytest.monkeypatch import MonkeyPatch +from uipath.core.guardrails import BaseGuardrail, GuardrailSelector +from uipath.platform.guardrails import GuardrailScope + +import uipath_langchain.agent.react.guardrails.guardrails_subgraph as mod +from tests.agent.guardrails.test_guardrail_utils import ( + FakeStateGraph, + fake_action, + fake_factory, +) +from uipath_langchain.agent.guardrails.actions import GuardrailAction + + +class TestAgentInitGuardrailsSubgraph: + def test_no_applicable_guardrails_returns_original_node(self): + """If no guardrails match the AGENT scope, the original node should be returned.""" + inner = ("inner", lambda s: s) + guardrails: Sequence[tuple[BaseGuardrail, GuardrailAction]] = [] + + # Case with empty guardrails + result = mod.create_agent_init_guardrails_subgraph( + init_node=inner, guardrails=guardrails + ) + assert result == inner[1] + + # Case with None guardrails + result_none = mod.create_agent_init_guardrails_subgraph( + init_node=inner, guardrails=None + ) + assert result_none == inner[1] + + # Case with guardrails but none matching AGENT scope + non_matching_guardrail = MagicMock() + non_matching_guardrail.selector = GuardrailSelector(scopes=[GuardrailScope.LLM]) + guardrails_non_match = [(non_matching_guardrail, MagicMock())] + + result_non_match = mod.create_agent_init_guardrails_subgraph( + init_node=inner, guardrails=guardrails_non_match + ) + assert result_non_match == inner[1] + + def test_two_guardrails_build_post_chain(self, monkeypatch: MonkeyPatch) -> None: + """Two AGENT guardrails should create a POST_EXECUTION chain with failure edges.""" + monkeypatch.setattr(mod, "StateGraph", FakeStateGraph) + monkeypatch.setattr(mod, "START", "START") + monkeypatch.setattr(mod, "END", "END") + monkeypatch.setattr( + mod, "create_agent_init_guardrail_node", fake_factory("eval") + ) + + guardrail1 = MagicMock() + guardrail1.name = "guardrail1" + guardrail1.selector = GuardrailSelector(scopes=[GuardrailScope.AGENT]) + + guardrail2 = MagicMock() + guardrail2.name = "guardrail2" + guardrail2.selector = GuardrailSelector(scopes=[GuardrailScope.AGENT]) + + non_matching = MagicMock() + non_matching.name = "llm_guardrail" + non_matching.selector = GuardrailSelector(scopes=[GuardrailScope.LLM]) + + guardrails: Sequence[tuple[BaseGuardrail, GuardrailAction]] = [ + (guardrail1, fake_action("log")), + (guardrail2, fake_action("block")), + (non_matching, fake_action("noop")), + ] + + inner = ("inner", lambda s: s) + result_graph = mod.create_agent_init_guardrails_subgraph( + init_node=inner, + guardrails=guardrails, + ) + + post_g1 = "eval_post_execution_guardrail1" + log_post_g1 = "log_post_execution_guardrail1" + post_g2 = "eval_post_execution_guardrail2" + block_post_g2 = "block_post_execution_guardrail2" + + expected_edges = { + ("START", "inner"), + ("inner", post_g1), + (log_post_g1, post_g2), + (block_post_g2, "END"), + } + assert expected_edges.issubset(set(result_graph.edges)) + + node_names = {name for name, _ in result_graph.nodes} + for name in [ + "inner", + post_g1, + post_g2, + log_post_g1, + block_post_g2, + ]: + assert name in node_names + assert "eval_post_execution_llm_guardrail" not in node_names + assert "noop_post_execution_llm_guardrail" not in node_names diff --git a/tests/agent/guardrails/test_agent_terminate_guardrails_subgraph.py b/tests/agent/guardrails/test_agent_terminate_guardrails_subgraph.py new file mode 100644 index 00000000..7d0689b1 --- /dev/null +++ b/tests/agent/guardrails/test_agent_terminate_guardrails_subgraph.py @@ -0,0 +1,128 @@ +"""Tests for terminate-node (agent-scope) guardrails subgraph construction.""" + +from __future__ import annotations + +import types +from typing import Any, Sequence +from unittest.mock import MagicMock + +from _pytest.monkeypatch import MonkeyPatch +from uipath.core.guardrails import BaseGuardrail, GuardrailSelector +from uipath.platform.guardrails import GuardrailScope + +import uipath_langchain.agent.react.guardrails.guardrails_subgraph as mod +from tests.agent.guardrails.test_guardrail_utils import ( + FakeStateGraphWithAinvoke, + fake_action, + fake_factory, +) +from uipath_langchain.agent.guardrails.actions import GuardrailAction + + +class TestAgentTerminateGuardrailsSubgraph: + def test_no_applicable_guardrails_returns_original_node(self): + """If no guardrails match the AGENT scope, the original node should be returned.""" + inner = ("inner", lambda s: s) + guardrails: Sequence[tuple[BaseGuardrail, GuardrailAction]] = [] + + # Case with empty guardrails + result = mod.create_agent_terminate_guardrails_subgraph( + terminate_node=inner, guardrails=guardrails + ) + assert result == inner[1] + + # Case with None guardrails + result_none = mod.create_agent_terminate_guardrails_subgraph( + terminate_node=inner, guardrails=None + ) + assert result_none == inner[1] + + # Case with guardrails but none matching AGENT scope + non_matching_guardrail = MagicMock() + non_matching_guardrail.selector = GuardrailSelector(scopes=[GuardrailScope.LLM]) + guardrails_non_match = [(non_matching_guardrail, MagicMock())] + + result_non_match = mod.create_agent_terminate_guardrails_subgraph( + terminate_node=inner, guardrails=guardrails_non_match + ) + assert result_non_match == inner[1] + + async def test_two_guardrails_build_post_chain_and_return_result( + self, monkeypatch: MonkeyPatch + ): + """Two AGENT guardrails should create a POST_EXECUTION chain and returns terminate result.""" + + monkeypatch.setattr(mod, "StateGraph", FakeStateGraphWithAinvoke) + monkeypatch.setattr(mod, "START", "START") + monkeypatch.setattr(mod, "END", "END") + monkeypatch.setattr( + mod, "create_agent_terminate_guardrail_node", fake_factory("eval") + ) + + captured: dict[str, Any] = {} + original_create = mod._create_guardrails_subgraph + + def _capture_create(*args: Any, **kwargs: Any) -> Any: + compiled = original_create(*args, **kwargs) + captured["compiled"] = compiled + return compiled + + monkeypatch.setattr(mod, "_create_guardrails_subgraph", _capture_create) + + guardrail1 = MagicMock() + guardrail1.name = "guardrail1" + guardrail1.selector = GuardrailSelector(scopes=[GuardrailScope.AGENT]) + + guardrail2 = MagicMock() + guardrail2.name = "guardrail2" + guardrail2.selector = GuardrailSelector(scopes=[GuardrailScope.AGENT]) + + non_matching = MagicMock() + non_matching.name = "llm_guardrail" + non_matching.selector = GuardrailSelector(scopes=[GuardrailScope.LLM]) + + guardrails: Sequence[tuple[BaseGuardrail, GuardrailAction]] = [ + (guardrail1, fake_action("log")), + (guardrail2, fake_action("block")), + (non_matching, fake_action("noop")), + ] + + expected_result: dict[str, Any] = {"done": 1} + + def terminate_fn(_state: Any) -> dict[str, Any]: + return expected_result + + run_terminate = mod.create_agent_terminate_guardrails_subgraph( + terminate_node=("terminate", terminate_fn), + guardrails=guardrails, + ) + + result = await run_terminate(types.SimpleNamespace(messages=["m"])) + assert result == expected_result + + result_graph = captured["compiled"] + + post_g1 = "eval_post_execution_guardrail1" + log_post_g1 = "log_post_execution_guardrail1" + post_g2 = "eval_post_execution_guardrail2" + block_post_g2 = "block_post_execution_guardrail2" + + expected_edges = { + ("START", "terminate"), + ("terminate", post_g1), + (log_post_g1, post_g2), + (block_post_g2, "END"), + } + assert expected_edges.issubset(set(result_graph.edges)) + + node_names = {name for name, _ in result_graph.nodes} + for name in [ + "terminate", + post_g1, + post_g2, + log_post_g1, + block_post_g2, + ]: + assert name in node_names + assert "eval_post_execution_llm_guardrail" not in node_names + assert "noop_post_execution_llm_guardrail" not in node_names diff --git a/tests/agent/guardrails/test_guardrail_nodes.py b/tests/agent/guardrails/test_guardrail_nodes.py index 3b80965a..9a9144c6 100644 --- a/tests/agent/guardrails/test_guardrail_nodes.py +++ b/tests/agent/guardrails/test_guardrail_nodes.py @@ -1,18 +1,23 @@ """Tests for guardrail node creation and routing.""" +import json import types from unittest.mock import MagicMock import pytest -from langchain_core.messages import HumanMessage, SystemMessage +from _pytest.monkeypatch import MonkeyPatch +from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage 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 uipath_langchain.agent.guardrails.types import ( - AgentGuardrailsGraphState, ExecutionStage, ) +from uipath_langchain.agent.react.types import AgentGuardrailsGraphState class FakeGuardrails: @@ -32,7 +37,7 @@ def __init__(self, result): self.guardrails = FakeGuardrails(result) -def _patch_uipath(monkeypatch, validation_passed=True, reason=None): +def _patch_uipath(monkeypatch, *, validation_passed=True, reason=None): result = types.SimpleNamespace(validation_passed=validation_passed, reason=reason) fake = FakeUiPath(result) monkeypatch.setattr( @@ -102,3 +107,246 @@ async def test_llm_failure_pre_and_post( cmd = await node(state) assert cmd.goto == "nope" assert cmd.update == {"guardrail_validation_result": "policy_violation"} + + +class TestAgentInitGuardrailNodes: + @pytest.mark.asyncio + @pytest.mark.parametrize( + "execution_stage,expected_name", + [ + (ExecutionStage.PRE_EXECUTION, "agent_pre_execution_example"), + (ExecutionStage.POST_EXECUTION, "agent_post_execution_example"), + ], + ids=["pre-success", "post-success"], + ) + async def test_agent_init_success_pre_and_post( + self, + monkeypatch: MonkeyPatch, + execution_stage: ExecutionStage, + expected_name: str, + ) -> None: + """Agent init node: routes to success and passes message payload to evaluator.""" + guardrail = MagicMock() + guardrail.name = "Example" + fake = _patch_uipath(monkeypatch, validation_passed=True, reason=None) + + node_name, node = create_agent_init_guardrail_node( + guardrail=guardrail, + execution_stage=execution_stage, + success_node="ok", + failure_node="nope", + ) + assert node_name == expected_name + + state = AgentGuardrailsGraphState(messages=[HumanMessage("payload")]) + cmd = await node(state) + assert cmd.goto == "ok" + assert cmd.update == {"guardrail_validation_result": None} + assert fake.guardrails.last_text == "payload" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "execution_stage,expected_name", + [ + (ExecutionStage.PRE_EXECUTION, "agent_pre_execution_example"), + (ExecutionStage.POST_EXECUTION, "agent_post_execution_example"), + ], + ids=["pre-fail", "post-fail"], + ) + async def test_agent_init_failure_pre_and_post( + self, + monkeypatch: MonkeyPatch, + execution_stage: ExecutionStage, + expected_name: str, + ) -> None: + """Agent init node: routes to failure and sets guardrail_validation_result.""" + guardrail = MagicMock() + guardrail.name = "Example" + _patch_uipath(monkeypatch, validation_passed=False, reason="policy_violation") + + node_name, node = create_agent_init_guardrail_node( + guardrail=guardrail, + execution_stage=execution_stage, + success_node="ok", + failure_node="nope", + ) + assert node_name == expected_name + + state = AgentGuardrailsGraphState(messages=[SystemMessage("payload")]) + cmd = await node(state) + assert cmd.goto == "nope" + assert cmd.update == {"guardrail_validation_result": "policy_violation"} + + +class TestAgentTerminateGuardrailNodes: + @pytest.mark.asyncio + @pytest.mark.parametrize( + "execution_stage,expected_name", + [ + (ExecutionStage.PRE_EXECUTION, "agent_pre_execution_example"), + (ExecutionStage.POST_EXECUTION, "agent_post_execution_example"), + ], + ids=["pre-success", "post-success"], + ) + async def test_agent_terminate_success_pre_and_post( + self, + monkeypatch: MonkeyPatch, + execution_stage: ExecutionStage, + expected_name: str, + ) -> None: + """Agent terminate node: routes to success and passes agent_result payload to evaluator.""" + guardrail = MagicMock() + guardrail.name = "Example" + fake = _patch_uipath(monkeypatch, validation_passed=True, reason=None) + + node_name, node = create_agent_terminate_guardrail_node( + guardrail=guardrail, + execution_stage=execution_stage, + success_node="ok", + failure_node="nope", + ) + assert node_name == expected_name + + agent_result = {"ok": True} + state = AgentGuardrailsGraphState(messages=[], agent_result=agent_result) + cmd = await node(state) + assert cmd.goto == "ok" + assert cmd.update == {"guardrail_validation_result": None} + assert fake.guardrails.last_text == str(agent_result) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "execution_stage,expected_name", + [ + (ExecutionStage.PRE_EXECUTION, "agent_pre_execution_example"), + (ExecutionStage.POST_EXECUTION, "agent_post_execution_example"), + ], + ids=["pre-fail", "post-fail"], + ) + async def test_agent_terminate_failure_pre_and_post( + self, + monkeypatch: MonkeyPatch, + execution_stage: ExecutionStage, + expected_name: str, + ) -> None: + """Agent terminate node: routes to failure and sets guardrail_validation_result.""" + guardrail = MagicMock() + guardrail.name = "Example" + _patch_uipath(monkeypatch, validation_passed=False, reason="policy_violation") + + node_name, node = create_agent_terminate_guardrail_node( + guardrail=guardrail, + execution_stage=execution_stage, + success_node="ok", + failure_node="nope", + ) + assert node_name == expected_name + + state = AgentGuardrailsGraphState(messages=[], agent_result={"ok": False}) + cmd = await node(state) + assert cmd.goto == "nope" + assert cmd.update == {"guardrail_validation_result": "policy_violation"} + + +class TestToolGuardrailNodes: + @pytest.mark.asyncio + @pytest.mark.parametrize( + "execution_stage,expected_name", + [ + (ExecutionStage.PRE_EXECUTION, "tool_pre_execution_example"), + (ExecutionStage.POST_EXECUTION, "tool_post_execution_example"), + ], + ids=["pre-success", "post-success"], + ) + async def test_tool_success_pre_and_post( + self, + monkeypatch: MonkeyPatch, + execution_stage: ExecutionStage, + expected_name: str, + ) -> None: + """Tool node: routes to success and passes the expected payload to evaluator.""" + guardrail = MagicMock() + guardrail.name = "Example" + fake = _patch_uipath(monkeypatch, validation_passed=True, reason=None) + + node_name, node = create_tool_guardrail_node( + guardrail=guardrail, + execution_stage=execution_stage, + success_node="ok", + failure_node="nope", + tool_name="my_tool", + ) + assert node_name == expected_name + + if execution_stage == ExecutionStage.PRE_EXECUTION: + state = AgentGuardrailsGraphState( + messages=[ + AIMessage( + content="", + tool_calls=[ + {"name": "my_tool", "args": {"x": 1}, "id": "call_1"} + ], + ) + ] + ) + cmd = await node(state) + assert cmd.goto == "ok" + assert cmd.update == {"guardrail_validation_result": None} + assert json.loads(fake.guardrails.last_text or "{}") == {"x": 1} + else: + state = AgentGuardrailsGraphState( + messages=[ToolMessage(content="tool output", tool_call_id="call_1")] + ) + cmd = await node(state) + assert cmd.goto == "ok" + assert cmd.update == {"guardrail_validation_result": None} + assert fake.guardrails.last_text == "tool output" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "execution_stage,expected_name", + [ + (ExecutionStage.PRE_EXECUTION, "tool_pre_execution_example"), + (ExecutionStage.POST_EXECUTION, "tool_post_execution_example"), + ], + ids=["pre-fail", "post-fail"], + ) + async def test_tool_failure_pre_and_post( + self, + monkeypatch: MonkeyPatch, + execution_stage: ExecutionStage, + expected_name: str, + ) -> None: + """Tool node: routes to failure and sets guardrail_validation_result.""" + guardrail = MagicMock() + guardrail.name = "Example" + _patch_uipath(monkeypatch, validation_passed=False, reason="policy_violation") + + node_name, node = create_tool_guardrail_node( + guardrail=guardrail, + execution_stage=execution_stage, + success_node="ok", + failure_node="nope", + tool_name="my_tool", + ) + assert node_name == expected_name + + if execution_stage == ExecutionStage.PRE_EXECUTION: + state = AgentGuardrailsGraphState( + messages=[ + AIMessage( + content="", + tool_calls=[ + {"name": "my_tool", "args": {"x": 1}, "id": "call_1"} + ], + ) + ] + ) + else: + state = AgentGuardrailsGraphState( + messages=[ToolMessage(content="tool output", tool_call_id="call_1")] + ) + + cmd = await node(state) + assert cmd.goto == "nope" + assert cmd.update == {"guardrail_validation_result": "policy_violation"} diff --git a/tests/agent/guardrails/test_guardrail_stage_filtering.py b/tests/agent/guardrails/test_guardrail_stage_filtering.py index d068ecc4..fbea6824 100644 --- a/tests/agent/guardrails/test_guardrail_stage_filtering.py +++ b/tests/agent/guardrails/test_guardrail_stage_filtering.py @@ -4,7 +4,7 @@ from uipath.platform.guardrails import BuiltInValidatorGuardrail, DeterministicGuardrail -import uipath_langchain.agent.guardrails.guardrails_subgraph +import uipath_langchain.agent.react.guardrails.guardrails_subgraph from uipath_langchain.agent.guardrails.actions.base_action import GuardrailAction from uipath_langchain.agent.guardrails.types import ExecutionStage @@ -39,7 +39,7 @@ def test_filter_by_stage_prompt_injection(self): # --- PRE_EXECUTION --- # Should get ALL guardrails - pre_filtered = uipath_langchain.agent.guardrails.guardrails_subgraph._filter_guardrails_by_stage( + pre_filtered = uipath_langchain.agent.react.guardrails.guardrails_subgraph._filter_guardrails_by_stage( guardrails, ExecutionStage.PRE_EXECUTION ) assert len(pre_filtered) == 3 @@ -49,7 +49,7 @@ def test_filter_by_stage_prompt_injection(self): # --- POST_EXECUTION --- # Should SKIP Prompt Injection - post_filtered = uipath_langchain.agent.guardrails.guardrails_subgraph._filter_guardrails_by_stage( + post_filtered = uipath_langchain.agent.react.guardrails.guardrails_subgraph._filter_guardrails_by_stage( guardrails, ExecutionStage.POST_EXECUTION ) assert len(post_filtered) == 2 diff --git a/tests/agent/guardrails/test_guardrail_utils.py b/tests/agent/guardrails/test_guardrail_utils.py new file mode 100644 index 00000000..4025c7ae --- /dev/null +++ b/tests/agent/guardrails/test_guardrail_utils.py @@ -0,0 +1,105 @@ +"""Shared utilities for guardrail tests.""" + +from __future__ import annotations + +import types +from typing import Any, Callable, Protocol + +from uipath.platform.guardrails import BaseGuardrail + +from uipath_langchain.agent.guardrails.actions.base_action import ( + GuardrailAction, + GuardrailActionNode, +) + + +class _CompiledGraph(Protocol): + """Protocol for compiled fake graphs used in tests.""" + + nodes: list[tuple[str, Any]] + edges: list[tuple[str, str]] + + async def ainvoke(self, state: Any) -> dict[str, Any]: + """Invoke the compiled graph.""" + + +class FakeStateGraph: + """Minimal fake of LangGraph's StateGraph used by guardrail-subgraph unit tests.""" + + def __init__(self, _state_type: Any) -> None: + """Create a fake graph container.""" + self.added_nodes: list[tuple[str, Any]] = [] + self.added_edges: list[tuple[str, str]] = [] + + def add_node(self, name: str, node: Any) -> None: + """Record a node added to the graph.""" + self.added_nodes.append((name, node)) + + def add_edge(self, src: str, dst: str) -> None: + """Record an edge added to the graph.""" + self.added_edges.append((src, dst)) + + def compile(self) -> Any: + """Compile the fake graph into an inspectable object.""" + # Return a simple object we can inspect if needed. + return types.SimpleNamespace(nodes=self.added_nodes, edges=self.added_edges) + + +class FakeStateGraphWithAinvoke(FakeStateGraph): + """Fake graph that exposes an async `ainvoke` on the compiled output. + + Useful for wrapper functions that expect `compiled_graph.ainvoke(...)` + (e.g., terminate subgraph wrappers). + """ + + def __init__(self, _state_type: Any) -> None: + """Create a fake graph container, tracking the first-added node as the entry node.""" + super().__init__(_state_type) + self._main_node_name: str | None = None + + def add_node(self, name: str, node: Any) -> None: + """Record a node and remember the first-added node name as the main node.""" + if self._main_node_name is None: + self._main_node_name = name + super().add_node(name, node) + + def compile(self) -> _CompiledGraph: + """Compile the fake graph and attach an async `ainvoke` that runs the main node.""" + compiled = super().compile() + compiled.main_node_name = self._main_node_name + + async def ainvoke(state: Any) -> dict[str, Any]: + node_map: dict[str, Callable[[Any], dict[str, Any]]] = { + n: fn for n, fn in compiled.nodes + } + main_node_name = compiled.main_node_name + if main_node_name is None: + raise RuntimeError("FakeStateGraphWithAinvoke has no nodes to invoke.") + return node_map[main_node_name](state) + + compiled.ainvoke = ainvoke + return compiled + + +def fake_action(fail_prefix: str) -> GuardrailAction: + class _Action(GuardrailAction): + def action_node( + self, + *, + guardrail: BaseGuardrail, + scope, + execution_stage, + guarded_component_name: str, + ) -> GuardrailActionNode: + name = f"{fail_prefix}_{execution_stage.name.lower()}_{guardrail.name}" + return name, lambda s: s + + return _Action() + + +def fake_factory(eval_prefix): + def _factory(guardrail, execution_stage, success_node, failure_node, **_kwargs): + name = f"{eval_prefix}_{execution_stage.name.lower()}_{guardrail.name}" + return name, (lambda s: s) # node function not invoked in this test + + return _factory diff --git a/tests/agent/guardrails/test_guardrails_subgraph.py b/tests/agent/guardrails/test_guardrails_subgraph.py deleted file mode 100644 index 7ea1987f..00000000 --- a/tests/agent/guardrails/test_guardrails_subgraph.py +++ /dev/null @@ -1,230 +0,0 @@ -"""Tests for guardrails subgraph construction.""" - -import types -from typing import Sequence -from unittest.mock import MagicMock - -from uipath.platform.guardrails import ( - BaseGuardrail, - GuardrailScope, -) - -import uipath_langchain.agent.guardrails.guardrails_subgraph as mod -from uipath_langchain.agent.guardrails.actions.base_action import ( - GuardrailAction, - GuardrailActionNode, -) -from uipath_langchain.agent.guardrails.types import ExecutionStage - - -class FakeStateGraph: - def __init__(self, _state_type): - self.added_nodes = [] - self.added_edges = [] - - def add_node(self, name, node): - self.added_nodes.append((name, node)) - - def add_edge(self, src, dst): - self.added_edges.append((src, dst)) - - def compile(self): - # Return a simple object we can inspect if needed - return types.SimpleNamespace(nodes=self.added_nodes, edges=self.added_edges) - - -def _fake_action(fail_prefix: str) -> GuardrailAction: - class _Action(GuardrailAction): - def action_node( - self, - *, - guardrail: BaseGuardrail, - scope, - execution_stage, - guarded_component_name: str, - ) -> GuardrailActionNode: - name = f"{fail_prefix}_{execution_stage.name.lower()}_{guardrail.name}" - return name, lambda s: s - - return _Action() - - -def _fake_factory(eval_prefix): - def _factory(guardrail, execution_stage, success_node, failure_node): - name = f"{eval_prefix}_{execution_stage.name.lower()}_{guardrail.name}" - return name, (lambda s: s) # node function not invoked in this test - - return _factory - - -class TestLlmGuardrailsSubgraph: - def test_no_applicable_guardrails_returns_original_node(self): - """If no guardrails match the scope, the original node should be returned.""" - inner = ("inner", lambda s: s) - guardrails: Sequence[tuple[BaseGuardrail, GuardrailAction]] = [] - - # Should return the inner node function directly, not a compiled graph - result = mod.create_llm_guardrails_subgraph( - llm_node=inner, guardrails=guardrails - ) - assert result == inner[1] - - # Case with None guardrails - result_none = mod.create_llm_guardrails_subgraph( - llm_node=inner, guardrails=None - ) - assert result_none == inner[1] - - # Case with guardrails but none matching LLM scope - non_matching_guardrail = MagicMock() - non_matching_guardrail.selector = types.SimpleNamespace( - scopes=[GuardrailScope.TOOL] - ) - guardrails_non_match = [(non_matching_guardrail, MagicMock())] - - result_non_match = mod.create_llm_guardrails_subgraph( - llm_node=inner, guardrails=guardrails_non_match - ) - assert result_non_match == inner[1] - - def test_two_guardrails_build_chains_pre_and_post(self, monkeypatch): - """Two guardrails should create reverse-ordered pre/post chains with failure edges.""" - monkeypatch.setattr(mod, "StateGraph", FakeStateGraph) - monkeypatch.setattr(mod, "START", "START") - monkeypatch.setattr(mod, "END", "END") - # Use fake factory to control eval node names - monkeypatch.setattr(mod, "create_llm_guardrail_node", _fake_factory("eval")) - - # Guardrails g1 (first), g2 (second); builder processes last first - guardrail1 = MagicMock() - guardrail1.name = "guardrail1" - guardrail1.selector = types.SimpleNamespace(scopes=[GuardrailScope.LLM]) - - guardrail2 = MagicMock() - guardrail2.name = "guardrail2" - guardrail2.selector = types.SimpleNamespace(scopes=[GuardrailScope.LLM]) - - a1 = _fake_action("log") - a2 = _fake_action("block") - guardrails = [(guardrail1, a1), (guardrail2, a2)] - - inner = ("inner", lambda s: s) - compiled = mod.create_llm_guardrails_subgraph( - llm_node=inner, - guardrails=guardrails, - ) - - # Expected node names - pre_g1 = "eval_pre_execution_guardrail1" - log_pre_g1 = "log_pre_execution_guardrail1" - pre_g2 = "eval_pre_execution_guardrail2" - block_pre_g2 = "block_pre_execution_guardrail2" - post_g1 = "eval_post_execution_guardrail1" - log_post_g1 = "log_post_execution_guardrail1" - post_g2 = "eval_post_execution_guardrail2" - block_post_g2 = "block_post_execution_guardrail2" - - # Edges (order not guaranteed; compare as a set) - expected_edges = { - # Pre-execution chain - ("START", pre_g1), - (log_pre_g1, pre_g2), - (block_pre_g2, "inner"), - # Inner to post-execution chain - ("inner", post_g1), - # Post-execution failure routing to END - (log_post_g1, post_g2), - (block_post_g2, "END"), - } - assert expected_edges.issubset(set(compiled.edges)) - - # Ensure expected nodes are present - node_names = {name for name, _ in compiled.nodes} - for name in [ - pre_g1, - pre_g2, - post_g1, - post_g2, - log_pre_g1, - block_pre_g2, - log_post_g1, - block_post_g2, - "inner", - ]: - assert name in node_names - - -class TestAgentGuardrailsSubgraph: - def test_no_applicable_guardrails_returns_original_node(self): - """If no guardrails match the AGENT scope, the original node should be returned.""" - inner = ("inner", lambda s: s) - guardrails: Sequence[tuple[BaseGuardrail, GuardrailAction]] = [] - stage = ExecutionStage.PRE_EXECUTION - - # Should return the inner node function directly - result = mod.create_agent_guardrails_subgraph( - agent_node=inner, guardrails=guardrails, execution_stage=stage - ) - assert result == inner[1] - - # Case with None guardrails - result_none = mod.create_agent_guardrails_subgraph( - agent_node=inner, guardrails=None, execution_stage=stage - ) - assert result_none == inner[1] - - # Case with guardrails but none matching AGENT scope - non_matching_guardrail = MagicMock() - non_matching_guardrail.selector = types.SimpleNamespace( - scopes=[GuardrailScope.LLM] - ) - guardrails_non_match = [(non_matching_guardrail, MagicMock())] - - result_non_match = mod.create_agent_guardrails_subgraph( - agent_node=inner, guardrails=guardrails_non_match, execution_stage=stage - ) - assert result_non_match == inner[1] - - -class TestToolGuardrailsSubgraph: - def test_no_applicable_guardrails_returns_original_node(self): - """If no guardrails match the TOOL scope/name, the original node should be returned.""" - tool_name = "test_tool" - inner = (tool_name, lambda s: s) - guardrails: Sequence[tuple[BaseGuardrail, GuardrailAction]] = [] - - # Should return the inner node function directly - result = mod.create_tool_guardrails_subgraph( - tool_node=inner, guardrails=guardrails - ) - assert result == inner[1] - - # Case with None guardrails - result_none = mod.create_tool_guardrails_subgraph( - tool_node=inner, guardrails=None - ) - assert result_none == inner[1] - - # Case with guardrails matching TOOL scope but NOT the tool name - wrong_name_guardrail = MagicMock() - wrong_name_guardrail.selector = types.SimpleNamespace( - scopes=[GuardrailScope.TOOL], match_names=["other_tool"] - ) - guardrails_wrong_name = [(wrong_name_guardrail, MagicMock())] - - result_wrong_name = mod.create_tool_guardrails_subgraph( - tool_node=inner, guardrails=guardrails_wrong_name - ) - assert result_wrong_name == inner[1] - - # Case with guardrails matching tool name but NOT TOOL scope (unlikely but good to test) - wrong_scope_guardrail = MagicMock() - wrong_scope_guardrail.selector = types.SimpleNamespace( - scopes=[GuardrailScope.LLM], match_names=[tool_name] - ) - guardrails_wrong_scope = [(wrong_scope_guardrail, MagicMock())] - - result_wrong_scope = mod.create_tool_guardrails_subgraph( - tool_node=inner, guardrails=guardrails_wrong_scope - ) - assert result_wrong_scope == inner[1] diff --git a/tests/agent/guardrails/test_llm_guardrails_subgraph.py b/tests/agent/guardrails/test_llm_guardrails_subgraph.py new file mode 100644 index 00000000..11dd72fd --- /dev/null +++ b/tests/agent/guardrails/test_llm_guardrails_subgraph.py @@ -0,0 +1,113 @@ +"""Tests for LLM guardrails subgraph construction.""" + +import types +from typing import Sequence +from unittest.mock import MagicMock + +from uipath.core.guardrails import BaseGuardrail, GuardrailSelector +from uipath.platform.guardrails import GuardrailScope + +import uipath_langchain.agent.react.guardrails.guardrails_subgraph as mod +from tests.agent.guardrails.test_guardrail_utils import ( + FakeStateGraph, + fake_action, + fake_factory, +) +from uipath_langchain.agent.guardrails.actions import GuardrailAction + + +class TestLlmGuardrailsSubgraph: + def test_no_applicable_guardrails_returns_original_node(self): + """If no guardrails match the scope, the original node should be returned.""" + inner = ("inner", lambda s: s) + guardrails: Sequence[tuple[BaseGuardrail, GuardrailAction]] = [] + + # Case with empty guardrails + result = mod.create_llm_guardrails_subgraph( + llm_node=inner, guardrails=guardrails + ) + assert result == inner[1] + + # Case with None guardrails + result_none = mod.create_llm_guardrails_subgraph( + llm_node=inner, guardrails=None + ) + assert result_none == inner[1] + + # Case with guardrails but none matching LLM scope + non_matching_guardrail = MagicMock() + non_matching_guardrail.selector = types.SimpleNamespace( + scopes=[GuardrailScope.TOOL] + ) + guardrails_non_match = [(non_matching_guardrail, MagicMock())] + + result_non_match = mod.create_llm_guardrails_subgraph( + llm_node=inner, guardrails=guardrails_non_match + ) + assert result_non_match == inner[1] + + def test_two_guardrails_build_chains_pre_and_post(self, monkeypatch): + """Two guardrails should create reverse-ordered pre/post chains with failure edges.""" + monkeypatch.setattr(mod, "StateGraph", FakeStateGraph) + monkeypatch.setattr(mod, "START", "START") + monkeypatch.setattr(mod, "END", "END") + # Use fake factory to control eval node names + monkeypatch.setattr(mod, "create_llm_guardrail_node", fake_factory("eval")) + + # Guardrails g1 (first), g2 (second); builder processes last first + guardrail1 = MagicMock() + guardrail1.name = "guardrail1" + guardrail1.selector = GuardrailSelector(scopes=[GuardrailScope.LLM]) + + guardrail2 = MagicMock() + guardrail2.name = "guardrail2" + guardrail2.selector = GuardrailSelector(scopes=[GuardrailScope.LLM]) + + a1 = fake_action("log") + a2 = fake_action("block") + guardrails = [(guardrail1, a1), (guardrail2, a2)] + + inner = ("inner", lambda s: s) + compiled = mod.create_llm_guardrails_subgraph( + llm_node=inner, + guardrails=guardrails, + ) + + # Expected node names + pre_g1 = "eval_pre_execution_guardrail1" + log_pre_g1 = "log_pre_execution_guardrail1" + pre_g2 = "eval_pre_execution_guardrail2" + block_pre_g2 = "block_pre_execution_guardrail2" + post_g1 = "eval_post_execution_guardrail1" + log_post_g1 = "log_post_execution_guardrail1" + post_g2 = "eval_post_execution_guardrail2" + block_post_g2 = "block_post_execution_guardrail2" + + # Edges (order not guaranteed; compare as a set) + expected_edges = { + # Pre-execution chain + ("START", pre_g1), + (log_pre_g1, pre_g2), + (block_pre_g2, "inner"), + # Inner to post-execution chain + ("inner", post_g1), + # Post-execution failure routing to END + (log_post_g1, post_g2), + (block_post_g2, "END"), + } + assert expected_edges.issubset(set(compiled.edges)) + + # Ensure expected nodes are present + node_names = {name for name, _ in compiled.nodes} + for name in [ + pre_g1, + pre_g2, + post_g1, + post_g2, + log_pre_g1, + block_pre_g2, + log_post_g1, + block_post_g2, + "inner", + ]: + assert name in node_names diff --git a/tests/agent/guardrails/test_tool_guardrails_subgraph.py b/tests/agent/guardrails/test_tool_guardrails_subgraph.py new file mode 100644 index 00000000..e0d905fc --- /dev/null +++ b/tests/agent/guardrails/test_tool_guardrails_subgraph.py @@ -0,0 +1,131 @@ +"""Tests for Tool guardrails subgraph construction.""" + +import types +from typing import Sequence +from unittest.mock import MagicMock + +from uipath.core.guardrails import BaseGuardrail, GuardrailSelector +from uipath.platform.guardrails import GuardrailScope + +import uipath_langchain.agent.react.guardrails.guardrails_subgraph as mod +from tests.agent.guardrails.test_guardrail_utils import ( + FakeStateGraph, + fake_action, + fake_factory, +) +from uipath_langchain.agent.guardrails.actions import GuardrailAction + + +class TestToolGuardrailsSubgraph: + def test_no_applicable_guardrails_returns_original_node(self): + tool_name = "test_tool" + inner = (tool_name, lambda s: s) + guardrails: Sequence[tuple[BaseGuardrail, GuardrailAction]] = [] + + # Case with empty guardrails + result = mod.create_tool_guardrails_subgraph( + tool_node=inner, guardrails=guardrails + ) + assert result == inner[1] + + # Case with None guardrails + result_none = mod.create_tool_guardrails_subgraph( + tool_node=inner, guardrails=None + ) + assert result_none == inner[1] + + # Case with guardrails matching TOOL scope but NOT the tool name + wrong_name_guardrail = MagicMock() + wrong_name_guardrail.selector = types.SimpleNamespace( + scopes=[GuardrailScope.TOOL], match_names=["other_tool"] + ) + guardrails_wrong_name = [(wrong_name_guardrail, MagicMock())] + + result_wrong_name = mod.create_tool_guardrails_subgraph( + tool_node=inner, guardrails=guardrails_wrong_name + ) + assert result_wrong_name == inner[1] + + # Case with guardrails matching tool name but NOT TOOL scope (unlikely but good to test) + wrong_scope_guardrail = MagicMock() + wrong_scope_guardrail.selector = types.SimpleNamespace( + scopes=[GuardrailScope.LLM], match_names=[tool_name] + ) + guardrails_wrong_scope = [(wrong_scope_guardrail, MagicMock())] + + result_wrong_scope = mod.create_tool_guardrails_subgraph( + tool_node=inner, guardrails=guardrails_wrong_scope + ) + assert result_wrong_scope == inner[1] + + def test_two_guardrails_build_chains_pre_and_post(self, monkeypatch): + """Two guardrails should create reverse-ordered pre/post chains with failure edges.""" + monkeypatch.setattr(mod, "StateGraph", FakeStateGraph) + monkeypatch.setattr(mod, "START", "START") + monkeypatch.setattr(mod, "END", "END") + # Use fake factory to control eval node names + monkeypatch.setattr(mod, "create_tool_guardrail_node", fake_factory("eval")) + + tool_name = "test_tool" + + # Guardrails g1 (first), g2 (second); builder processes last first + guardrail1 = MagicMock() + guardrail1.name = "guardrail1" + guardrail1.selector = GuardrailSelector( + scopes=[GuardrailScope.TOOL], match_names=[tool_name] + ) + + guardrail2 = MagicMock() + guardrail2.name = "guardrail2" + guardrail2.selector = GuardrailSelector( + scopes=[GuardrailScope.TOOL], match_names=[tool_name] + ) + + a1 = fake_action("log") + a2 = fake_action("block") + guardrails = [(guardrail1, a1), (guardrail2, a2)] + + inner = (tool_name, lambda s: s) + compiled = mod.create_tool_guardrails_subgraph( + tool_node=inner, + guardrails=guardrails, + ) + + # Expected node names + pre_g1 = "eval_pre_execution_guardrail1" + log_pre_g1 = "log_pre_execution_guardrail1" + pre_g2 = "eval_pre_execution_guardrail2" + block_pre_g2 = "block_pre_execution_guardrail2" + post_g1 = "eval_post_execution_guardrail1" + log_post_g1 = "log_post_execution_guardrail1" + post_g2 = "eval_post_execution_guardrail2" + block_post_g2 = "block_post_execution_guardrail2" + + # Edges (order not guaranteed; compare as a set) + expected_edges = { + # Pre-execution chain + ("START", pre_g1), + (log_pre_g1, pre_g2), + (block_pre_g2, tool_name), + # Inner to post-execution chain + (tool_name, post_g1), + # Post-execution failure routing to END + (log_post_g1, post_g2), + (block_post_g2, "END"), + } + assert expected_edges.issubset(set(compiled.edges)) + + # Ensure expected nodes are present + node_names = {name for name, _ in compiled.nodes} + for name in [ + pre_g1, + pre_g2, + post_g1, + post_g2, + log_pre_g1, + block_pre_g2, + log_post_g1, + block_post_g2, + tool_name, + ]: + assert name in node_names diff --git a/uv.lock b/uv.lock index 27b0fc44..4e249ff3 100644 --- a/uv.lock +++ b/uv.lock @@ -2863,7 +2863,7 @@ wheels = [ [[package]] name = "uipath-langchain" -version = "0.1.29" +version = "0.1.30" source = { editable = "." } dependencies = [ { name = "aiosqlite" },