Skip to content

Commit d757edf

Browse files
author
Valentina Bojan
committed
feat(guardrails): forward agent_input to deterministic pre-evaluation
When an input_schema is configured, the agent's validated input parameters are extracted from graph state and passed to DeterministicGuardrailsService.evaluate_pre_deterministic_guardrail so rules can reference FieldSource.AGENT_INPUT (uipath-core 0.5.14+). Plumbs input_schema through create_tool_guardrail_node / _create_guardrail_node; extraction reuses extract_input_data_from_state. Post-execution is unchanged — agent_input rules are pre-only by design. AL-410 / AL-405 (Phase A).
1 parent a3c7f53 commit d757edf

3 files changed

Lines changed: 140 additions & 3 deletions

File tree

src/uipath_langchain/agent/guardrails/guardrail_nodes.py

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
from typing import Any, Callable
55

66
from langgraph.types import Command
7+
from pydantic import BaseModel
78
from uipath.core.guardrails import (
89
DeterministicGuardrail,
910
DeterministicGuardrailsService,
@@ -26,18 +27,35 @@
2627
get_message_content,
2728
)
2829
from uipath_langchain.agent.react.types import AgentGuardrailsGraphState
30+
from uipath_langchain.agent.react.utils import extract_input_data_from_state
2931

3032
from ..exceptions import AgentRuntimeError, AgentRuntimeErrorCode
3133

3234
logger = logging.getLogger(__name__)
3335

3436

