feat(guardrails): let tool-scope judge guardrails read the files a tool call mentions

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"

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
 
 

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 @@
     )
 
 
+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 @@
                 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,7 +392,14 @@
     )
 
 
+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,
     success_node: str,
@@ -368,6 +421,51 @@
         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 @@
             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)