diff --git a/pyproject.toml b/pyproject.toml index d2ac71462..e74b80dfd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "uipath-langchain" -version = "0.18.12" +version = "0.18.13" 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/attachment_refs.py b/src/uipath_langchain/agent/guardrails/attachment_refs.py index eace068f1..a0353cab0 100644 --- a/src/uipath_langchain/agent/guardrails/attachment_refs.py +++ b/src/uipath_langchain/agent/guardrails/attachment_refs.py @@ -1,18 +1,30 @@ -"""Project a run's job-attachment registry into guardrail attachment references. +"""Project a run's attachments into guardrail attachment references. -Any built-in guardrail forwards the run's attachments unless it is scoped to prompts. The -runtime forwards id, file name and mime type only; the backend's feature flag decides -whether they are used at all, and the backend decides which validators and file types it -can inspect and resolves the id through Orchestrator. +Two sources feed the same reference shape: -Nothing in this module raises: the guardrail node re-raises any exception, which would end -the run over a single malformed attachment. +* Agent- and LLM-scope guardrails judge the conversation so far, so they read the whole + job-attachment registry (:func:`resolve_guardrail_attachments`). +* Tool-scope guardrails judge one tool call, so they read only the attachments that call + mentions: the ``{"ID": ...}`` objects in its arguments before the tool runs, or in its + result afterwards (:func:`resolve_referenced_attachments`). A tool call that names no + file forwards nothing, even when the run holds files elsewhere; otherwise every tool + call would ship every file to the backend. + +Any built-in guardrail forwards references unless it is scoped to prompts. The runtime +forwards id, file name and mime type only; the backend's feature flag decides whether +they are used at all, the backend decides which validators and file types it can +inspect, and it resolves the id through Orchestrator. + +Nothing in this module raises: the guardrail node re-raises any exception, which would +end the run over a single malformed attachment. """ import logging import uuid +from collections.abc import Iterable, Iterator, Mapping from typing import Any +from pydantic import BaseModel from uipath.platform.attachments import Attachment from uipath.platform.guardrails import BuiltInValidatorGuardrail, GuardrailAttachment @@ -24,6 +36,11 @@ #: ``appliesTo`` guardrail parameter; only ``Prompts`` excludes files (default is ``Both``). _APPLIES_TO_PARAMETER = "appliesto" _PROMPTS_ONLY = "prompts" +#: Wire key of a job attachment reference as the model and the tools exchange it. +_ID_KEY = "ID" +#: Bounds for scanning a tool payload, which can be arbitrarily large or deep. +_MAX_SCAN_DEPTH = 32 +_MAX_SCAN_NODES = 10_000 def _scope_includes_files(guardrail: BuiltInValidatorGuardrail) -> bool: @@ -44,8 +61,9 @@ async def resolve_guardrail_attachments( job_attachments: dict[str, Attachment], guardrail: BuiltInValidatorGuardrail, ) -> list[GuardrailAttachment]: - """Return up to five attachment references for the guardrail, or an empty list. + """Return up to five references for every attachment the run knows about. + For Agent- and LLM-scope guardrails, which evaluate the conversation as a whole. Empty when the guardrail is scoped to prompts or the run has no attachments. Never raises. """ @@ -57,14 +75,119 @@ async def resolve_guardrail_attachments( guardrail.name, ) return [] + return _collect(job_attachments.values()) + + +def resolve_referenced_attachments( + data: Any, + job_attachments: dict[str, Attachment] | None, + guardrail: BuiltInValidatorGuardrail, +) -> list[GuardrailAttachment]: + """Return up to five references for the attachments ``data`` mentions. + + For Tool-scope guardrails: ``data`` is the tool call's arguments (before the tool + runs) or its parsed result (after). A mention is a mapping with an ``ID`` that parses + as a UUID, the shape the model emits and the tool wrapper expands. Only ids the run's + registry holds are forwarded, with the registry's name and type: the registry is the + set of files this run legitimately has (agent input plus files its tools returned), + so a mention the run never held, or a non-attachment resource id, is skipped rather + than sent to the backend for lookup. Empty when nothing is mentioned or the guardrail + is scoped to prompts. Never raises. + """ + if data is None: + return [] + if not _scope_includes_files(guardrail): + logger.debug( + "Guardrail '%s' is scoped to prompts; skipping attachment resolution.", + guardrail.name, + ) + return [] + try: + mentions = list(_iter_attachment_mentions(data)) + except Exception: + logger.warning( + "Could not scan the tool payload for attachments; guardrail '%s' evaluates " + "without files.", + guardrail.name, + exc_info=True, + ) + return [] + registry = job_attachments or {} + resolved = (_resolve_mention(mention, registry) for mention in mentions) + return _collect(attachment for attachment in resolved if attachment is not None) + + +def _iter_attachment_mentions(data: Any) -> Iterator[Any]: + """Yield every attachment-shaped value in ``data``, depth first, within bounds.""" + stack: list[tuple[Any, int]] = [(data, 0)] + visited = 0 + while stack: + value, depth = stack.pop() + visited += 1 + if visited > _MAX_SCAN_NODES: + logger.debug( + "Stopped scanning the tool payload for attachments after %d values.", + _MAX_SCAN_NODES, + ) + return + if depth > _MAX_SCAN_DEPTH: + continue + if isinstance(value, Attachment): + yield value + elif isinstance(value, BaseModel): + stack.append((value.model_dump(by_alias=True), depth + 1)) + elif isinstance(value, Mapping): + if _is_mention(value): + yield value + else: + stack.extend((item, depth + 1) for item in value.values()) + elif isinstance(value, (list, tuple, set, frozenset)): + stack.extend((item, depth + 1) for item in value) + + +def _is_mention(value: Mapping[Any, Any]) -> bool: + raw_id = value.get(_ID_KEY) + if raw_id is None or isinstance(raw_id, bool): + return False + try: + uuid.UUID(str(raw_id)) + except ValueError: + return False + return True + + +def _resolve_mention(mention: Any, registry: dict[str, Attachment]) -> Any | None: + """Look one mention up in the run's registry; None when the run never held it.""" + try: + raw_id = mention.id if isinstance(mention, Attachment) else mention[_ID_KEY] + attachment_id = str(uuid.UUID(str(raw_id))) + known = registry.get(attachment_id) + if known is not None: + return known + logger.debug( + "Skipping attachment reference '%s': this run does not hold it.", + attachment_id, + ) + except Exception: + logger.debug( + "Skipping a malformed attachment reference in the tool payload.", + exc_info=True, + ) + return None + +def _collect(attachments: Iterable[Any]) -> list[GuardrailAttachment]: + """Build references, dropping unusable and duplicate ones, capped to the API limit.""" references: list[GuardrailAttachment] = [] - for attachment in job_attachments.values(): + seen: set[str] = set() + for attachment in attachments: reference = _to_reference(attachment) - if reference is not None: - references.append(reference) - if len(references) == _MAX_ATTACHMENTS: - break + if reference is None or reference.id in seen: + continue + seen.add(reference.id) + references.append(reference) + if len(references) == _MAX_ATTACHMENTS: + break return references diff --git a/src/uipath_langchain/agent/guardrails/guardrail_nodes.py b/src/uipath_langchain/agent/guardrails/guardrail_nodes.py index be96eb695..5a918934d 100644 --- a/src/uipath_langchain/agent/guardrails/guardrail_nodes.py +++ b/src/uipath_langchain/agent/guardrails/guardrail_nodes.py @@ -4,6 +4,7 @@ import re from typing import Any, Callable +from langchain_core.messages import AIMessage, ToolMessage from langgraph.types import Command from uipath.core.guardrails import ( DeterministicGuardrail, @@ -23,6 +24,7 @@ from uipath_langchain.agent.guardrails.attachment_refs import ( resolve_guardrail_attachments, + resolve_referenced_attachments, ) from uipath_langchain.agent.guardrails.types import ExecutionStage from uipath_langchain.agent.guardrails.utils import ( @@ -32,14 +34,15 @@ get_message_content, ) from uipath_langchain.agent.react.types import AgentGuardrailsGraphState +from uipath_langchain.agent.react.utils import ( + extract_current_tool_call_index, + find_latest_ai_message, +) from ..exceptions import AgentRuntimeError, AgentRuntimeErrorCode logger = logging.getLogger(__name__) -#: Scopes whose guardrails may inspect attached files (tool scope excluded on purpose). -_ATTACHMENT_SCOPES = frozenset({GuardrailScope.AGENT, GuardrailScope.LLM}) - def _evaluate_deterministic_guardrail( state: AgentGuardrailsGraphState, @@ -168,6 +171,48 @@ def _create_validation_command( ) +async def _resolve_attachments( + state: AgentGuardrailsGraphState, + guardrail: BuiltInValidatorGuardrail, + scope: GuardrailScope, + execution_stage: ExecutionStage, + input_data_extractor: Callable[[AgentGuardrailsGraphState], dict[str, Any]] | None, + output_data_extractor: Callable[[AgentGuardrailsGraphState], dict[str, Any]] | None, +) -> list[GuardrailAttachment]: + """Attachment references for one built-in guardrail evaluation. + + Agent and LLM scope judge the conversation, so they read the run's whole attachment + registry. Tool scope judges one call, so it reads only the files that call's + arguments (pre-execution) or result (post-execution) mention: a tool call without a + file forwards nothing even when the run has files elsewhere. Never raises. + """ + registry = state.inner_state.job_attachments + if scope != GuardrailScope.TOOL: + return await resolve_guardrail_attachments(registry, guardrail) + + if execution_stage == ExecutionStage.PRE_EXECUTION: + extractor, source_name = input_data_extractor, "arguments" + else: + extractor, source_name = output_data_extractor, "result" + # A mention always carries the literal ``ID`` key; skip parsing plain-text + # results, which the output extractor would otherwise warn about. + if not state.messages or "ID" not in get_message_content(state.messages[-1]): + return [] + if extractor is None: + return [] + try: + source = extractor(state) + except Exception: + logger.warning( + "Could not read the tool %s for guardrail '%s'; evaluating without files.", + source_name, + guardrail.name, + exc_info=True, + ) + return [] + return resolve_referenced_attachments(source, registry, guardrail) + + def _create_guardrail_node( guardrail: BaseGuardrail, scope: GuardrailScope, @@ -235,12 +280,13 @@ async def node( else: metadata["payload"]["output"] = payload - attachments = ( - await resolve_guardrail_attachments( - state.inner_state.job_attachments, guardrail - ) - if scope in _ATTACHMENT_SCOPES - else [] + attachments = await _resolve_attachments( + state, + guardrail, + scope, + execution_stage, + input_data_extractor, + output_data_extractor, ) result = await _evaluate_builtin_guardrail( @@ -346,6 +392,13 @@ def _payload_generator(state: AgentGuardrailsGraphState) -> str: ) +def _tool_call_field(tool_call: Any, field: str) -> Any: + """Read ``field`` from a tool call given as a dict or an object.""" + if isinstance(tool_call, dict): + return tool_call.get(field) + return getattr(tool_call, field, None) + + def create_tool_guardrail_node( guardrail: BaseGuardrail, execution_stage: ExecutionStage, @@ -368,6 +421,51 @@ def create_tool_guardrail_node( A tuple of (node_name, node_function) for the guardrail evaluation node. """ + def _current_call_args(state: AgentGuardrailsGraphState) -> dict[str, Any]: + """Arguments of the tool call this evaluation is about. + + One AI message can carry several calls to the same tool, executed one after + another with a ToolMessage appended after each. The call under evaluation is + therefore not "the first call named ``tool_name``" but, before execution, the + first one without a ToolMessage yet (the same selection the tool node makes), + and after execution, the one the last ToolMessage answers. Falls back to the + first matching call when the history has no ToolMessage bookkeeping. + """ + messages = state.messages + if not messages: + return {} + + ai_message = find_latest_ai_message(messages) + if ai_message is None or not ai_message.tool_calls: + return {} + + selected: Any = None + if execution_stage == ExecutionStage.PRE_EXECUTION: + try: + index = extract_current_tool_call_index(messages, tool_name) + except AgentRuntimeError: + index = None + if index is not None and index < len(ai_message.tool_calls): + selected = ai_message.tool_calls[index] + else: # POST_EXECUTION + last_message = messages[-1] + if isinstance(last_message, ToolMessage): + selected = next( + ( + call + for call in ai_message.tool_calls + if _tool_call_field(call, "id") == last_message.tool_call_id + ), + None, + ) + + if selected is None: + return _extract_tool_args_from_message(ai_message, tool_name) + + return _extract_tool_args_from_message( + AIMessage(content="", tool_calls=[selected]), tool_name + ) + def _payload_generator(state: AgentGuardrailsGraphState) -> str: """Extract tool call arguments for the specified tool name. @@ -381,24 +479,13 @@ def _payload_generator(state: AgentGuardrailsGraphState) -> str: return "" if execution_stage == ExecutionStage.PRE_EXECUTION: - last_message = state.messages[-1] - args_dict = _extract_tool_args_from_message(last_message, tool_name) - return json.dumps(args_dict) + return json.dumps(_current_call_args(state)) return get_message_content(state.messages[-1]) # 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 _current_call_args(state) def _output_data_extractor(state: AgentGuardrailsGraphState) -> dict[str, Any]: return _extract_tool_output_data(state) diff --git a/tests/agent/guardrails/test_attachment_refs.py b/tests/agent/guardrails/test_attachment_refs.py index 738c209d5..b8e67b976 100644 --- a/tests/agent/guardrails/test_attachment_refs.py +++ b/tests/agent/guardrails/test_attachment_refs.py @@ -1,10 +1,11 @@ -"""Tests for projecting the job-attachment registry into guardrail attachment refs.""" +"""Tests for projecting a run's attachments into guardrail attachment references.""" import uuid from typing import Any from unittest.mock import MagicMock import pytest +from pydantic import BaseModel from uipath.platform.attachments import Attachment from uipath.platform.guardrails import BuiltInValidatorGuardrail from uipath.platform.guardrails.guardrails import EnumParameterValue @@ -12,6 +13,7 @@ from uipath_langchain.agent.guardrails.attachment_refs import ( _MAX_ATTACHMENTS, resolve_guardrail_attachments, + resolve_referenced_attachments, ) _UUID = "7f2c1e44-0b3a-4a1e-9d55-2f9a1c3b8e10" @@ -178,3 +180,177 @@ async def test_resolves_when_the_scope_parameter_is_malformed(self, monkeypatch) result = await resolve_guardrail_attachments(_registry(), guardrail) assert [r.file_name for r in result] == ["a.csv"] + + +def _other_attachment(name: str = "b.pdf", mime: str = "application/pdf"): + attachment_id = str(uuid.uuid4()) + return attachment_id, Attachment( + id=uuid.UUID(attachment_id), full_name=name, mime_type=mime + ) + + +class TestResolveReferencedAttachments: + """Tool scope: only the attachments a tool call mentions are forwarded.""" + + async def test_resolves_the_attachment_the_tool_call_names(self): + """The model passes ``{"ID": ...}`` only; name and type come from the registry.""" + result = resolve_referenced_attachments( + {"attachment": {"ID": _UUID}, "question": "summarize"}, + _registry(), + _judge(), + ) + + assert [r.model_dump(by_alias=True) for r in result] == [ + {"id": _UUID, "fileName": "a.csv", "mimeType": "text/csv"} + ] + + async def test_forwards_nothing_when_the_call_names_no_file(self): + """The run holds a file, but this call does not touch it: forwarding it would + ship every file on every tool call.""" + result = resolve_referenced_attachments( + {"query": "cats", "limit": 5}, _registry(), _judge() + ) + + assert result == [] + + async def test_finds_mentions_nested_in_lists_and_dicts(self): + other_id, other = _other_attachment() + registry = _registry() | {other_id: other} + data = { + "files": [{"ID": _UUID}], + "options": {"inner": {"ref": {"ID": other_id}, "flag": True}}, + } + + result = resolve_referenced_attachments(data, registry, _judge()) + + assert sorted(r.id for r in result) == sorted([_UUID, other_id]) + + async def test_skips_an_id_the_run_never_held_even_with_inline_name_and_type( + self, + ): + """A tool result (or a prompt-injected argument) can name any id; only files the + run legitimately holds are sent to the backend, which would otherwise look the + id up in Orchestrator and judge a file the run never had.""" + other_id = str(uuid.uuid4()) + data = {"file": {"ID": other_id, "FullName": "out.csv", "MimeType": "text/csv"}} + + assert resolve_referenced_attachments(data, _registry(), _judge()) == [] + + async def test_uses_the_registry_name_and_type_not_the_inline_ones(self): + data = {"file": {"ID": _UUID, "FullName": "renamed.csv", "MimeType": "x/y"}} + + result = resolve_referenced_attachments(data, _registry(), _judge()) + + assert [(r.file_name, r.mime_type) for r in result] == [("a.csv", "text/csv")] + + async def test_skips_a_uuid_that_is_not_one_of_the_run_attachments(self): + """A UUID under ``ID`` that is not an attachment (a queue item, a job) must not + reach the backend, which would look it up in Orchestrator.""" + data = {"item": {"ID": str(uuid.uuid4())}} + + assert resolve_referenced_attachments(data, _registry(), _judge()) == [] + + @pytest.mark.parametrize("raw_id", ["queue-item-42", 12, True, None, ""]) + async def test_skips_a_non_uuid_id(self, raw_id): + data = {"item": {"ID": raw_id, "FullName": "a.csv", "MimeType": "text/csv"}} + + assert resolve_referenced_attachments(data, _registry(), _judge()) == [] + + async def test_deduplicates_the_same_attachment(self): + data = {"a": {"ID": _UUID}, "b": [{"ID": _UUID.upper()}]} + + result = resolve_referenced_attachments(data, _registry(), _judge()) + + assert [r.id for r in result] == [_UUID] + + async def test_caps_at_the_api_limit(self): + registry = dict(_other_attachment() for _ in range(_MAX_ATTACHMENTS + 2)) + data = {"files": [{"ID": attachment_id} for attachment_id in registry]} + + result = resolve_referenced_attachments(data, registry, _judge()) + + assert len(result) == _MAX_ATTACHMENTS + + @pytest.mark.parametrize("applies_to", ["Prompts", "prompts"]) + async def test_returns_empty_when_scoped_to_prompts(self, applies_to): + result = resolve_referenced_attachments( + {"attachment": {"ID": _UUID}}, _registry(), _scoped_judge(applies_to) + ) + + assert result == [] + + async def test_accepts_attachment_instances_and_models(self): + """Arguments may already carry expanded objects, not only wire dicts; they are + still looked up in the registry by id.""" + + class ToolArgs(BaseModel): + attachment: Attachment + + registry = _registry() + data = { + "direct": registry[_UUID], + "wrapped": ToolArgs(attachment=registry[_UUID]), + } + + result = resolve_referenced_attachments(data, registry, _judge()) + + assert [r.id for r in result] == [_UUID] + + async def test_bounds_the_scan_on_deep_and_large_payloads(self): + """A tool result can be arbitrarily large; the scan stops instead of stalling the + guardrail node, and never raises.""" + deep: dict[str, Any] = {"ID": _UUID} + for _ in range(100): + deep = {"child": deep} + wide = {"rows": [{"n": i} for i in range(20_000)], "file": {"ID": _UUID}} + + assert resolve_referenced_attachments(deep, _registry(), _judge()) == [] + wide_result = resolve_referenced_attachments(wide, _registry(), _judge()) + assert isinstance(wide_result, list) + + @pytest.mark.parametrize("data", [None, 42, "just text", object(), [1, "two"]]) + async def test_never_raises_on_non_structured_payloads(self, data): + assert resolve_referenced_attachments(data, _registry(), _judge()) == [] + + async def test_matches_the_registry_path_for_the_same_attachment(self): + """Agent/LLM scope and tool scope must send the backend identical references.""" + registry = _registry() + + via_registry = await resolve_guardrail_attachments(registry, _judge()) + via_mention = resolve_referenced_attachments( + {"attachment": {"ID": _UUID}}, registry, _judge() + ) + + assert via_registry == via_mention + + +class _UnscannablePayload(dict[str, Any]): + """A mapping whose lookups blow up, as a broken tool result might.""" + + def get(self, key, default=None): # noqa: D401 + raise RuntimeError("boom") + + +class _MentionWithUnreadableId(dict[str, Any]): + """Looks like a mention to the scanner but fails when the id is read.""" + + def get(self, key, default=None): + return _UUID if key == "ID" else default + + def __getitem__(self, key): + raise RuntimeError("boom") + + +class TestResolveReferencedAttachmentsErrorPaths: + async def test_returns_empty_when_the_payload_cannot_be_scanned(self): + """The guardrail node re-raises, so a broken payload must degrade to no files.""" + data = {"result": _UnscannablePayload(ID=_UUID)} + + assert resolve_referenced_attachments(data, _registry(), _judge()) == [] + + async def test_skips_a_mention_whose_id_cannot_be_read(self): + data = {"file": _MentionWithUnreadableId(), "other": {"ID": _UUID}} + + result = resolve_referenced_attachments(data, _registry(), _judge()) + + assert [r.id for r in result] == [_UUID] diff --git a/tests/agent/guardrails/test_guardrail_nodes.py b/tests/agent/guardrails/test_guardrail_nodes.py index fc030bb85..5ae0fafd3 100644 --- a/tests/agent/guardrails/test_guardrail_nodes.py +++ b/tests/agent/guardrails/test_guardrail_nodes.py @@ -14,6 +14,7 @@ from uipath.platform.guardrails import BuiltInValidatorGuardrail from uipath_langchain.agent.guardrails.guardrail_nodes import ( + _create_guardrail_node, create_agent_init_guardrail_node, create_agent_terminate_guardrail_node, create_llm_guardrail_node, @@ -1003,7 +1004,8 @@ async def test_tool_guardrail_payload_populated_post_execution(self, monkeypatch class TestGuardrailNodeAttachments: - """Agent- and LLM-scope nodes forward the run's job attachments to the judge.""" + """Built-in guardrail nodes forward attachment references to the judge: the run's + registry at Agent and LLM scope, the files one tool call mentions at Tool scope.""" _UUID = "7f2c1e44-0b3a-4a1e-9d55-2f9a1c3b8e10" @@ -1087,17 +1089,58 @@ async def test_llm_node_forwards_resolved_attachments(self, monkeypatch): assert fake.guardrails.last_attachments == [attachment] - @pytest.mark.asyncio - async def test_tool_scope_node_never_resolves_attachments(self, monkeypatch): - """Tool scope is excluded by product decision: a tool-scope judge would ship file - contents on every tool call.""" - from unittest.mock import AsyncMock + def _tool_pre_state(self, args): + return AgentGuardrailsGraphState( + messages=[ + AIMessage( + content="", + tool_calls=[{"name": "my_tool", "args": args, "id": "c1"}], + ) + ], + inner_state=InnerAgentGuardrailsGraphState( + job_attachments=self._state_with_attachment().inner_state.job_attachments + ), + ) + + def _tool_post_state(self, content): + return AgentGuardrailsGraphState( + messages=[ + AIMessage( + content="", + tool_calls=[{"name": "my_tool", "args": {}, "id": "c1"}], + ), + ToolMessage(content=content, tool_call_id="c1"), + ], + inner_state=InnerAgentGuardrailsGraphState( + job_attachments=self._state_with_attachment().inner_state.job_attachments + ), + ) + @pytest.mark.asyncio + async def test_tool_pre_node_judges_the_current_call_when_a_tool_is_called_twice( + self, monkeypatch + ): + """One AI message, two calls to the same tool: after the first call's ToolMessage + lands, the second evaluation must judge the second call's arguments and file, + not the first call's (the tool node selects the call the same way).""" fake = _patch_uipath(monkeypatch, reason="ok") - resolver = AsyncMock(return_value=[]) - monkeypatch.setattr( - "uipath_langchain.agent.guardrails.guardrail_nodes.resolve_guardrail_attachments", - resolver, + other = "00000000-0000-4000-8000-000000000001" + first_args = {"attachment": {"ID": other}, "question": "first"} + second_args = {"attachment": {"ID": self._UUID}, "question": "second"} + state = AgentGuardrailsGraphState( + messages=[ + AIMessage( + content="", + tool_calls=[ + {"name": "my_tool", "args": first_args, "id": "c1"}, + {"name": "my_tool", "args": second_args, "id": "c2"}, + ], + ), + ToolMessage(content="first done", tool_call_id="c1"), + ], + inner_state=InnerAgentGuardrailsGraphState( + job_attachments=self._state_with_attachment().inner_state.job_attachments + ), ) _, node = create_tool_guardrail_node( @@ -1107,22 +1150,224 @@ async def test_tool_scope_node_never_resolves_attachments(self, monkeypatch): failure_node="nope", tool_name="my_tool", ) + cmd = await node(state) + + assert cmd.goto == "ok" + assert json.loads(fake.guardrails.last_text) == second_args + assert [a.id for a in fake.guardrails.last_attachments] == [self._UUID] + + @pytest.mark.asyncio + async def test_tool_post_node_judges_the_answered_call_when_a_tool_is_called_twice( + self, monkeypatch + ): + """After the second of two calls returns, the post evaluation reads that + call's result (the last ToolMessage), which here names the file.""" + fake = _patch_uipath(monkeypatch, reason="ok") + second_result = json.dumps( + { + "file": { + "ID": self._UUID, + "FullName": "Tickets.csv", + "MimeType": "text/csv", + } + } + ) state = AgentGuardrailsGraphState( messages=[ AIMessage( content="", - tool_calls=[{"name": "my_tool", "args": {"q": 1}, "id": "c1"}], - ) + tool_calls=[ + {"name": "my_tool", "args": {"n": 1}, "id": "c1"}, + {"name": "my_tool", "args": {"n": 2}, "id": "c2"}, + ], + ), + ToolMessage(content="first done", tool_call_id="c1"), + ToolMessage(content=second_result, tool_call_id="c2"), ], inner_state=InnerAgentGuardrailsGraphState( job_attachments=self._state_with_attachment().inner_state.job_attachments ), ) + + _, node = create_tool_guardrail_node( + guardrail=self._judge_guardrail(), + execution_stage=ExecutionStage.POST_EXECUTION, + success_node="ok", + failure_node="nope", + tool_name="my_tool", + ) await node(state) - resolver.assert_not_awaited() + assert fake.guardrails.last_text == second_result + assert [a.id for a in fake.guardrails.last_attachments] == [self._UUID] + + @pytest.mark.asyncio + async def test_tool_pre_node_forwards_attachments_referenced_in_tool_args( + self, monkeypatch + ): + """Before the tool runs, the judge reads the file the call names. The model + passes ``{"ID": ...}`` only, so name and type come from the run's registry.""" + fake = _patch_uipath(monkeypatch, reason="ok") + args = {"attachment": {"ID": self._UUID}, "question": "summarize"} + + _, node = create_tool_guardrail_node( + guardrail=self._judge_guardrail(), + execution_stage=ExecutionStage.PRE_EXECUTION, + success_node="ok", + failure_node="nope", + tool_name="my_tool", + ) + cmd = await node(self._tool_pre_state(args)) + + assert cmd.goto == "ok" + assert json.loads(fake.guardrails.last_text) == args + assert [ + a.model_dump(by_alias=True) for a in fake.guardrails.last_attachments + ] == [{"id": self._UUID, "fileName": "Tickets.csv", "mimeType": "text/csv"}] + + @pytest.mark.asyncio + async def test_tool_pre_node_ignores_registry_when_args_reference_nothing( + self, monkeypatch + ): + """A tool call without a file forwards nothing even though the run holds one; + otherwise every tool call would ship every file to the backend.""" + fake = _patch_uipath(monkeypatch, reason="ok") + + _, node = create_tool_guardrail_node( + guardrail=self._judge_guardrail(), + execution_stage=ExecutionStage.PRE_EXECUTION, + success_node="ok", + failure_node="nope", + tool_name="my_tool", + ) + await node(self._tool_pre_state({"q": 1})) + + assert fake.guardrails.last_attachments == [] + + @pytest.mark.asyncio + async def test_tool_post_node_forwards_attachment_returned_by_the_tool( + self, monkeypatch + ): + """After the tool runs, the judge reads the file the result names.""" + fake = _patch_uipath(monkeypatch, reason="ok") + content = json.dumps( + { + "file": { + "ID": self._UUID, + "FullName": "Tickets.csv", + "MimeType": "text/csv", + } + } + ) + + _, node = create_tool_guardrail_node( + guardrail=self._judge_guardrail(), + execution_stage=ExecutionStage.POST_EXECUTION, + success_node="ok", + failure_node="nope", + tool_name="my_tool", + ) + await node(self._tool_post_state(content)) + + assert fake.guardrails.last_text == content + assert [a.id for a in fake.guardrails.last_attachments] == [self._UUID] + + @pytest.mark.asyncio + async def test_tool_post_node_with_plain_text_result_forwards_nothing( + self, monkeypatch + ): + fake = _patch_uipath(monkeypatch, reason="ok") + + _, node = create_tool_guardrail_node( + guardrail=self._judge_guardrail(), + execution_stage=ExecutionStage.POST_EXECUTION, + success_node="ok", + failure_node="nope", + tool_name="my_tool", + ) + await node(self._tool_post_state("tool output")) + assert fake.guardrails.last_attachments == [] + @pytest.mark.asyncio + async def test_tool_node_without_extractors_forwards_nothing(self, monkeypatch): + """A tool-scope built-in node built without the argument/result extractors has + no source to scan and must not fall back to the whole registry.""" + from uipath.platform.guardrails import GuardrailScope + + fake = _patch_uipath(monkeypatch, reason="ok") + + _, node = _create_guardrail_node( + self._judge_guardrail(), + GuardrailScope.TOOL, + ExecutionStage.PRE_EXECUTION, + lambda state: "payload", + "ok", + "nope", + ) + cmd = await node(self._tool_pre_state({"attachment": {"ID": self._UUID}})) + + assert cmd.goto == "ok" + assert fake.guardrails.last_attachments == [] + + @pytest.mark.asyncio + async def test_tool_post_node_evaluates_without_files_when_the_result_cannot_be_read( + self, monkeypatch + ): + """A file must never fail the run: if the result extractor blows up, the judge + still sees the text payload, just without attachments.""" + fake = _patch_uipath(monkeypatch, reason="ok") + + def broken(_state): + raise RuntimeError("boom") + + monkeypatch.setattr( + "uipath_langchain.agent.guardrails.guardrail_nodes._extract_tool_output_data", + broken, + ) + + _, node = create_tool_guardrail_node( + guardrail=self._judge_guardrail(), + execution_stage=ExecutionStage.POST_EXECUTION, + success_node="ok", + failure_node="nope", + tool_name="my_tool", + ) + cmd = await node(self._tool_post_state('{"ID": "not-json-but-mentions-ID"')) + + assert cmd.goto == "ok" + assert fake.guardrails.last_text == '{"ID": "not-json-but-mentions-ID"' + assert fake.guardrails.last_attachments == [] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "factory", [create_agent_init_guardrail_node, create_llm_guardrail_node] + ) + async def test_agent_and_llm_pre_nodes_forward_files_scoped_attachments( + self, monkeypatch, factory + ): + """Pre-execution at Agent and LLM scope with ``appliesTo = Files`` still hands + the run's files to the judge (real resolver, nothing patched).""" + from uipath.platform.guardrails.guardrails import EnumParameterValue + + fake = _patch_uipath(monkeypatch, reason="ok") + guardrail = self._judge_guardrail() + guardrail.validator_parameters = [ + EnumParameterValue.model_validate( + {"$parameterType": "enum", "id": "appliesTo", "value": "Files"} + ) + ] + + _, node = factory( + guardrail=guardrail, + execution_stage=ExecutionStage.PRE_EXECUTION, + success_node="ok", + failure_node="nope", + ) + await node(self._state_with_attachment()) + + assert [a.id for a in fake.guardrails.last_attachments] == [self._UUID] + @pytest.mark.asyncio async def test_attachment_rejection_falls_back_to_text_only(self, monkeypatch): """A 400 on a request that carried attachments must not kill the run: the backend diff --git a/tests/agent/guardrails/test_tool_guardrails_subgraph_run.py b/tests/agent/guardrails/test_tool_guardrails_subgraph_run.py new file mode 100644 index 000000000..06336d455 --- /dev/null +++ b/tests/agent/guardrails/test_tool_guardrails_subgraph_run.py @@ -0,0 +1,117 @@ +"""A compiled tool-guardrail subgraph, run end to end with a judge that reads files.""" + +import json +import uuid +from typing import Any + +import pytest +from langchain_core.messages import AIMessage, ToolMessage +from langchain_core.tools import StructuredTool +from pydantic import BaseModel +from uipath.core.guardrails import GuardrailValidationResultType +from uipath.platform.attachments import Attachment +from uipath.platform.guardrails import BuiltInValidatorGuardrail + +from tests.agent.guardrails.test_guardrail_nodes import _patch_uipath +from uipath_langchain.agent.exceptions import AgentRuntimeError, AgentRuntimeErrorCode +from uipath_langchain.agent.guardrails.actions import BlockAction +from uipath_langchain.agent.react.guardrails.guardrails_subgraph import ( + create_tool_guardrails_subgraph, +) +from uipath_langchain.agent.tools.tool_node import UiPathToolNode + +_UUID = "7f2c1e44-0b3a-4a1e-9d55-2f9a1c3b8e10" +_TOOL = "analyze_files" + + +class _Args(BaseModel): + attachment: dict[str, Any] + question: str + + +def _judge() -> BuiltInValidatorGuardrail: + return BuiltInValidatorGuardrail.model_validate( + { + "$guardrailType": "builtInValidator", + "id": "judge-files", + "name": "Judge files", + "description": "Blocks poisoned files before a tool reads them.", + "enabledForEvals": True, + "selector": {"scopes": ["Tool"], "matchNames": [_TOOL]}, + "validatorType": "llm_as_judge", + "validatorParameters": [], + } + ) + + +def _build(monkeypatch, *, result: GuardrailValidationResultType): + calls: list[dict[str, Any]] = [] + + async def analyze(attachment: dict[str, Any], question: str) -> str: + calls.append({"attachment": attachment, "question": question}) + return "The file lists 12 tickets." + + tool = StructuredTool.from_function( + coroutine=analyze, name=_TOOL, description="Reads a file.", args_schema=_Args + ) + graph = create_tool_guardrails_subgraph( + tool_node=(_TOOL, UiPathToolNode(tool)), + guardrails=[(_judge(), BlockAction("unsafe file"))], + ) + fake = _patch_uipath(monkeypatch, result=result, reason="judged") + return graph, fake, calls + + +def _input() -> dict[str, Any]: + args = {"attachment": {"ID": _UUID}, "question": "summarize"} + return { + "messages": [ + AIMessage( + content="", tool_calls=[{"name": _TOOL, "args": args, "id": "c1"}] + ) + ], + "inner_state": { + "job_attachments": { + _UUID: Attachment( + id=uuid.UUID(_UUID), full_name="Tickets.csv", mime_type="text/csv" + ) + } + }, + } + + +@pytest.mark.asyncio +async def test_pre_execution_judge_blocks_before_the_tool_runs(monkeypatch): + """The judge saw the referenced file and the tool implementation never ran.""" + graph, fake, calls = _build( + monkeypatch, result=GuardrailValidationResultType.VALIDATION_FAILED + ) + + payload = _input() + + with pytest.raises(AgentRuntimeError) as raised: + await graph.ainvoke(payload) + + assert raised.value.error_info.code == AgentRuntimeError.full_code( + AgentRuntimeErrorCode.TERMINATION_GUARDRAIL_VIOLATION + ) + assert calls == [] + assert fake.guardrails.call_count == 1 + assert [a.id for a in fake.guardrails.last_attachments] == [_UUID] + assert json.loads(fake.guardrails.last_text)["attachment"] == {"ID": _UUID} + + +@pytest.mark.asyncio +async def test_passing_judge_lets_the_tool_run_and_judges_its_result(monkeypatch): + graph, fake, calls = _build( + monkeypatch, result=GuardrailValidationResultType.PASSED + ) + + final = await graph.ainvoke(_input()) + + assert [call["question"] for call in calls] == ["summarize"] + assert fake.guardrails.call_count == 2 + assert isinstance(final["messages"][-1], ToolMessage) + # The post-execution judge read the plain-text result and no file with it. + assert fake.guardrails.last_text == "The file lists 12 tickets." + assert fake.guardrails.last_attachments == [] diff --git a/uv.lock b/uv.lock index b6cf4ec82..a1a2cb0fa 100644 --- a/uv.lock +++ b/uv.lock @@ -4828,7 +4828,7 @@ wheels = [ [[package]] name = "uipath-langchain" -version = "0.18.12" +version = "0.18.13" source = { editable = "." } dependencies = [ { name = "a2a-sdk" },