diff --git a/src/uipath_langchain/agent/advanced/agent.py b/src/uipath_langchain/agent/advanced/agent.py index ea6e1d564..65b101b87 100644 --- a/src/uipath_langchain/agent/advanced/agent.py +++ b/src/uipath_langchain/agent/advanced/agent.py @@ -23,9 +23,28 @@ from langgraph.graph.state import CompiledStateGraph, StateGraph from pydantic import BaseModel, ConfigDict, Field, create_model from uipath.core.chat import UiPathConversationMessageData +from uipath.runtime.errors import UiPathErrorCategory from uipath_langchain._utils import get_unique_model_field_name from uipath_langchain.agent.attachments.job_attachments import get_job_attachment_paths +from uipath_langchain.agent.attachments.output_files import ( + DEFAULT_MAX_OUTPUT_FILE_RETRIES, + diagnose_output_files, + get_output_file_fields, +) +from uipath_langchain.agent.exceptions import ( + AgentRuntimeError, + AgentRuntimeErrorCode, +) +from uipath_langchain.agent.react.conversational_output_node import ( + create_conversational_output_extractor, +) +from uipath_langchain.agent.react.utils import ( + has_custom_conversational_output_fields, +) +from uipath_langchain.agent.tools.internal_tools.output_file_tool import ( + OUTPUT_FILE_TOOL_NAME, +) from uipath_langchain.runtime.messages import UiPathChatMessagesMapper from .types import ( @@ -175,6 +194,11 @@ def create_advanced_agent_graph( ``FilesystemBackend`` also enables workspace memory: deepagents' ``MemoryMiddleware`` reads ``/memory/MEMORY.md`` from the backend each turn. Memory stays disabled for non-filesystem backends, which carry no workspace. + + When the output schema declares a job-attachment field, a verification node + gates the typed output: an unfilled required file field, or a reference to an + attachment that is not linked to this job, sends the agent back for another + turn with a corrective message instead of emitting an output it cannot honor. """ memory_sources = ( [MEMORY_INDEX_VIRTUAL_PATH] if isinstance(backend, FilesystemBackend) else [] @@ -182,6 +206,13 @@ def create_advanced_agent_graph( runtime_prompt = _resolve_runtime_system_prompt( system_prompt, AdvancedAgentGraphState, input_schema ) + # Gated on the tool being present, for the same reason as the standard graph: + # an agent with no way to create a file must not be faulted for lacking one. + output_file_fields = ( + get_output_file_fields(output_schema) + if any(tool.name == OUTPUT_FILE_TOOL_NAME for tool in tools) + else [] + ) inner_graph = create_advanced_agent( model=model, @@ -194,16 +225,25 @@ def create_advanced_agent_graph( skills=skills, ) + output_file_retries_key = get_unique_model_field_name( + "uipath__output_file_retries", AdvancedAgentGraphState, input_schema + ) + output_file_problem_key = get_unique_model_field_name( + "uipath__output_file_problem", AdvancedAgentGraphState, input_schema + ) + state_fields: dict[str, Any] = dict(runtime_prompt.state_fields) + if output_file_fields: + state_fields[output_file_retries_key] = (int, 0) + state_fields[output_file_problem_key] = (str | None, None) + wrapper_state = create_state_with_input(input_schema) - if runtime_prompt.state_fields: + if state_fields: wrapper_state = create_model( "RuntimeAdvancedAgentGraphState", __base__=wrapper_state, - **runtime_prompt.state_fields, + **state_fields, ) - internal_fields = set(AdvancedAgentGraphState.model_fields) | set( - runtime_prompt.state_fields - ) + internal_fields = set(AdvancedAgentGraphState.model_fields) | set(state_fields) attachment_paths = ( get_job_attachment_paths(input_schema) if input_schema is not None else [] ) @@ -231,6 +271,41 @@ def transform_output(state: BaseModel) -> dict[str, Any]: structured = getattr(state, "structured_response", {}) return output_schema.model_validate(structured).model_dump() + async def verify_output_files(state: BaseModel) -> dict[str, Any]: + structured = getattr(state, "structured_response", {}) or {} + problem = await diagnose_output_files(output_file_fields, structured) + if problem is None: + return {output_file_problem_key: None} + + retries = getattr(state, output_file_retries_key, 0) or 0 + if retries >= DEFAULT_MAX_OUTPUT_FILE_RETRIES: + raise AgentRuntimeError( + code=AgentRuntimeErrorCode.OUTPUT_VALIDATION_ERROR, + title="Agent did not produce the required output file", + detail=( + f"{problem} The agent was given " + f"{DEFAULT_MAX_OUTPUT_FILE_RETRIES} chance(s) to correct this " + "and did not. Verify the agent's prompt asks for the file, and " + "that the output schema's file fields are the ones you intend." + ), + category=UiPathErrorCategory.USER, + ) + + # The structured-output tool call is already answered by this point, so the + # correction goes in as a new user turn rather than a tool result. + return { + "messages": [HumanMessage(content=problem)], + output_file_retries_key: retries + 1, + output_file_problem_key: problem, + } + + def route_after_verification(state: BaseModel) -> str: + return ( + "advanced_agent" + if getattr(state, output_file_problem_key, None) + else "transform_output" + ) + wrapper: StateGraph[Any, Any, Any, Any] = StateGraph( wrapper_state, input_schema=input_schema, output_schema=output_schema ) @@ -239,7 +314,16 @@ def transform_output(state: BaseModel) -> dict[str, Any]: wrapper.add_node("transform_output", transform_output) wrapper.add_edge(START, "transform_input") wrapper.add_edge("transform_input", "advanced_agent") - wrapper.add_edge("advanced_agent", "transform_output") + if output_file_fields: + wrapper.add_node("verify_output_files", verify_output_files) + wrapper.add_edge("advanced_agent", "verify_output_files") + wrapper.add_conditional_edges( + "verify_output_files", + route_after_verification, + ["transform_output", "advanced_agent"], + ) + else: + wrapper.add_edge("advanced_agent", "transform_output") wrapper.add_edge("transform_output", END) return wrapper @@ -252,6 +336,7 @@ def create_conversational_advanced_agent_graph( backend: BackendProtocol | BackendFactory | None, skills: Sequence[str] | None = None, input_schema: type[BaseModel] | None = None, + output_schema: type[BaseModel] | None = None, ) -> StateGraph[Any, Any, Any, Any]: """Wrap the advanced agent in a parent graph that speaks the conversational contract. @@ -260,6 +345,11 @@ def create_conversational_advanced_agent_graph( messages as ``uipath__agent_response_messages``. Callable system prompts are resolved once from the exchange input and used by the deep agent for that invocation. + + When ``output_schema`` declares fields beyond the response messages, they are + filled the same way the standard conversational agent fills them: a focused + extraction call over the exchange's messages, after the loop has finished. + The loop itself produces messages, so nothing in it can produce those fields. """ memory_sources = ( [MEMORY_INDEX_VIRTUAL_PATH] if isinstance(backend, FilesystemBackend) else [] @@ -288,6 +378,13 @@ class ConversationalAdvancedAgentOutput(BaseModel): default_factory=list ) + with_output_extraction = has_custom_conversational_output_fields(output_schema) + graph_output: type[BaseModel] = ( + output_schema + if with_output_extraction and output_schema is not None + else ConversationalAdvancedAgentOutput + ) + graph_input: type[BaseModel] = _ConversationalAdvancedAgentGraphInput wrapper_input: type[BaseModel] = _ConversationalAdvancedAgentGraphInput if input_schema: @@ -311,10 +408,17 @@ class ConversationalAdvancedAgentOutput(BaseModel): input_schema if "messages" in input_schema.model_fields else wrapper_input ) + conversational_output_key = get_unique_model_field_name( + "uipath__conversational_output", + _ConversationalAdvancedAgentGraphInput, + input_schema, + ) state_fields: dict[str, Any] = { initial_message_count_key: (int | None, None), **runtime_prompt.state_fields, } + if with_output_extraction: + state_fields[conversational_output_key] = (dict[str, Any] | None, None) wrapper_state = cast( type[BaseModel], create_model( @@ -347,10 +451,13 @@ def capture_exchange_start(state: BaseModel) -> dict[str, Any]: update.update(runtime_prompt.resolve(declared_input(state))) return update - def transform_output(state: BaseModel) -> dict[str, Any]: + def _new_messages(state: BaseModel) -> list[Any]: initial_count = getattr(state, initial_message_count_key) or 0 messages = cast(ConversationalAdvancedAgentGraphState, state).messages - new_messages = messages[initial_count:] + return list(messages[initial_count:]) + + def transform_output(state: BaseModel) -> dict[str, Any]: + new_messages = _new_messages(state) converted = ( UiPathChatMessagesMapper.map_langchain_messages_to_uipath_message_data_list( messages=new_messages, include_tool_results=False @@ -358,19 +465,53 @@ def transform_output(state: BaseModel) -> dict[str, Any]: if new_messages else [] ) - return {"uipath__agent_response_messages": converted} + if not with_output_extraction or output_schema is None: + return {"uipath__agent_response_messages": converted} + + custom_fields = getattr(state, conversational_output_key, None) or {} + output = { + **custom_fields, + "uipath__agent_response_messages": [ + message.model_dump(by_alias=True) for message in converted + ], + } + return output_schema.model_validate(output).model_dump( + by_alias=True, exclude_none=True + ) + + # Built once: binding tools is stateless, and the node runs every exchange. + extract_output = ( + create_conversational_output_extractor(model, output_schema) + if with_output_extraction and output_schema is not None + else None + ) + + async def generate_conversational_output(state: BaseModel) -> dict[str, Any]: + assert extract_output is not None # guarded by with_output_extraction + # The whole transcript, as the standard path passes: a declared field's + # answer often lives in the user's message or an earlier exchange, not in + # what the agent just produced. + messages = cast(ConversationalAdvancedAgentGraphState, state).messages + return {conversational_output_key: await extract_output(messages)} wrapper: StateGraph[Any, Any, Any, Any] = StateGraph( wrapper_state, input_schema=graph_input, - output_schema=ConversationalAdvancedAgentOutput, + output_schema=graph_output, ) wrapper.add_node("capture_exchange_start", capture_exchange_start) wrapper.add_node("advanced_agent", inner_graph) wrapper.add_node("transform_output", transform_output) wrapper.add_edge(START, "capture_exchange_start") wrapper.add_edge("capture_exchange_start", "advanced_agent") - wrapper.add_edge("advanced_agent", "transform_output") + if with_output_extraction: + wrapper.add_node( + "generate_conversational_output", generate_conversational_output + ) + wrapper.add_edge("advanced_agent", "generate_conversational_output") + wrapper.add_edge("generate_conversational_output", "transform_output") + else: + wrapper.add_edge("advanced_agent", "transform_output") wrapper.add_edge("transform_output", END) return wrapper diff --git a/src/uipath_langchain/agent/attachments/constants.py b/src/uipath_langchain/agent/attachments/constants.py new file mode 100644 index 000000000..4a376c3d3 --- /dev/null +++ b/src/uipath_langchain/agent/attachments/constants.py @@ -0,0 +1,8 @@ +"""Names shared between an output-file's schema handling and its tool. + +A leaf module so neither side has to import the other: the tool module builds +the tool, and the attachments module builds the prompt and the corrective +messages that name it. +""" + +OUTPUT_FILE_TOOL_NAME = "create_output_file" diff --git a/src/uipath_langchain/agent/attachments/output_files.py b/src/uipath_langchain/agent/attachments/output_files.py new file mode 100644 index 000000000..a307703c5 --- /dev/null +++ b/src/uipath_langchain/agent/attachments/output_files.py @@ -0,0 +1,224 @@ +"""Discovery and verification of job-attachment fields in an agent's output schema. + +An output schema may declare fields that hold a file (a job attachment). The +agent has no way to fill such a field on its own, so the runtime injects the +``create_output_file`` tool and tells the agent, in the system prompt, which +fields expect a file and what to write into them. + +Verification closes the loop. Nothing stops a model from inventing an attachment +id, so at termination every attachment reference in the output is checked against +the attachments actually linked to this job. A reference that is not there did +not come from the tool. +""" + +import uuid +from typing import Any, NamedTuple + +from pydantic import BaseModel +from uipath.platform import UiPath +from uipath.platform.common import UiPathConfig + +from .constants import OUTPUT_FILE_TOOL_NAME +from .job_attachments import get_job_attachment_paths +from .pydantic_json import extract_values_by_paths + + +class OutputFileField(NamedTuple): + """One declared output field that holds a file.""" + + path: str + """JSONPath to the field, e.g. ``$.report`` or ``$.exports[*]``.""" + + name: str + """The field's name as the agent sees it.""" + + description: str + """The field's description from the schema; empty when none was authored.""" + + required: bool + """Whether the schema requires the field to be filled.""" + + +def get_output_file_fields(model: type[BaseModel]) -> list[OutputFileField]: + """Describe every job-attachment field declared by an output model. + + Only top-level fields carry a name, description, and required flag that are + meaningful to state in a prompt; a nested attachment still gets a path so it + is verified, described by its path alone. + """ + fields = [] + for path in get_job_attachment_paths(model): + name = _field_name_from_path(path) + field_info = model.model_fields.get(name) + fields.append( + OutputFileField( + path=path, + name=name, + description=(field_info.description or "") if field_info else "", + required=field_info.is_required() if field_info else False, + ) + ) + return fields + + +def _field_name_from_path(path: str) -> str: + """The first segment of a JSONPath, e.g. ``$.exports[*]`` -> ``exports``.""" + return path.removeprefix("$.").split(".")[0].split("[")[0] + + +def missing_output_files( + fields: list[OutputFileField], output: dict[str, Any] +) -> list[OutputFileField]: + """Required file fields the agent left empty. + + A path that resolves to ``None`` counts as empty: an optional-shaped field + the model declined to fill still matches its JSONPath. + """ + return [ + field + for field in fields + if field.required and not _filled_values(output, field.path) + ] + + +def _filled_values(output: dict[str, Any], path: str) -> list[dict[str, Any]]: + """Attachment-shaped values at ``path``, skipping empty ones.""" + return [ + value + for value in extract_values_by_paths(output, [path]) + if isinstance(value, dict) and value + ] + + +def output_attachment_ids( + fields: list[OutputFileField], output: dict[str, Any] +) -> list[str]: + """Every attachment id referenced by the output's file fields.""" + ids = [] + for field in fields: + for value in _filled_values(output, field.path): + if value.get("ID"): + ids.append(str(value["ID"])) + return ids + + +async def unlinked_output_attachment_ids( + fields: list[OutputFileField], output: dict[str, Any] +) -> list[str]: + """Referenced attachment ids that are not linked to the current job. + + Returns an empty list when there is no job to check against — a local run + stores attachments outside Orchestrator, so there is nothing to verify. + """ + referenced = output_attachment_ids(fields, output) + if not referenced or not UiPathConfig.job_key: + return [] + + uipath = UiPath() + linked = { + str(key).lower() + for key in await uipath.jobs.list_attachments_async( + job_key=uuid.UUID(str(UiPathConfig.job_key)), + folder_key=UiPathConfig.folder_key, + ) + } + return [id for id in referenced if id.lower() not in linked] + + +_PROMPT_HEADER = """\ +**Output files** +These output fields hold a file, and the only way to fill one is with the \ +reference returned by the `{tool}` tool. Put each returned reference in its \ +matching field. +""" + +_PROMPT_REQUIRED_RULE = """\ +Create every required file before you end execution.""" + +_PROMPT_OPTIONAL_RULE = """\ +Create an optional file only when it serves the request; leaving one empty is a \ +valid answer.""" + +_PROMPT_FORMAT_RULE = """\ +If a field's description names a file format, use that format. Otherwise choose \ +the format that best fits the content, and give the file an extension that \ +matches it.""" + +_PROMPT_WORKSPACE_RULE = """\ +For anything you have already written to a file, or any non-text file, pass its \ +workspace path as `file_path` rather than re-emitting the body as `content`.""" + + +def build_output_files_prompt( + fields: list[OutputFileField], + *, + tool_name: str, + with_workspace: bool = False, +) -> str: + """Describe the declared output file fields and how to fill them. + + Returns an empty string when the output schema declares no file field, so + the caller can append the result unconditionally. + """ + if not fields: + return "" + + lines = [_PROMPT_HEADER.format(tool=tool_name)] + for field in fields: + suffix = " (required)" if field.required else " (optional)" + description = f" — {field.description}" if field.description else "" + lines.append(f"- `{field.name}`{suffix}{description}") + lines.append("") + # Stated per kind that is actually declared, so an all-optional schema is + # never told to produce a file and an all-required one is never told it may + # skip one. + if any(field.required for field in fields): + lines.append(_PROMPT_REQUIRED_RULE) + if any(not field.required for field in fields): + lines.append(_PROMPT_OPTIONAL_RULE) + lines.append(_PROMPT_FORMAT_RULE) + if with_workspace: + lines.append(_PROMPT_WORKSPACE_RULE) + return "\n".join(lines) + + +DEFAULT_MAX_OUTPUT_FILE_RETRIES = 2 + + +def _missing_files_message(fields: list[OutputFileField]) -> str: + names = ", ".join(f"'{field.name}'" for field in fields) + return ( + f"Execution cannot end: the output field(s) {names} must hold a file and " + f"are empty. Call `{OUTPUT_FILE_TOOL_NAME}` once per field, put each returned " + f"reference in its field, then end execution again." + ) + + +def _unlinked_ids_message(ids: list[str]) -> str: + listed = ", ".join(f"'{id}'" for id in ids) + return ( + f"Execution cannot end: the attachment reference(s) {listed} in the " + f"output do not belong to this job. Only a reference returned by " + f"`{OUTPUT_FILE_TOOL_NAME}` (or by a tool that produced a file) is valid. Create " + f"the file with `{OUTPUT_FILE_TOOL_NAME}` and use the reference it returns." + ) + + +async def diagnose_output_files( + fields: list[OutputFileField], output: dict[str, Any] +) -> str | None: + """Why this output cannot be accepted yet, or None when it can. + + Checked in order: a required file field left empty, then a reference to an + attachment that is not linked to this job. The message is written for the + agent to act on, so it names the field and the tool to call. + """ + missing = missing_output_files(fields, output) + if missing: + return _missing_files_message(missing) + + unlinked = await unlinked_output_attachment_ids(fields, output) + if unlinked: + return _unlinked_ids_message(unlinked) + + return None diff --git a/src/uipath_langchain/agent/react/agent.py b/src/uipath_langchain/agent/react/agent.py index 20dda70b9..121bfb6b7 100644 --- a/src/uipath_langchain/agent/react/agent.py +++ b/src/uipath_langchain/agent/react/agent.py @@ -13,7 +13,12 @@ from uipath_langchain.chat.hitl import IS_CONVERSATIONAL_CLIENT_SIDE_TOOL from ...runtime._citations import cas_deep_rag_citation_wrapper +from ..attachments.output_files import ( + DEFAULT_MAX_OUTPUT_FILE_RETRIES, + get_output_file_fields, +) from ..guardrails.actions import GuardrailAction +from ..tools.internal_tools.output_file_tool import OUTPUT_FILE_TOOL_NAME from ..tools.structured_tool_with_output_type import StructuredToolWithOutputType from .conversational_output_node import ( create_conversational_output_node, @@ -31,6 +36,10 @@ create_llm_node, ) from .memory_node import create_memory_recall_node +from .output_files_node import ( + create_output_files_node, + output_files_verified, +) from .router import ( create_route_agent, ) @@ -81,6 +90,16 @@ def create_agent( config = AgentGraphConfig() agent_tools = list(tools) + # Verification is gated on the tool being present: without a way to create the + # file, faulting a run for not having one would only break agents that never + # could. The caller decides, by supplying the tool, that files are expected. + output_file_fields = ( + get_output_file_fields(output_schema) + if output_schema is not None + and not config.is_conversational + and any(tool.name == OUTPUT_FILE_TOOL_NAME for tool in agent_tools) + else [] + ) flow_control_tools: list[BaseTool] = ( [] if config.is_conversational else create_flow_control_tools(output_schema) ) @@ -161,6 +180,23 @@ def create_agent( ) builder.add_node(AgentGraphNode.TERMINATE, terminate_with_guardrails_subgraph) + if output_file_fields: + builder.add_node( + AgentGraphNode.VERIFY_OUTPUT_FILES, + create_output_files_node( + output_file_fields, DEFAULT_MAX_OUTPUT_FILE_RETRIES + ), + ) + builder.add_conditional_edges( + AgentGraphNode.VERIFY_OUTPUT_FILES, + lambda state: ( + AgentGraphNode.TERMINATE + if output_files_verified(state) + else AgentGraphNode.AGENT + ), + [AgentGraphNode.TERMINATE, AgentGraphNode.AGENT], + ) + if with_conversational_output_node and output_schema is not None: builder.add_node( AgentGraphNode.GENERATE_CONVERSATIONAL_OUTPUT, @@ -216,8 +252,11 @@ def create_agent( *tool_node_names, AgentGraphNode.TERMINATE, ] + if output_file_fields: + target_node_names.append(AgentGraphNode.VERIFY_OUTPUT_FILES) route_agent = create_route_agent( valid_targets=target_node_names, + verify_output_files=bool(output_file_fields), ) builder.add_conditional_edges( diff --git a/src/uipath_langchain/agent/react/conversational_output_node.py b/src/uipath_langchain/agent/react/conversational_output_node.py index 339052062..dbae81e9a 100644 --- a/src/uipath_langchain/agent/react/conversational_output_node.py +++ b/src/uipath_langchain/agent/react/conversational_output_node.py @@ -1,15 +1,22 @@ -"""GENERATE_CONVERSATIONAL_OUTPUT node for the Agent graph. - -This intermediate node runs after AGENT for conversational agents whose -output schema declares custom fields beyond `uipath__agent_response_messages`. -It performs a focused LLM call with only the `set_conversational_output` -tool bound and `tool_choice="any"` to extract the structured output for the turn. +"""Structured-output extraction for conversational agents. + +A conversational agent's loop produces messages, but its output schema may +declare fields as well. Nothing in the message stream fills those, so they are +extracted afterwards by a focused LLM call with only the +`set_conversational_output` tool bound and `tool_choice="any"`, which forces the +model to answer with the declared fields. + +`create_conversational_output_extractor` is that call, independent of any graph. +`create_conversational_output_node` wraps it as the react graph's +GENERATE_CONVERSATIONAL_OUTPUT node; the advanced agent's wrapper graph builds +its own node over the same extractor. """ -from typing import TypeVar +from collections.abc import Awaitable, Callable, Sequence +from typing import Any, TypeVar from langchain_core.language_models import BaseChatModel -from langchain_core.messages import AIMessage, HumanMessage +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage from langchain_core.runnables.config import var_child_runnable_config from pydantic import BaseModel from uipath.agent.react import SET_CONVERSATIONAL_OUTPUT_TOOL @@ -34,15 +41,19 @@ StateT = TypeVar("StateT", bound=AgentGraphState) -def create_conversational_output_node( +def create_conversational_output_extractor( model: BaseChatModel, agent_output_schema: type[BaseModel], -): - """Build the conversational structured-output node. +) -> Callable[[Sequence[BaseMessage]], Awaitable[dict[str, Any]]]: + """Build the focused call that extracts the declared output fields. + + The returned coroutine takes the exchange's messages and returns the + structured-output arguments the model produced. It is graph-agnostic: the + caller decides where those arguments are stored. Args: model: The chat model to invoke for the extraction call. Reused from - the AGENT loop; rebinding is stateless. + the agent loop; rebinding is stateless. agent_output_schema: The agent's declared output schema. Used to construct the `set_conversational_output` tool with the LLM-fillable fields (`uipath__agent_response_messages` stripped). @@ -63,8 +74,8 @@ def create_conversational_output_node( ) output_prompt = get_generate_output_prompt() - async def conversational_output_node(state: StateT): - messages = [*state.messages, HumanMessage(content=output_prompt)] + async def extract(exchange_messages: Sequence[BaseMessage]) -> dict[str, Any]: + messages = [*exchange_messages, HumanMessage(content=output_prompt)] config = config_without_streaming(var_child_runnable_config.get(None)) try: @@ -115,6 +126,19 @@ async def conversational_output_node(state: StateT): category=UiPathErrorCategory.SYSTEM, ) - return {"inner_state": {"conversational_output": set_output_call["args"]}} + return set_output_call["args"] + + return extract + + +def create_conversational_output_node( + model: BaseChatModel, + agent_output_schema: type[BaseModel], +): + """Build the react graph's GENERATE_CONVERSATIONAL_OUTPUT node.""" + extract = create_conversational_output_extractor(model, agent_output_schema) + + async def conversational_output_node(state: StateT): + return {"inner_state": {"conversational_output": await extract(state.messages)}} return conversational_output_node diff --git a/src/uipath_langchain/agent/react/output_files_node.py b/src/uipath_langchain/agent/react/output_files_node.py new file mode 100644 index 000000000..f419e4cf5 --- /dev/null +++ b/src/uipath_langchain/agent/react/output_files_node.py @@ -0,0 +1,92 @@ +"""Verification gate for output file fields, run just before termination. + +Sits between the agent loop and TERMINATE whenever the output schema declares a +job-attachment field. It inspects the pending ``end_execution`` arguments and +lets termination proceed only when every required file field carries a reference +to an attachment that is actually linked to this job. + +A failure is not fatal. The node answers the ``end_execution`` tool call with a +corrective message and hands control back to the agent, which can create the +missing file and end again. Only a run that keeps failing that check faults, so +a model that forgets the tool costs a turn rather than the job. +""" + +from typing import Any + +from langchain_core.messages import ToolMessage +from langchain_core.messages.tool import ToolCall +from uipath.agent.react import END_EXECUTION_TOOL +from uipath.runtime.errors import UiPathErrorCategory + +from ..attachments.output_files import OutputFileField, diagnose_output_files +from ..exceptions import AgentRuntimeError, AgentRuntimeErrorCode +from .types import AgentGraphState +from .utils import extract_current_tool_call_index, find_latest_ai_message + + +def _pending_end_execution(state: AgentGraphState) -> ToolCall | None: + """The ``end_execution`` tool call the agent is currently making, if any.""" + last_message = find_latest_ai_message(state.messages) + if last_message is None or not last_message.tool_calls: + return None + index = extract_current_tool_call_index(state.messages) + if index is None: + return None + tool_call = last_message.tool_calls[index] + if tool_call["name"] != END_EXECUTION_TOOL.name: + return None + return tool_call + + +def _cleared() -> dict[str, Any]: + """State update that records a passing verification.""" + return {"inner_state": {"output_file_problem": None}} + + +def create_output_files_node(fields: list[OutputFileField], max_retries: int): + """Create the node that gates termination on the declared output files.""" + + async def output_files_node(state: AgentGraphState) -> dict[str, Any]: + tool_call = _pending_end_execution(state) + if tool_call is None: + return _cleared() + + problem = await diagnose_output_files(fields, tool_call["args"]) + if problem is None: + return _cleared() + + retries = state.inner_state.output_file_retries + if retries >= max_retries: + raise AgentRuntimeError( + code=AgentRuntimeErrorCode.OUTPUT_VALIDATION_ERROR, + title="Agent did not produce the required output file", + detail=( + f"{problem} The agent was given {max_retries} chance(s) to " + "correct this and did not. Verify the agent's prompt asks for " + "the file, and that the output schema's file fields are the " + "ones you intend." + ), + category=UiPathErrorCategory.USER, + ) + + return { + "messages": [ + ToolMessage( + content=problem, + tool_call_id=tool_call["id"], + name=END_EXECUTION_TOOL.name, + status="error", + ) + ], + "inner_state": { + "output_file_retries": retries + 1, + "output_file_problem": problem, + }, + } + + return output_files_node + + +def output_files_verified(state: AgentGraphState) -> bool: + """Whether the last verification pass let termination proceed.""" + return state.inner_state.output_file_problem is None diff --git a/src/uipath_langchain/agent/react/router.py b/src/uipath_langchain/agent/react/router.py index eb30c1fd4..b1664d0f1 100644 --- a/src/uipath_langchain/agent/react/router.py +++ b/src/uipath_langchain/agent/react/router.py @@ -3,6 +3,7 @@ from collections.abc import Container from typing import Literal +from uipath.agent.react import END_EXECUTION_TOOL from uipath.runtime.errors import UiPathErrorCategory from ..exceptions import AgentRuntimeError, AgentRuntimeErrorCode @@ -15,11 +16,14 @@ def create_route_agent( valid_targets: Container[str] | None = None, + verify_output_files: bool = False, ): """Create the conditional-edge routing function. Args: valid_targets: Allowed routing destinations + verify_output_files: Send ``end_execution`` through the output-file + verification node instead of straight to TERMINATE. Returns: Routing function for LangGraph conditional edges @@ -81,6 +85,8 @@ def route_agent( current_tool_name = current_tool_call["name"] if current_tool_name in FLOW_CONTROL_TOOLS: + if verify_output_files and current_tool_name == END_EXECUTION_TOOL.name: + return AgentGraphNode.VERIFY_OUTPUT_FILES return AgentGraphNode.TERMINATE if valid_targets is not None and current_tool_name not in valid_targets: diff --git a/src/uipath_langchain/agent/react/types.py b/src/uipath_langchain/agent/react/types.py index 9a890e8c5..12d792fd2 100644 --- a/src/uipath_langchain/agent/react/types.py +++ b/src/uipath_langchain/agent/react/types.py @@ -29,6 +29,8 @@ class InnerAgentGraphState(BaseModel): tools_storage: Annotated[dict[Hashable, Any], merge_dicts] = {} memory_injection: str = "" conversational_output: dict[str, Any] | None = None + output_file_retries: int = 0 + output_file_problem: str | None = None class InnerAgentGuardrailsGraphState(InnerAgentGraphState): @@ -66,6 +68,7 @@ class AgentGraphNode(StrEnum): LLM = "llm" TOOLS = "tools" GENERATE_CONVERSATIONAL_OUTPUT = "generate-conversational-output" + VERIFY_OUTPUT_FILES = "verify-output-files" TERMINATE = "terminate" GUARDED_TERMINATE = "guarded-terminate" MEMORY_RECALL = "memory_recall" diff --git a/src/uipath_langchain/agent/tools/internal_tools/output_file_tool.py b/src/uipath_langchain/agent/tools/internal_tools/output_file_tool.py new file mode 100644 index 000000000..bd5f38270 --- /dev/null +++ b/src/uipath_langchain/agent/tools/internal_tools/output_file_tool.py @@ -0,0 +1,233 @@ +"""Internal tool that publishes agent-authored content as a job attachment. + +Injected automatically — never configured by the user — whenever the agent's +output schema declares a job-attachment field. The tool creates the attachment, +links it to the current job, and returns the attachment ticket; the agent then +places that ticket in the declared output field. + +Two content sources, and which one is offered depends on the agent flavour: + +- ``content`` — the file body inline. The only source a standard agent has, + since it owns no filesystem. Text formats only. +- ``file_path`` — a path in the agent's own workspace, offered only when a + filesystem backend is present (advanced agents). Preferred there: the body + never round-trips through the model, so large and binary files work. +""" + +import mimetypes +from pathlib import Path +from typing import Any, Protocol, runtime_checkable + +from uipath.eval.mocks import mockable +from uipath.platform import UiPath +from uipath.platform.common import UiPathConfig +from uipath.runtime.errors import UiPathErrorCategory + +from uipath_langchain.agent.exceptions import ( + AgentRuntimeError, + AgentRuntimeErrorCode, +) +from uipath_langchain.agent.react.jsonschema_pydantic_converter import create_model +from uipath_langchain.agent.tools.structured_tool_with_output_type import ( + StructuredToolWithOutputType, +) +from uipath_langchain.agent.tools.tool_node import ToolWrapperMixin + +from ...attachments.constants import OUTPUT_FILE_TOOL_NAME +from .schema_utils import JOB_ATTACHMENT_DEFINITION + +__all__ = ["OUTPUT_FILE_TOOL_NAME", "create_output_file_tool", "guess_mime_type"] + + +_DEFAULT_MIME_TYPE = "application/octet-stream" + +# mimetypes has no entry for these on every supported Python, and they are among +# the formats an agent is most likely to pick for a generated document. +_EXTRA_MIME_TYPES = { + ".md": "text/markdown", + ".markdown": "text/markdown", + ".yaml": "application/yaml", + ".yml": "application/yaml", + ".jsonl": "application/jsonl", +} + +_TOOL_DESCRIPTION = ( + "Create a file and attach it to this job, then return the attachment " + "reference to put in the agent output field that expects a file. Call this " + "before ending execution: an output file field can only be filled with a " + "reference this tool returned." +) + +_FILE_NAME_DESCRIPTION = ( + "File name including the extension, e.g. 'summary.md' or 'accounts.csv'. " + "The extension determines the file's MIME type, so it must match the " + "format of the content." +) + +_CONTENT_DESCRIPTION = "The full text content of the file." + +_FILE_PATH_DESCRIPTION = ( + "Path of an existing file in your workspace to publish, e.g. '/report.md'. " + "Prefer this over 'content' for anything you have already written to a " + "file, and use it for any non-text file." +) + + +@runtime_checkable +class _BoundedPathBackend(Protocol): + """The part of a filesystem backend this tool needs: bounded path resolution.""" + + def _resolve_path(self, file_path: str) -> Path: ... + + +def output_file_tool_output_schema() -> dict[str, Any]: + """The tool's output schema: a single job-attachment ticket under ``file``.""" + return { + "type": "object", + "properties": { + "file": { + "$ref": "#/definitions/job-attachment", + "description": "Reference to the created file. Use this value for the output file field.", + } + }, + "required": ["file"], + "definitions": {"job-attachment": JOB_ATTACHMENT_DEFINITION}, + } + + +def _input_schema(*, with_file_path: bool) -> dict[str, Any]: + properties: dict[str, Any] = { + "file_name": {"type": "string", "description": _FILE_NAME_DESCRIPTION}, + "content": {"type": "string", "description": _CONTENT_DESCRIPTION}, + } + if with_file_path: + properties["file_path"] = { + "type": "string", + "description": _FILE_PATH_DESCRIPTION, + } + return { + "type": "object", + "properties": properties, + "required": ["file_name"], + } + + +def guess_mime_type(file_name: str) -> str: + """Resolve a file's MIME type from its extension.""" + suffix = Path(file_name).suffix.lower() + if suffix in _EXTRA_MIME_TYPES: + return _EXTRA_MIME_TYPES[suffix] + guessed, _ = mimetypes.guess_type(file_name) + return guessed or _DEFAULT_MIME_TYPE + + +def _resolve_source_path(backend: Any, file_path: str) -> Path: + """Resolve a workspace-relative path, rejecting anything outside the workspace.""" + if not isinstance(backend, _BoundedPathBackend): + raise AgentRuntimeError( + code=AgentRuntimeErrorCode.FILE_ERROR, + title="Workspace file paths are not available", + detail=( + f"'{OUTPUT_FILE_TOOL_NAME}' received a 'file_path' but this agent " + "has no workspace to read it from. Pass the file body in 'content' instead." + ), + category=UiPathErrorCategory.SYSTEM, + ) + # The backend's own resolver keeps the path inside the workspace root; it + # raises on traversal, so no separate containment check is needed here. + return backend._resolve_path(file_path) + + +class _OutputFileTool(StructuredToolWithOutputType, ToolWrapperMixin): + """Output type plus a state-updating wrapper, as the other attachment-producing tools have.""" + + +def create_output_file_tool(backend: Any | None = None) -> _OutputFileTool: + """Create the ``create_output_file`` tool. + + Args: + backend: The agent's filesystem backend, when it has one. ``file_path`` + is offered only for a backend that can resolve a workspace path; + for anything else the tool accepts inline ``content`` only, rather + than advertising an argument every use of which would fail. + """ + with_file_path = isinstance(backend, _BoundedPathBackend) + input_model = create_model(_input_schema(with_file_path=with_file_path)) + output_model = create_model(output_file_tool_output_schema()) + + async def create_output_file_fn(**kwargs: Any) -> dict[str, Any]: + file_name = kwargs.get("file_name") + content = kwargs.get("content") + file_path = kwargs.get("file_path") + + if not file_name: + raise ValueError("'file_name' is required.") + if not content and not file_path: + raise ValueError( + "Provide the file body in 'content'" + + ( + ", or an existing workspace path in 'file_path'." + if with_file_path + else "." + ) + ) + if content and file_path: + raise ValueError("'content' and 'file_path' are mutually exclusive.") + + # basename only: file_name reaches us from the model, and it names the + # attachment rather than a path on disk. + attachment_name = Path(file_name).name + + @mockable( + name=OUTPUT_FILE_TOOL_NAME, + description=_TOOL_DESCRIPTION, + input_schema=input_model.model_json_schema(), + output_schema=output_model.model_json_schema(), + example_calls=[], + ) + async def publish_output_file(**_tool_kwargs: Any) -> dict[str, Any]: + source_path = ( + _resolve_source_path(backend, file_path) if file_path else None + ) + if source_path is not None and not source_path.is_file(): + raise ValueError( + f"'{file_path}' does not exist in your workspace. Write the " + "file first, or pass its body in 'content'." + ) + + uipath = UiPath() + attachment_id = await uipath.jobs.create_attachment_async( + name=attachment_name, + content=content if source_path is None else None, + source_path=str(source_path) if source_path is not None else None, + job_key=UiPathConfig.job_key, + folder_key=UiPathConfig.folder_key, + ) + return { + "ID": str(attachment_id), + "FullName": attachment_name, + "MimeType": guess_mime_type(attachment_name), + } + + return {"file": await publish_output_file(**kwargs)} + + # Imported here to avoid a circular import at module load. + from uipath_langchain.agent.wrappers import get_job_attachment_wrapper + + tool = _OutputFileTool( + name=OUTPUT_FILE_TOOL_NAME, + description=_TOOL_DESCRIPTION, + args_schema=input_model, + coroutine=create_output_file_fn, + output_type=output_model, + metadata={ + "tool_type": "internal", + "display_name": OUTPUT_FILE_TOOL_NAME, + "args_schema": input_model, + "output_schema": output_model, + }, + ) + tool.set_tool_wrappers( + awrapper=get_job_attachment_wrapper(output_type=output_model) + ) + return tool diff --git a/src/uipath_langchain/agent/tools/internal_tools/schema_utils.py b/src/uipath_langchain/agent/tools/internal_tools/schema_utils.py index 2cc75ffaa..17b9e8b4f 100644 --- a/src/uipath_langchain/agent/tools/internal_tools/schema_utils.py +++ b/src/uipath_langchain/agent/tools/internal_tools/schema_utils.py @@ -2,6 +2,28 @@ from typing import Any +# The `job-attachment` definitions key is load-bearing: the JSON-schema-to-Pydantic +# converter derives the `__Job_attachment` marker type from it, and that marker is +# what `get_job_attachment_paths` looks for when discovering attachment fields. +JOB_ATTACHMENT_DEFINITION: dict[str, Any] = { + "type": "object", + "properties": { + "ID": {"type": "string", "description": "Orchestrator attachment key"}, + "FullName": {"type": "string", "description": "File name"}, + "MimeType": { + "type": "string", + "description": "The MIME type of the content", + }, + "Metadata": { + "type": "object", + "description": "Dictionary of metadata", + "additionalProperties": {"type": "string"}, + }, + }, + "required": ["ID", "FullName", "MimeType"], + "x-uipath-resource-kind": "JobAttachment", +} + # BatchTransform output schema with file attachment BATCH_TRANSFORM_OUTPUT_SCHEMA: dict[str, Any] = { "type": "object", @@ -12,26 +34,7 @@ } }, "required": ["result"], - "definitions": { - "job-attachment": { - "type": "object", - "properties": { - "ID": {"type": "string", "description": "Orchestrator attachment key"}, - "FullName": {"type": "string", "description": "File name"}, - "MimeType": { - "type": "string", - "description": "The MIME type of the content", - }, - "Metadata": { - "type": "object", - "description": "Dictionary of metadata", - "additionalProperties": {"type": "string"}, - }, - }, - "required": ["ID", "FullName", "MimeType"], - "x-uipath-resource-kind": "JobAttachment", - } - }, + "definitions": {"job-attachment": JOB_ATTACHMENT_DEFINITION}, } diff --git a/tests/agent/advanced/test_conversational_advanced_agent_graph.py b/tests/agent/advanced/test_conversational_advanced_agent_graph.py index cdb546d8e..ab9aa0aba 100644 --- a/tests/agent/advanced/test_conversational_advanced_agent_graph.py +++ b/tests/agent/advanced/test_conversational_advanced_agent_graph.py @@ -343,3 +343,150 @@ async def test_empty_history_still_produces_response() -> None: result = await graph.ainvoke({"messages": [HumanMessage(content="hi", id="u1")]}) assert len(result["uipath__agent_response_messages"]) == 1 + + +def _conversational_output_model(**properties: dict[str, Any]) -> type[BaseModel]: + """Build an output model the way the runtime does, from the agent's JSON schema.""" + from uipath_langchain.agent.react.jsonschema_pydantic_converter import ( + create_model as create_model_from_schema, + ) + + return create_model_from_schema( + { + "type": "object", + "properties": { + "uipath__agent_response_messages": {"type": "array"}, + **properties, + }, + } + ) + + +_OutputWithCustomFields = _conversational_output_model( + ticketId={"type": "string"}, resolved={"type": "boolean"} +) +_OutputMessagesOnly = _conversational_output_model() + + +class TestCustomOutputFields: + """Declared output fields are filled by the same extraction call standard + conversational agents use: the loop produces messages, not fields.""" + + def test_custom_fields_insert_the_extraction_node(self) -> None: + graph = create_conversational_advanced_agent_graph( + model=_mock_model(), + tools=[], + system_prompt="sys", + backend=None, + output_schema=_OutputWithCustomFields, + ) + + assert "generate_conversational_output" in set(graph.nodes) + + def test_messages_only_output_skips_the_extraction_node(self) -> None: + graph = create_conversational_advanced_agent_graph( + model=_mock_model(), + tools=[], + system_prompt="sys", + backend=None, + output_schema=_OutputMessagesOnly, + ) + + assert "generate_conversational_output" not in set(graph.nodes) + + def test_no_output_schema_skips_the_extraction_node(self) -> None: + graph = create_conversational_advanced_agent_graph( + model=_mock_model(), tools=[], system_prompt="sys", backend=None + ) + + assert "generate_conversational_output" not in set(graph.nodes) + + def test_extraction_state_key_does_not_collide_with_input(self) -> None: + class _Colliding(BaseModel): + messages: list[Any] = Field(default_factory=list) + uipath__conversational_output: str = "" + + graph = create_conversational_advanced_agent_graph( + model=_mock_model(), + tools=[], + system_prompt="sys", + backend=None, + input_schema=_Colliding, + output_schema=_OutputWithCustomFields, + ) + + assert "uipath__conversational_output_1" in graph.state_schema.model_fields + + @pytest.mark.asyncio + async def test_extracted_fields_are_merged_into_the_output(self) -> None: + with ( + patch( + "uipath_langchain.agent.advanced.agent.create_advanced_agent", + return_value=_fake_inner_agent(), + ), + patch( + "uipath_langchain.agent.advanced.agent.create_conversational_output_extractor", + return_value=_extractor({"ticketId": "INC-42", "resolved": True}), + ), + ): + graph = create_conversational_advanced_agent_graph( + model=_mock_model(), + tools=[], + system_prompt="sys", + backend=None, + output_schema=_OutputWithCustomFields, + ).compile() + result = await graph.ainvoke( + {"messages": [HumanMessage(content="hi", id="u1")]} + ) + + assert result["ticketId"] == "INC-42" + assert result["resolved"] is True + assert len(result["uipath__agent_response_messages"]) == 1 + + @pytest.mark.asyncio + async def test_extraction_sees_the_whole_transcript(self) -> None: + """A declared field's answer often lives in an earlier turn, so the + extraction gets the full history, as the standard path does.""" + seen: list[list[Any]] = [] + + async def record(messages: Any) -> dict[str, Any]: + seen.append(list(messages)) + return {"ticketId": "INC-1"} + + with ( + patch( + "uipath_langchain.agent.advanced.agent.create_advanced_agent", + return_value=_fake_inner_agent(), + ), + patch( + "uipath_langchain.agent.advanced.agent.create_conversational_output_extractor", + return_value=record, + ), + ): + graph = create_conversational_advanced_agent_graph( + model=_mock_model(), + tools=[], + system_prompt="sys", + backend=None, + output_schema=_OutputWithCustomFields, + ).compile() + await graph.ainvoke( + { + "messages": [ + HumanMessage(content="older turn", id="u0"), + HumanMessage(content="hi", id="u1"), + ] + } + ) + + assert [message.id for message in seen[0]] == ["u0", "u1", "ai-1"] + + +def _extractor(args: dict[str, Any]) -> Any: + """An extraction callable that always returns ``args``.""" + + async def extract(messages: Any) -> dict[str, Any]: + return args + + return extract diff --git a/tests/agent/advanced/test_create_advanced_agent_graph.py b/tests/agent/advanced/test_create_advanced_agent_graph.py index 154b3fc66..4cb480d1f 100644 --- a/tests/agent/advanced/test_create_advanced_agent_graph.py +++ b/tests/agent/advanced/test_create_advanced_agent_graph.py @@ -366,3 +366,82 @@ async def handler(prepared: ModelRequest[Any]) -> ModelResponse[Any]: assert captured[0].system_message is not None assert captured[0].system_message.text == ("runtime prompt\n\ndeepagents prompt") + + +class TestOutputFileVerification: + """The wrapper gates typed output on the declared output file fields.""" + + ATTACHMENT_ID = "11111111-1111-1111-1111-111111111111" + + @staticmethod + def _output_model(required: bool = True) -> type[BaseModel]: + from uipath_langchain.agent.react.jsonschema_pydantic_converter import ( + create_model as create_model_from_schema, + ) + from uipath_langchain.agent.tools.internal_tools.schema_utils import ( + JOB_ATTACHMENT_DEFINITION, + ) + + return create_model_from_schema( + { + "type": "object", + "properties": { + "summary": {"type": "string"}, + "report": {"$ref": "#/definitions/job-attachment"}, + }, + "required": ["report"] if required else [], + "definitions": {"job-attachment": JOB_ATTACHMENT_DEFINITION}, + } + ) + + @staticmethod + def _tools() -> list[Any]: + from uipath_langchain.agent.tools.internal_tools.output_file_tool import ( + create_output_file_tool, + ) + + return [create_output_file_tool()] + + def test_file_output_inserts_the_verification_node(self) -> None: + graph = _build(output_schema=self._output_model(), tools=self._tools()) + + assert "verify_output_files" in set(graph.nodes) + + def test_no_file_output_keeps_the_direct_edge(self) -> None: + graph = _build(output_schema=_Output, tools=self._tools()) + + assert "verify_output_files" not in set(graph.nodes) + + def test_file_output_without_the_tool_is_not_verified(self) -> None: + """With the feature off the tool is absent, so nothing gates the output.""" + graph = _build(output_schema=self._output_model(), tools=[]) + + assert "verify_output_files" not in set(graph.nodes) + + def test_retry_and_problem_state_fields_are_added(self) -> None: + graph = _build(output_schema=self._output_model(), tools=self._tools()) + fields = set(graph.state_schema.model_fields) + + assert "uipath__output_file_retries" in fields + assert "uipath__output_file_problem" in fields + + def test_no_file_output_adds_no_verification_state(self) -> None: + graph = _build(output_schema=_Output, tools=self._tools()) + fields = set(graph.state_schema.model_fields) + + assert "uipath__output_file_retries" not in fields + assert "uipath__output_file_problem" not in fields + + async def test_verification_state_is_not_forwarded_as_agent_input(self) -> None: + """The keys are internal, so transform_input must not treat them as inputs.""" + graph = _build( + input_schema=_Input, + output_schema=self._output_model(), + tools=self._tools(), + ) + state = graph.state_schema(book={"title": "x"}, question="q") + + update = await graph.nodes["transform_input"].runnable.ainvoke(state) + + assert "messages" in update + assert "uipath__output_file_retries" not in update diff --git a/tests/agent/attachments/test_output_files.py b/tests/agent/attachments/test_output_files.py new file mode 100644 index 000000000..f688e0b3b --- /dev/null +++ b/tests/agent/attachments/test_output_files.py @@ -0,0 +1,293 @@ +"""Tests for output-schema file field discovery, prompting, and verification.""" + +from typing import Any + +import pytest + +from uipath_langchain.agent.attachments.output_files import ( + build_output_files_prompt, + get_output_file_fields, + missing_output_files, + output_attachment_ids, + unlinked_output_attachment_ids, +) +from uipath_langchain.agent.react.jsonschema_pydantic_converter import create_model +from uipath_langchain.agent.tools.internal_tools.schema_utils import ( + JOB_ATTACHMENT_DEFINITION, +) + +ATTACHMENT_ID = "11111111-1111-1111-1111-111111111111" +OTHER_ATTACHMENT_ID = "22222222-2222-2222-2222-222222222222" + + +def build_output_model(properties: dict[str, Any], required: list[str] | None = None): + return create_model( + { + "type": "object", + "properties": properties, + "required": required or [], + "definitions": {"job-attachment": JOB_ATTACHMENT_DEFINITION}, + } + ) + + +def ticket(attachment_id: str = ATTACHMENT_ID) -> dict[str, str]: + return { + "ID": attachment_id, + "FullName": "report.md", + "MimeType": "text/markdown", + } + + +class TestGetOutputFileFields: + def test_no_attachment_fields_returns_empty(self): + model = build_output_model({"summary": {"type": "string"}}) + assert get_output_file_fields(model) == [] + + def test_discovers_name_description_and_required(self): + model = build_output_model( + { + "summary": {"type": "string"}, + "report": { + "$ref": "#/definitions/job-attachment", + "description": "The generated report", + }, + }, + required=["summary", "report"], + ) + + fields = get_output_file_fields(model) + + assert len(fields) == 1 + assert fields[0].path == "$.report" + assert fields[0].name == "report" + assert fields[0].description == "The generated report" + assert fields[0].required is True + + def test_optional_field_is_not_required(self): + model = build_output_model( + {"report": {"$ref": "#/definitions/job-attachment"}}, + ) + + assert get_output_file_fields(model)[0].required is False + + def test_array_of_attachments_keeps_the_field_name(self): + model = build_output_model( + { + "exports": { + "type": "array", + "items": {"$ref": "#/definitions/job-attachment"}, + "description": "Every exported file", + } + }, + required=["exports"], + ) + + field = get_output_file_fields(model)[0] + + assert field.path == "$.exports[*]" + assert field.name == "exports" + assert field.description == "Every exported file" + + +class TestBuildOutputFilesPrompt: + def test_empty_fields_produce_no_prompt(self): + assert build_output_files_prompt([], tool_name="create_output_file") == "" + + def test_lists_each_field_with_its_description(self): + model = build_output_model( + { + "report": { + "$ref": "#/definitions/job-attachment", + "description": "The generated report", + }, + "extras": { + "type": "array", + "items": {"$ref": "#/definitions/job-attachment"}, + }, + }, + required=["report"], + ) + + prompt = build_output_files_prompt( + get_output_file_fields(model), tool_name="create_output_file" + ) + + assert "create_output_file" in prompt + assert "`report` (required) — The generated report" in prompt + assert "`extras` (optional)" in prompt + assert "choose the format that best fits the content" in prompt + + def test_required_field_is_told_to_produce_the_file(self): + model = build_output_model( + {"report": {"$ref": "#/definitions/job-attachment"}}, required=["report"] + ) + + prompt = build_output_files_prompt( + get_output_file_fields(model), tool_name="create_output_file" + ) + + assert "Create every required file" in prompt + assert "only when it serves the request" not in prompt + + def test_optional_field_is_not_told_to_produce_the_file(self): + """The runtime does not require it, so the prompt must not demand it.""" + model = build_output_model({"report": {"$ref": "#/definitions/job-attachment"}}) + + prompt = build_output_files_prompt( + get_output_file_fields(model), tool_name="create_output_file" + ) + + assert "only when it serves the request" in prompt + assert "Create every required file" not in prompt + + def test_mixed_fields_state_both_rules(self): + model = build_output_model( + { + "report": {"$ref": "#/definitions/job-attachment"}, + "extras": { + "type": "array", + "items": {"$ref": "#/definitions/job-attachment"}, + }, + }, + required=["report"], + ) + + prompt = build_output_files_prompt( + get_output_file_fields(model), tool_name="create_output_file" + ) + + assert "Create every required file" in prompt + assert "only when it serves the request" in prompt + + def test_workspace_rule_only_when_requested(self): + model = build_output_model({"report": {"$ref": "#/definitions/job-attachment"}}) + fields = get_output_file_fields(model) + + assert "file_path" not in build_output_files_prompt( + fields, tool_name="create_output_file" + ) + assert "file_path" in build_output_files_prompt( + fields, tool_name="create_output_file", with_workspace=True + ) + + +class TestMissingOutputFiles: + @pytest.fixture + def fields(self): + model = build_output_model( + { + "report": {"$ref": "#/definitions/job-attachment"}, + "optional_export": {"$ref": "#/definitions/job-attachment"}, + }, + required=["report"], + ) + return get_output_file_fields(model) + + def test_required_field_absent_is_reported(self, fields): + missing = missing_output_files(fields, {"summary": "done"}) + + assert [field.name for field in missing] == ["report"] + + def test_required_field_null_is_reported(self, fields): + missing = missing_output_files(fields, {"report": None}) + + assert [field.name for field in missing] == ["report"] + + def test_required_field_filled_is_not_reported(self, fields): + assert missing_output_files(fields, {"report": ticket()}) == [] + + def test_optional_field_absent_is_not_reported(self, fields): + assert missing_output_files(fields, {"report": ticket()}) == [] + + +class TestOutputAttachmentIds: + @pytest.fixture + def fields(self): + model = build_output_model( + { + "report": {"$ref": "#/definitions/job-attachment"}, + "exports": { + "type": "array", + "items": {"$ref": "#/definitions/job-attachment"}, + }, + } + ) + return get_output_file_fields(model) + + def test_collects_ids_from_scalar_and_array_fields(self, fields): + ids = output_attachment_ids( + fields, + {"report": ticket(), "exports": [ticket(OTHER_ATTACHMENT_ID)]}, + ) + + assert sorted(ids) == sorted([ATTACHMENT_ID, OTHER_ATTACHMENT_ID]) + + def test_ignores_empty_and_malformed_values(self, fields): + ids = output_attachment_ids( + fields, {"report": None, "exports": [{"FullName": "x.md"}]} + ) + + assert ids == [] + + +class TestUnlinkedOutputAttachmentIds: + @pytest.fixture + def fields(self): + model = build_output_model( + {"report": {"$ref": "#/definitions/job-attachment"}}, required=["report"] + ) + return get_output_file_fields(model) + + async def test_no_job_key_skips_verification(self, fields, monkeypatch): + monkeypatch.delenv("UIPATH_JOB_KEY", raising=False) + + assert await unlinked_output_attachment_ids(fields, {"report": ticket()}) == [] + + async def test_linked_attachment_passes(self, fields, monkeypatch): + _patch_job(monkeypatch, linked=[ATTACHMENT_ID]) + + assert await unlinked_output_attachment_ids(fields, {"report": ticket()}) == [] + + async def test_linked_attachment_matches_case_insensitively( + self, fields, monkeypatch + ): + _patch_job(monkeypatch, linked=[ATTACHMENT_ID.upper()]) + + assert await unlinked_output_attachment_ids(fields, {"report": ticket()}) == [] + + async def test_unknown_attachment_is_reported(self, fields, monkeypatch): + _patch_job(monkeypatch, linked=[OTHER_ATTACHMENT_ID]) + + unlinked = await unlinked_output_attachment_ids(fields, {"report": ticket()}) + + assert unlinked == [ATTACHMENT_ID] + + async def test_empty_output_does_not_call_the_platform(self, fields, monkeypatch): + calls: list[Any] = [] + _patch_job(monkeypatch, linked=[], calls=calls) + + assert await unlinked_output_attachment_ids(fields, {}) == [] + assert calls == [] + + +def _patch_job( + monkeypatch, *, linked: list[str], calls: list[Any] | None = None +) -> None: + """Point the verification at a fake job with ``linked`` attachments.""" + monkeypatch.setenv("UIPATH_JOB_KEY", "33333333-3333-3333-3333-333333333333") + monkeypatch.delenv("UIPATH_FOLDER_KEY", raising=False) + + class FakeJobs: + async def list_attachments_async(self, **kwargs: Any) -> list[str]: + if calls is not None: + calls.append(kwargs) + return linked + + class FakeUiPath: + jobs = FakeJobs() + + monkeypatch.setattr( + "uipath_langchain.agent.attachments.output_files.UiPath", + lambda *args, **kwargs: FakeUiPath(), + ) diff --git a/tests/agent/react/test_output_files_node.py b/tests/agent/react/test_output_files_node.py new file mode 100644 index 000000000..2b35ff0e6 --- /dev/null +++ b/tests/agent/react/test_output_files_node.py @@ -0,0 +1,213 @@ +"""Tests for the output-file verification node and its graph wiring.""" + +from typing import Any + +import pytest +from langchain_core.language_models.fake_chat_models import GenericFakeChatModel +from langchain_core.messages import AIMessage, HumanMessage, SystemMessage, ToolMessage +from uipath.agent.react import END_EXECUTION_TOOL, RAISE_ERROR_TOOL +from uipath.runtime.errors import UiPathErrorCategory + +from uipath_langchain.agent.attachments.output_files import get_output_file_fields +from uipath_langchain.agent.exceptions import ( + AgentRuntimeError, + AgentRuntimeErrorCode, +) +from uipath_langchain.agent.react.agent import create_agent +from uipath_langchain.agent.react.jsonschema_pydantic_converter import create_model +from uipath_langchain.agent.react.output_files_node import ( + create_output_files_node, + output_files_verified, +) +from uipath_langchain.agent.react.types import AgentGraphNode, AgentGraphState +from uipath_langchain.agent.tools.internal_tools.output_file_tool import ( + OUTPUT_FILE_TOOL_NAME, + create_output_file_tool, +) +from uipath_langchain.agent.tools.internal_tools.schema_utils import ( + JOB_ATTACHMENT_DEFINITION, +) + +ATTACHMENT_ID = "11111111-1111-1111-1111-111111111111" +OTHER_ATTACHMENT_ID = "22222222-2222-2222-2222-222222222222" +JOB_KEY = "33333333-3333-3333-3333-333333333333" + + +def output_schema(required: list[str] | None = None) -> dict[str, Any]: + return { + "type": "object", + "properties": { + "summary": {"type": "string"}, + "report": { + "$ref": "#/definitions/job-attachment", + "description": "The generated report", + }, + }, + "required": required if required is not None else ["summary", "report"], + "definitions": {"job-attachment": JOB_ATTACHMENT_DEFINITION}, + } + + +def ticket(attachment_id: str = ATTACHMENT_ID) -> dict[str, str]: + return { + "ID": attachment_id, + "FullName": "report.md", + "MimeType": "text/markdown", + } + + +def state_ending_with(args: dict[str, Any], *, tool_name: str | None = None) -> Any: + """State whose latest AI message calls a flow-control tool with ``args``.""" + return AgentGraphState( + messages=[ + HumanMessage(content="go"), + AIMessage( + content="", + tool_calls=[ + { + "name": tool_name or END_EXECUTION_TOOL.name, + "args": args, + "id": "call-1", + } + ], + ), + ] + ) + + +@pytest.fixture +def fields(): + return get_output_file_fields(create_model(output_schema())) + + +@pytest.fixture +def linked_job(monkeypatch): + """A current job whose only linked attachment is ATTACHMENT_ID.""" + monkeypatch.setenv("UIPATH_JOB_KEY", JOB_KEY) + monkeypatch.delenv("UIPATH_FOLDER_KEY", raising=False) + + class FakeJobs: + async def list_attachments_async(self, **kwargs: Any) -> list[str]: + return [ATTACHMENT_ID] + + class FakeUiPath: + jobs = FakeJobs() + + monkeypatch.setattr( + "uipath_langchain.agent.attachments.output_files.UiPath", + lambda *args, **kwargs: FakeUiPath(), + ) + + +class TestOutputFilesNode: + async def test_valid_output_clears_the_problem(self, fields, linked_job): + node = create_output_files_node(fields, max_retries=2) + + update = await node(state_ending_with({"summary": "s", "report": ticket()})) + + assert update["inner_state"]["output_file_problem"] is None + assert "messages" not in update + + async def test_missing_required_file_returns_a_corrective_tool_message( + self, fields, linked_job + ): + node = create_output_files_node(fields, max_retries=2) + + update = await node(state_ending_with({"summary": "s"})) + + message = update["messages"][0] + assert isinstance(message, ToolMessage) + assert message.tool_call_id == "call-1" + assert message.status == "error" + assert OUTPUT_FILE_TOOL_NAME in message.content + assert "'report'" in message.content + assert update["inner_state"]["output_file_retries"] == 1 + + async def test_unlinked_attachment_returns_a_corrective_tool_message( + self, fields, linked_job + ): + node = create_output_files_node(fields, max_retries=2) + + update = await node( + state_ending_with({"summary": "s", "report": ticket(OTHER_ATTACHMENT_ID)}) + ) + + assert OTHER_ATTACHMENT_ID in update["messages"][0].content + assert update["inner_state"]["output_file_retries"] == 1 + + async def test_retries_are_capped_then_the_run_faults(self, fields, linked_job): + node = create_output_files_node(fields, max_retries=2) + state = state_ending_with({"summary": "s"}) + state.inner_state.output_file_retries = 2 + + with pytest.raises(AgentRuntimeError) as exc_info: + await node(state) + + assert exc_info.value.error_info.code == AgentRuntimeError.full_code( + AgentRuntimeErrorCode.OUTPUT_VALIDATION_ERROR + ) + assert exc_info.value.error_info.category == UiPathErrorCategory.USER + + async def test_optional_file_field_left_empty_passes(self, linked_job): + fields = get_output_file_fields( + create_model(output_schema(required=["summary"])) + ) + node = create_output_files_node(fields, max_retries=2) + + update = await node(state_ending_with({"summary": "s"})) + + assert update["inner_state"]["output_file_problem"] is None + + async def test_non_end_execution_call_is_left_alone(self, fields, linked_job): + node = create_output_files_node(fields, max_retries=2) + + update = await node( + state_ending_with({"message": "boom"}, tool_name=RAISE_ERROR_TOOL.name) + ) + + assert update["inner_state"]["output_file_problem"] is None + assert "messages" not in update + + +class TestOutputFilesVerified: + def test_cleared_problem_is_verified(self): + assert output_files_verified(AgentGraphState()) is True + + def test_recorded_problem_is_not_verified(self): + state = AgentGraphState() + state.inner_state.output_file_problem = "missing" + + assert output_files_verified(state) is False + + +class TestGraphWiring: + def build(self, schema: dict[str, Any]): + return create_agent( + model=GenericFakeChatModel(messages=iter([])), + tools=[create_output_file_tool()], + messages=[SystemMessage(content="sys"), HumanMessage(content="go")], + output_schema=create_model(schema), + ).compile() + + def test_file_output_adds_the_verification_node(self): + graph = self.build(output_schema()) + + assert AgentGraphNode.VERIFY_OUTPUT_FILES in graph.get_graph().nodes + + def test_no_file_output_leaves_the_graph_unchanged(self): + graph = self.build( + {"type": "object", "properties": {"summary": {"type": "string"}}} + ) + + assert AgentGraphNode.VERIFY_OUTPUT_FILES not in graph.get_graph().nodes + + def test_verification_can_reach_both_terminate_and_agent(self): + edges = self.build(output_schema()).get_graph().edges + targets = { + edge.target + for edge in edges + if edge.source == AgentGraphNode.VERIFY_OUTPUT_FILES + } + + assert AgentGraphNode.TERMINATE in targets + assert AgentGraphNode.AGENT in targets diff --git a/tests/agent/tools/internal_tools/test_output_file_tool.py b/tests/agent/tools/internal_tools/test_output_file_tool.py new file mode 100644 index 000000000..f43666865 --- /dev/null +++ b/tests/agent/tools/internal_tools/test_output_file_tool.py @@ -0,0 +1,196 @@ +"""Tests for the create_output_file internal tool.""" + +from pathlib import Path +from typing import Any + +import pytest +from langchain_core.tools import StructuredTool +from pydantic import BaseModel + +from uipath_langchain.agent.tools.internal_tools.output_file_tool import ( + OUTPUT_FILE_TOOL_NAME, + create_output_file_tool, + guess_mime_type, +) + +ATTACHMENT_ID = "11111111-1111-1111-1111-111111111111" + + +def args_schema(tool: StructuredTool) -> type[BaseModel]: + """The tool's argument model, narrowed from the permissive declared union.""" + schema = tool.args_schema + assert isinstance(schema, type) and issubclass(schema, BaseModel) + return schema + + +async def call(tool: StructuredTool, **kwargs: Any) -> dict[str, Any]: + """Invoke the tool's coroutine directly, bypassing argument validation.""" + coroutine = tool.coroutine + assert coroutine is not None + result = await coroutine(**kwargs) + assert isinstance(result, dict) + return result + + +class FakeBackend: + """Stands in for a filesystem backend with bounded path resolution.""" + + def __init__(self, root: Path) -> None: + self.root = root + + def _resolve_path(self, file_path: str) -> Path: + resolved = (self.root / file_path.lstrip("/")).resolve() + if resolved != self.root and self.root not in resolved.parents: + raise ValueError(f"Workspace path escapes root: {file_path}") + return resolved + + +@pytest.fixture +def created(monkeypatch) -> list[dict[str, Any]]: + """Capture every attachment the tool creates.""" + calls: list[dict[str, Any]] = [] + + class FakeJobs: + async def create_attachment_async(self, **kwargs: Any) -> str: + calls.append(kwargs) + return ATTACHMENT_ID + + class FakeUiPath: + jobs = FakeJobs() + + monkeypatch.setattr( + "uipath_langchain.agent.tools.internal_tools.output_file_tool.UiPath", + lambda *args, **kwargs: FakeUiPath(), + ) + return calls + + +class TestGuessMimeType: + @pytest.mark.parametrize( + ("file_name", "expected"), + [ + ("report.md", "text/markdown"), + ("accounts.csv", "text/csv"), + ("data.json", "application/json"), + ("notes.txt", "text/plain"), + ("config.yaml", "application/yaml"), + ("book.pdf", "application/pdf"), + ("mystery", "application/octet-stream"), + ("REPORT.MD", "text/markdown"), + ], + ) + def test_extension_drives_the_mime_type(self, file_name, expected): + assert guess_mime_type(file_name) == expected + + +class TestToolSchema: + def test_content_only_without_a_backend(self): + properties = args_schema(create_output_file_tool()).model_json_schema()[ + "properties" + ] + + assert set(properties) == {"file_name", "content"} + + def test_backend_adds_file_path(self, tmp_path): + tool = create_output_file_tool(FakeBackend(tmp_path)) + properties = args_schema(tool).model_json_schema()["properties"] + + assert set(properties) == {"file_name", "content", "file_path"} + + def test_only_file_name_is_required(self): + schema = args_schema(create_output_file_tool()).model_json_schema() + + assert schema["required"] == ["file_name"] + + def test_tool_is_named_for_the_prompt(self): + assert create_output_file_tool().name == OUTPUT_FILE_TOOL_NAME + + +class TestCreateFromContent: + async def test_uploads_the_content_and_returns_a_ticket(self, created): + tool = create_output_file_tool() + + result = await call(tool, file_name="report.md", content="# Report") + + assert result == { + "file": { + "ID": ATTACHMENT_ID, + "FullName": "report.md", + "MimeType": "text/markdown", + } + } + assert created[0]["name"] == "report.md" + assert created[0]["content"] == "# Report" + assert created[0]["source_path"] is None + + async def test_file_name_is_reduced_to_its_basename(self, created): + tool = create_output_file_tool() + + result = await call(tool, file_name="../../etc/passwd.txt", content="nope") + + assert result["file"]["FullName"] == "passwd.txt" + assert created[0]["name"] == "passwd.txt" + + async def test_no_source_is_rejected(self, created): + tool = create_output_file_tool() + + with pytest.raises(ValueError, match="'content'"): + await call(tool, file_name="report.md") + + assert created == [] + + +class TestCreateFromWorkspacePath: + async def test_uploads_the_workspace_file(self, created, tmp_path): + (tmp_path / "report.md").write_text("# Report") + tool = create_output_file_tool(FakeBackend(tmp_path)) + + result = await call(tool, file_name="report.md", file_path="/report.md") + + assert result["file"]["ID"] == ATTACHMENT_ID + assert created[0]["source_path"] == str(tmp_path / "report.md") + assert created[0]["content"] is None + + async def test_missing_workspace_file_is_rejected(self, created, tmp_path): + tool = create_output_file_tool(FakeBackend(tmp_path)) + + with pytest.raises(ValueError, match="does not exist in your workspace"): + await call(tool, file_name="report.md", file_path="/absent.md") + + assert created == [] + + async def test_path_escaping_the_workspace_is_rejected(self, created, tmp_path): + tool = create_output_file_tool(FakeBackend(tmp_path)) + + with pytest.raises(ValueError, match="escapes root"): + await call(tool, file_name="passwd.txt", file_path="../../etc/passwd") + + assert created == [] + + async def test_content_and_file_path_together_are_rejected(self, created, tmp_path): + tool = create_output_file_tool(FakeBackend(tmp_path)) + + with pytest.raises(ValueError, match="mutually exclusive"): + await call(tool, file_name="report.md", content="x", file_path="/report.md") + + assert created == [] + + +class _BackendWithoutPaths: + """A backend that cannot resolve a workspace path.""" + + +class TestBackendWithoutPathResolution: + def test_file_path_is_not_offered(self): + """Advertising it would give the model an argument that always fails.""" + tool = create_output_file_tool(_BackendWithoutPaths()) + properties = args_schema(tool).model_json_schema()["properties"] + + assert set(properties) == {"file_name", "content"} + + async def test_content_still_works(self, created): + tool = create_output_file_tool(_BackendWithoutPaths()) + + result = await call(tool, file_name="report.md", content="# Report") + + assert result["file"]["ID"] == ATTACHMENT_ID