37+
def _resolve_agent_input(
38+
state: AgentGuardrailsGraphState,
39+
input_schema: type[BaseModel] | None,
40+
) -> dict[str, Any] | None:
41+
if input_schema is None:
42+
return None
43+
try:
44+
return extract_input_data_from_state(state, input_schema)
45+
except Exception:
46+
# The state may not yet carry agent-input fields (e.g., very early
47+
# subgraphs, or schemas whose required fields aren't seeded yet); fall
48+
# back to "no agent_input available" rather than crashing the run.
49+
return None
50+
51+
3552
def _evaluate_deterministic_guardrail(
3653
state: AgentGuardrailsGraphState,
3754
guardrail: DeterministicGuardrail,
3855
execution_stage: ExecutionStage,
3956
input_data_extractor: Callable[[AgentGuardrailsGraphState], dict[str, Any]],
4057
output_data_extractor: Callable[[AgentGuardrailsGraphState], dict[str, Any]] | None,
58+
input_schema: type[BaseModel] | None = None,
4159
):
4260
"""Evaluate deterministic guardrail.
4361
@@ -47,6 +65,10 @@ def _evaluate_deterministic_guardrail(
4765
execution_stage: The execution stage (PRE_EXECUTION or POST_EXECUTION).
4866
input_data_extractor: Function to extract input data from state.
4967
output_data_extractor: Function to extract output data from state (optional).
68+
input_schema: Optional input schema; when provided, the agent's
69+
validated input parameters are extracted from state and passed to
70+
pre-execution evaluation so rules can reference
71+
``FieldSource.AGENT_INPUT``.
5072
5173
Returns:
5274
The guardrail evaluation result.
@@ -56,7 +78,9 @@ def _evaluate_deterministic_guardrail(
5678

5779
if execution_stage == ExecutionStage.PRE_EXECUTION:
5880
return service.evaluate_pre_deterministic_guardrail(
59-
input_data=input_data, guardrail=guardrail
81+
input_data=input_data,
82+
guardrail=guardrail,
83+
agent_input=_resolve_agent_input(state, input_schema),
6084
)
6185
else: # POST_EXECUTION
6286
output_data = output_data_extractor(state) if output_data_extractor else {}
@@ -150,6 +174,7 @@ def _create_guardrail_node(
150174
| None = None,
151175
tool_name: str | None = None,
152176
tool_type: str | None = None,
177+
input_schema: type[BaseModel] | None = None,
153178
) -> tuple[str, Callable[[AgentGuardrailsGraphState], Any]]:
154179
"""Private factory for guardrail evaluation nodes.
155180
@@ -195,6 +220,7 @@ async def node(
195220
execution_stage,
196221
input_data_extractor,
197222
output_data_extractor,
223+
input_schema,
198224
)
199225
elif isinstance(guardrail, BuiltInValidatorGuardrail):
200226
# Generate and store payload for observability
@@ -314,6 +340,7 @@ def create_tool_guardrail_node(
314340
failure_node: str,
315341
tool_name: str,
316342
tool_type: str | None = None,
343+
input_schema: type[BaseModel] | None = None,
317344
) -> tuple[str, Callable[[AgentGuardrailsGraphState], Any]]:
318345
"""Create a guardrail node for TOOL scope guardrails.
319346
@@ -324,6 +351,8 @@ def create_tool_guardrail_node(
324351
failure_node: Node to route to on validation fail.
325352
tool_name: Name of the tool to extract arguments from.
326353
tool_type: Optional type of the tool (e.g., "process", "escalation", "mcp").
354+
input_schema: Optional agent input schema; enables rules to reference
355+
``FieldSource.AGENT_INPUT`` during pre-execution evaluation.
327356
328357
Returns:
329358
A tuple of (node_name, node_function) for the guardrail evaluation node.
@@ -375,4 +404,5 @@ def _output_data_extractor(state: AgentGuardrailsGraphState) -> dict[str, Any]:
375404
_output_data_extractor,
376405
tool_name,
377406
tool_type,
407+
input_schema,
378408
)

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

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -452,7 +452,10 @@ def create_tool_guardrails_subgraph(
452452
scope=GuardrailScope.TOOL,
453453
execution_stages=[ExecutionStage.PRE_EXECUTION, ExecutionStage.POST_EXECUTION],
454454
node_factory=partial(
455-
create_tool_guardrail_node, tool_name=tool_name, tool_type=tool_type
455+
create_tool_guardrail_node,
456+
tool_name=tool_name,
457+
tool_type=tool_type,
458+
input_schema=input_schema,
456459
),
457460
input_schema=input_schema,
458461
)

tests/agent/guardrails/test_guardrail_nodes.py

Lines changed: 105 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -498,7 +498,7 @@ async def test_evaluate_deterministic_guardrail_pre_execution(self, monkeypatch)
498498

499499
assert result.result == GuardrailValidationResultType.PASSED
500500
mock_service.evaluate_pre_deterministic_guardrail.assert_called_once_with(
501-
input_data={"test": "data"}, guardrail=guardrail
501+
input_data={"test": "data"}, guardrail=guardrail, agent_input=None
502502
)
503503

504504
@pytest.mark.asyncio
@@ -544,6 +544,110 @@ async def test_evaluate_deterministic_guardrail_post_execution(self, monkeypatch
544544
guardrail=guardrail,
545545
)
546546

547+
@pytest.mark.asyncio
548+
async def test_evaluate_deterministic_guardrail_pre_passes_agent_input_when_schema_supplied(
549+
self, monkeypatch
550+
):
551+
"""When input_schema is supplied, the agent input dict is extracted from
552+
state and forwarded to evaluate_pre_deterministic_guardrail."""
553+
from pydantic import BaseModel
554+
from uipath.core.guardrails import DeterministicGuardrail
555+
556+
from uipath_langchain.agent.guardrails.guardrail_nodes import (
557+
_evaluate_deterministic_guardrail,
558+
)
559+
from uipath_langchain.agent.react.utils import (
560+
create_guardrails_state_with_input,
561+
)
562+
563+
class AgentInputSchema(BaseModel):
564+
user_identity: str
565+
role: str
566+
567+
mock_result = GuardrailValidationResult(
568+
result=GuardrailValidationResultType.PASSED, reason=""
569+
)
570+
mock_service = MagicMock()
571+
mock_service.evaluate_pre_deterministic_guardrail.return_value = mock_result
572+
monkeypatch.setattr(
573+
"uipath_langchain.agent.guardrails.guardrail_nodes.DeterministicGuardrailsService",
574+
lambda: mock_service,
575+
)
576+
577+
guardrail = MagicMock(spec=DeterministicGuardrail)
578+
state_cls = create_guardrails_state_with_input(AgentInputSchema)
579+
state = state_cls(messages=[], user_identity="U157877", role="admin")
580+
input_extractor = MagicMock(return_value={"target_user_id": "U157878"})
581+
output_extractor = MagicMock()
582+
583+
result = _evaluate_deterministic_guardrail(
584+
state,
585+
guardrail,
586+
ExecutionStage.PRE_EXECUTION,
587+
input_extractor,
588+
output_extractor,
589+
AgentInputSchema,
590+
)
591+
592+
assert result.result == GuardrailValidationResultType.PASSED
593+
mock_service.evaluate_pre_deterministic_guardrail.assert_called_once_with(
594+
input_data={"target_user_id": "U157878"},
595+
guardrail=guardrail,
596+
agent_input={"user_identity": "U157877", "role": "admin"},
597+
)
598+
599+
@pytest.mark.asyncio
600+
async def test_evaluate_deterministic_guardrail_post_does_not_pass_agent_input(
601+
self, monkeypatch
602+
):
603+
"""Even with input_schema supplied, post-execution does NOT receive
604+
agent_input — agent-input rules are pre-execution only by design."""
605+
from pydantic import BaseModel
606+
from uipath.core.guardrails import DeterministicGuardrail
607+
608+
from uipath_langchain.agent.guardrails.guardrail_nodes import (
609+
_evaluate_deterministic_guardrail,
610+
)
611+
from uipath_langchain.agent.react.utils import (
612+
create_guardrails_state_with_input,
613+
)
614+
615+
class AgentInputSchema(BaseModel):
616+
user_identity: str
617+
618+
mock_service = MagicMock()
619+
mock_service.evaluate_post_deterministic_guardrail.return_value = (
620+
GuardrailValidationResult(
621+
result=GuardrailValidationResultType.PASSED, reason=""
622+
)
623+
)
624+
monkeypatch.setattr(
625+
"uipath_langchain.agent.guardrails.guardrail_nodes.DeterministicGuardrailsService",
626+
lambda: mock_service,
627+
)
628+
629+
guardrail = MagicMock(spec=DeterministicGuardrail)
630+
state_cls = create_guardrails_state_with_input(AgentInputSchema)
631+
state = state_cls(messages=[], user_identity="U157877")
632+
633+
_evaluate_deterministic_guardrail(
634+
state,
635+
guardrail,
636+
ExecutionStage.POST_EXECUTION,
637+
MagicMock(return_value={"input": "data"}),
638+
MagicMock(return_value={"output": "data"}),
639+
AgentInputSchema,
640+
)
641+
642+
mock_service.evaluate_post_deterministic_guardrail.assert_called_once_with(
643+
input_data={"input": "data"},
644+
output_data={"output": "data"},
645+
guardrail=guardrail,
646+
)
647+
# agent_input must NOT be present in the post call
648+
kwargs = mock_service.evaluate_post_deterministic_guardrail.call_args.kwargs
649+
assert "agent_input" not in kwargs
650+
547651
@pytest.mark.asyncio
548652
async def test_evaluate_builtin_guardrail(self, monkeypatch):
549653
"""Test built-in guardrail evaluation."""

0 commit comments

Comments
 (0)