feat(guardrails): give post-execution guardrails only the files the output references [AL-611]

pyproject.toml

@@ -1,13 +1,13 @@
 [project]
 name = "uipath-langchain"
-version = "0.18.33"
+version = "0.18.34"
 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"
 dependencies = [
     "uipath>=2.14.35, <2.15.0",
     "uipath-core>=0.5.33, <0.6.0",
-    "uipath-platform>=0.2.34, <0.3.0",
+    "uipath-platform>=0.2.36, <0.3.0",
     "uipath-runtime>=0.13.0, <0.14.0",
     "uipath-llm-client>=1.21.0, <1.22.0",
     "langgraph>=1.2.11, <2.0.0",

src/uipath_langchain/agent/attachments/output_files.py

@@ -201,6 +201,31 @@ async def resolve_output_attachments(
     return resolved, []
 
 
+def verified_output_attachments(
+    fields: list[OutputFileField], output: dict[str, Any]
+) -> dict[str, Attachment]:
+    """The output's file references as registry entries, keyed by attachment id.
+
+    For an output :func:`check_output_files` accepted, whose references were rebuilt
+    from the attachments Orchestrator holds. Registering them lets the agent-output
+    guardrails resolve a file whatever tool produced it, including tools that do not
+    register their attachments. Without a job nothing was looked up, so nothing is
+    returned.
+    """
+    if not UiPathConfig.job_key:
+        return {}
+    verified: dict[str, Attachment] = {}
+    for field in fields:
+        for value in _filled_values(output, field.path):
+            try:
+                attachment = Attachment.model_validate(value, from_attributes=True)
+            except ValidationError:
+                continue
+            if attachment.id is not None:
+                verified[str(attachment.id)] = attachment
+    return verified
+
+
 class OutputFilesCheck(NamedTuple):
     """The verdict on an output's file fields."""
 

src/uipath_langchain/agent/guardrails/attachment_refs.py

@@ -2,15 +2,23 @@
 
 Two sources feed the same reference shape:
 
-* 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 its ``appliesTo`` scope is text only.
+* Pre-execution guardrails at Agent and LLM scope judge the conversation so far, so they
+  read the whole job-attachment registry (:func:`resolve_guardrail_attachments`).
+* Every other guardrail judges one payload, so it reads only the attachments that payload
+  mentions (:func:`resolve_referenced_attachments`): the ``{"ID": ...}`` objects in a tool
+  call's arguments or result, in the tool-call arguments an LLM produced, or in the
+  agent's result. A payload that names no file forwards none, even when the run holds
+  files elsewhere; otherwise every check would ship every file to the backend, and an
+  input file would be judged again as if the agent had produced it.
+
+Both return one of three answers, which the validate API reads differently:
+
+* ``None``: say nothing about files. The guardrail is scoped to text only, or the
+  payload could not be scanned; the backend evaluates the text as before.
+* ``[]``: files apply and there are none to inspect. A files-only guardrail passes
+  without judging the text.
+* a list: the files to inspect.
+
 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.
@@ -38,7 +46,7 @@
 _FILE_SCOPES: frozenset[str] = frozenset({"files", "both"})
 #: 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.
+#: Bounds for scanning a payload, which can be arbitrarily large or deep.
 _MAX_SCAN_DEPTH = 32
 _MAX_SCAN_NODES = 10_000
 
@@ -60,76 +68,83 @@ def _scope_includes_files(guardrail: BuiltInValidatorGuardrail) -> bool:
     return True
 
 
+def _is_text_only(guardrail: BuiltInValidatorGuardrail) -> bool:
+    if _scope_includes_files(guardrail):
+        return False
+    logger.debug(
+        "Guardrail '%s' is scoped to text only; skipping attachment resolution.",
+        guardrail.name,
+    )
+    return True
+
+
 async def resolve_guardrail_attachments(
     job_attachments: dict[str, Attachment],
     guardrail: BuiltInValidatorGuardrail,
-) -> list[GuardrailAttachment]:
+) -> list[GuardrailAttachment] | None:
     """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 text only or the run has no attachments. Never
-    raises.
+    For pre-execution guardrails at Agent and LLM scope, which evaluate the conversation
+    as a whole. ``None`` when the guardrail is scoped to text only, empty when the run
+    has no attachments. Never raises.
     """
+    if _is_text_only(guardrail):
+        return None
     if not job_attachments:
         return []
-    if not _scope_includes_files(guardrail):
-        logger.debug(
-            "Guardrail '%s' is scoped to text only; skipping attachment resolution.",
-            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 text only. Never raises.
+) -> list[GuardrailAttachment] | None:
+    """Return up to five references for the attachments ``data`` mentions, in order.
+
+    ``data`` is the payload one guardrail judges: a tool call's arguments or parsed
+    result, an LLM's tool-call arguments, or the agent's result. 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 or its output verified), so a mention the run never
+    held, or a non-attachment resource id, is skipped rather than sent to the backend
+    for lookup. ``None`` when the guardrail is scoped to text only or ``data`` cannot be
+    scanned; empty when nothing is mentioned. Never raises.
     """
+    if _is_text_only(guardrail):
+        return None
     if data is None:
         return []
-    if not _scope_includes_files(guardrail):
-        logger.debug(
-            "Guardrail '%s' is scoped to text only; 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 "
+            "Could not scan the payload for attachments; guardrail '%s' evaluates "
             "without files.",
             guardrail.name,
             exc_info=True,
         )
-        return []
+        return None
     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."""
+    """Yield every attachment-shaped value in ``data`` in document order, within bounds.
+
+    Children are pushed reversed so the stack pops them first to last: when more files
+    are mentioned than the API accepts, the ones written first are the ones sent.
+    """
     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.",
+                "Stopped scanning the payload for attachments after %d values.",
                 _MAX_SCAN_NODES,
             )
             return
@@ -140,12 +155,13 @@ def _iter_attachment_mentions(data: Any) -> Iterator[Any]:
         elif isinstance(value, BaseModel):
             stack.append((value.model_dump(by_alias=True), depth + 1))
         elif isinstance(value, Mapping):
+            # A UUID ``ID`` may name a record other than a file, so its children are
+            # still scanned; a mention that is not in the registry is dropped later.
             if _is_mention(value):
                 yield value
-            else:
-                stack.extend((item, depth + 1) for item in value.values())
+            stack.extend((item, depth + 1) for item in reversed(list(value.values())))
         elif isinstance(value, (list, tuple, set, frozenset)):
-            stack.extend((item, depth + 1) for item in value)
+            stack.extend((item, depth + 1) for item in reversed(list(value)))
 
 
 def _is_mention(value: Mapping[Any, Any]) -> bool:
@@ -173,7 +189,7 @@ def _resolve_mention(mention: Any, registry: dict[str, Attachment]) -> Any | Non
         )
     except Exception:
         logger.debug(
-            "Skipping a malformed attachment reference in the tool payload.",
+            "Skipping a malformed attachment reference in the payload.",
             exc_info=True,
         )
     return None

src/uipath_langchain/agent/guardrails/guardrail_nodes.py

@@ -205,45 +205,35 @@ def _create_validation_command(
     )
 
 
+AttachmentSource = Callable[[AgentGuardrailsGraphState], Any]
+"""Reads, from the state, the payload whose file mentions a guardrail inspects."""
+
+
 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_source: AttachmentSource | None,
+) -> list[GuardrailAttachment] | None:
     """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.
+    Without a source the guardrail judges the conversation so far, so it reads the run's
+    whole attachment registry. With one it judges a single payload, so it reads only the
+    files that payload mentions. ``None`` lets the backend fall back to the text; a
+    source that raises gets that too. Never raises.
     """
     registry = state.inner_state.job_attachments
-    if scope != GuardrailScope.TOOL:
+    if attachment_source is None:
         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)
+        source = attachment_source(state)
     except Exception:
         logger.warning(
-            "Could not read the tool %s for guardrail '%s'; evaluating without files.",
-            source_name,
+            "Could not read the payload of guardrail '%s' for files; evaluating "
+            "without them.",
             guardrail.name,
             exc_info=True,
         )
-        return []
+        return None
     return resolve_referenced_attachments(source, registry, guardrail)
 
 
@@ -261,12 +251,16 @@ def _create_guardrail_node(
     tool_name: str | None = None,
     tool_type: str | None = None,
     termination_mode: GuardrailTerminationMode | None = None,
+    attachment_source: AttachmentSource | None = None,
 ) -> tuple[str, Callable[[AgentGuardrailsGraphState], Any]]:
     """Private factory for guardrail evaluation nodes.
 
     Returns a node with observability metadata attached as __metadata__ attribute:
     - goto success_node on validation pass
     - goto failure_node on validation fail
+
+    ``attachment_source`` picks the files a built-in guardrail inspects; see
+    :func:`_resolve_attachments`.
     """
     raw_node_name = f"{scope.name}_{execution_stage.name}_{guardrail.name}"
     node_name = re.sub(r"\W+", "_", raw_node_name.lower()).strip("_")
@@ -317,12 +311,7 @@ async def node(
                     metadata["payload"]["output"] = payload
 
                 attachments = await _resolve_attachments(
-                    state,
-                    guardrail,
-                    scope,
-                    execution_stage,
-                    input_data_extractor,
-                    output_data_extractor,
+                    state, guardrail, attachment_source
                 )
 
                 result = await _evaluate_builtin_guardrail(
@@ -382,6 +371,11 @@ def _payload_generator(state: AgentGuardrailsGraphState) -> str:
             case ExecutionStage.POST_EXECUTION:
                 return json.dumps(_extract_tools_args_from_message(state.messages[-1]))
 
+    def _tool_call_args(state: AgentGuardrailsGraphState) -> Any:
+        if not state.messages:
+            return None
+        return _extract_tools_args_from_message(state.messages[-1])
+
     return _create_guardrail_node(
         guardrail,
         GuardrailScope.LLM,
@@ -390,6 +384,11 @@ def _payload_generator(state: AgentGuardrailsGraphState) -> str:
         success_node,
         failure_node,
         termination_mode=termination_mode,
+        attachment_source=(
+            _tool_call_args
+            if execution_stage == ExecutionStage.POST_EXECUTION
+            else None
+        ),
     )
 
 
@@ -428,6 +427,9 @@ def create_agent_terminate_guardrail_node(
     def _payload_generator(state: AgentGuardrailsGraphState) -> str:
         return str(state.inner_state.agent_result)
 
+    def _agent_result(state: AgentGuardrailsGraphState) -> Any:
+        return state.inner_state.agent_result
+
     return _create_guardrail_node(
         guardrail,
         GuardrailScope.AGENT,
@@ -436,6 +438,7 @@ def _payload_generator(state: AgentGuardrailsGraphState) -> str:
         success_node,
         failure_node,
         termination_mode=termination_mode,
+        attachment_source=_agent_result,
     )
 
 
@@ -539,6 +542,13 @@ def _input_data_extractor(state: AgentGuardrailsGraphState) -> dict[str, Any]:
     def _output_data_extractor(state: AgentGuardrailsGraphState) -> dict[str, Any]:
         return _extract_tool_output_data(state)
 
+    def _tool_result(state: AgentGuardrailsGraphState) -> Any:
+        # 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 None
+        return _extract_tool_output_data(state)
+
     return _create_guardrail_node(
         guardrail,
         GuardrailScope.TOOL,
@@ -551,4 +561,9 @@ def _output_data_extractor(state: AgentGuardrailsGraphState) -> dict[str, Any]:
         tool_name,
         tool_type,
         termination_mode=termination_mode,
+        attachment_source=(
+            _current_call_args
+            if execution_stage == ExecutionStage.PRE_EXECUTION
+            else _tool_result
+        ),
     )