feat: let agents produce output files as job attachments

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 (
@@ -157,7 +176,7 @@
     )
 
 
 def create_advanced_agent_graph(
     model: BaseChatModel,
     tools: Sequence[BaseTool],
     system_prompt: str | Callable[[dict[str, Any]], str],
@@ -175,13 +194,25 @@
     ``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 []
     )
     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 @@
         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 @@
         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,19 +314,29 @@
     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
 
 
 def create_conversational_advanced_agent_graph(
     model: BaseChatModel,
     tools: Sequence[BaseTool],
     system_prompt: str | Callable[[dict[str, Any]], str],
     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 @@
     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 @@
             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 @@
             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,30 +451,67 @@
             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
             )
             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

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"