From c19f253882065aaa84ab7840a9ce3d33ef276b66 Mon Sep 17 00:00:00 2001 From: Maxwell Du <60411452+maxduu@users.noreply.github.com> Date: Tue, 30 Jun 2026 23:34:03 -0400 Subject: [PATCH 1/4] feat: convo agent output --- src/uipath_langchain/agent/react/agent.py | 30 +++- src/uipath_langchain/agent/react/constants.py | 2 + .../agent/react/conversational_output_node.py | 92 +++++++++++ .../agent/react/router_conversational.py | 46 ++++-- .../agent/react/terminate_node.py | 84 +++++++++- .../agent/react/tools/tools.py | 21 +++ src/uipath_langchain/agent/react/types.py | 13 +- src/uipath_langchain/agent/react/utils.py | 41 +++++ .../internal_tools/analyze_files_tool.py | 21 +-- src/uipath_langchain/agent/tools/utils.py | 17 ++- .../react/test_conversational_output_node.py | 143 ++++++++++++++++++ tests/agent/react/test_create_agent.py | 72 +++++++++ tests/agent/react/test_flow_control_tools.py | 44 ++++++ .../agent/react/test_router_conversational.py | 63 ++++++++ tests/agent/react/test_terminate_node.py | 73 +++++++-- tests/agent/react/test_utils.py | 74 +++++++++ .../internal_tools/test_analyze_files_tool.py | 26 ---- tests/agent/tools/test_utils.py | 29 ++++ 18 files changed, 813 insertions(+), 78 deletions(-) create mode 100644 src/uipath_langchain/agent/react/conversational_output_node.py create mode 100644 tests/agent/react/test_conversational_output_node.py create mode 100644 tests/agent/react/test_flow_control_tools.py diff --git a/src/uipath_langchain/agent/react/agent.py b/src/uipath_langchain/agent/react/agent.py index 44f09052c..7e2342521 100644 --- a/src/uipath_langchain/agent/react/agent.py +++ b/src/uipath_langchain/agent/react/agent.py @@ -15,6 +15,9 @@ from ...runtime._citations import cas_deep_rag_citation_wrapper from ..guardrails.actions import GuardrailAction from ..tools.structured_tool_with_output_type import StructuredToolWithOutputType +from .conversational_output_node import ( + create_conversational_output_node, +) from .guardrails.guardrails_subgraph import ( create_agent_init_guardrails_subgraph, create_agent_terminate_guardrails_subgraph, @@ -42,7 +45,10 @@ AgentGraphState, MemoryConfig, ) -from .utils import create_state_with_input +from .utils import ( + create_state_with_input, + has_custom_conversational_output_fields, +) InputT = TypeVar("InputT", bound=BaseModel) OutputT = TypeVar("OutputT", bound=BaseModel) @@ -126,6 +132,11 @@ def create_agent( terminate_node = create_terminate_node(output_schema, config.is_conversational) + with_conversational_output_node = ( + config.is_conversational + and has_custom_conversational_output_fields(output_schema) + ) + CompleteAgentGraphState = create_state_with_input( input_schema if input_schema is not None else BaseModel ) @@ -150,6 +161,16 @@ def create_agent( ) builder.add_node(AgentGraphNode.TERMINATE, terminate_with_guardrails_subgraph) + if with_conversational_output_node and output_schema is not None: + builder.add_node( + AgentGraphNode.GENERATE_CONVERSATIONAL_OUTPUT, + create_conversational_output_node(model, output_schema), + ) + builder.add_edge( + AgentGraphNode.GENERATE_CONVERSATIONAL_OUTPUT, + AgentGraphNode.TERMINATE, + ) + if memory: memory_recall = create_memory_recall_node(memory, input_schema=input_schema) builder.add_node(AgentGraphNode.MEMORY_RECALL, memory_recall) @@ -182,7 +203,12 @@ def create_agent( *tool_node_names, AgentGraphNode.TERMINATE, ] - route_agent = create_route_agent_conversational(valid_targets=target_node_names) + if with_conversational_output_node: + target_node_names.append(AgentGraphNode.GENERATE_CONVERSATIONAL_OUTPUT) + route_agent = create_route_agent_conversational( + valid_targets=target_node_names, + with_generate_output_node=with_conversational_output_node, + ) else: target_node_names = [ AgentGraphNode.AGENT, diff --git a/src/uipath_langchain/agent/react/constants.py b/src/uipath_langchain/agent/react/constants.py index aec74daa6..d28d7b789 100644 --- a/src/uipath_langchain/agent/react/constants.py +++ b/src/uipath_langchain/agent/react/constants.py @@ -1,2 +1,4 @@ DEFAULT_MAX_CONSECUTIVE_THINKING_MESSAGES = 0 DEFAULT_MAX_LLM_MESSAGES = 25 + +UIPATH_CONVERSATIONAL_AGENT_RESPONSE_MESSAGES_FIELD = "uipath__agent_response_messages" diff --git a/src/uipath_langchain/agent/react/conversational_output_node.py b/src/uipath_langchain/agent/react/conversational_output_node.py new file mode 100644 index 000000000..e167a67ce --- /dev/null +++ b/src/uipath_langchain/agent/react/conversational_output_node.py @@ -0,0 +1,92 @@ +"""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 reliably extract the structured +output for the turn — decoupling conversational quality from schema +compliance. + +The LLM call is tagged with `TAG_NOSTREAM` so its tokens / events never +reach the chat-UI message stream. TERMINATE then reads the tool call's +args from `state.messages[-1]`. +""" + +from typing import TypeVar + +from langchain_core.language_models import BaseChatModel +from langchain_core.messages import AIMessage, HumanMessage +from langchain_core.runnables.config import var_child_runnable_config +from pydantic import BaseModel +from uipath.agent.react.conversational_prompts import ( + get_generate_output_prompt, +) +from uipath.runtime.errors import UiPathErrorCategory + +from uipath_langchain.chat.handlers import get_payload_handler + +from ..exceptions import AgentRuntimeError, AgentRuntimeErrorCode +from ..exceptions.licensing import raise_for_provider_http_error +from ..tools.utils import config_without_streaming +from .tools.tools import create_set_conversational_output_tool +from .types import AgentGraphState + +StateT = TypeVar("StateT", bound=AgentGraphState) + + +def create_conversational_output_node( + model: BaseChatModel, + agent_output_schema: type[BaseModel], +): + """Build the focused structured-output extraction node. + + Args: + model: The chat model to invoke for the extraction call. Reused from + 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). + """ + set_conversational_output_tool = create_set_conversational_output_tool( + agent_output_schema + ) + # Disable streaming on this internal LLM call + non_streaming_model = model.model_copy(update={"disable_streaming": True}) + payload_handler = get_payload_handler(non_streaming_model) + binding_kwargs = payload_handler.get_tool_binding_kwargs( + tools=[set_conversational_output_tool], + tool_choice="any", + parallel_tool_calls=False, + ) + llm = non_streaming_model.bind_tools( + [set_conversational_output_tool], **binding_kwargs + ) + output_prompt = get_generate_output_prompt() + + async def conversational_output_node(state: StateT): + # The appended HumanMessage stays local to this LLM call — only the + # response is returned to state, so the framework instruction never + # enters the persisted conversation history. + messages = [*state.messages, HumanMessage(content=output_prompt)] + config = config_without_streaming(var_child_runnable_config.get(None)) + + try: + response = await llm.ainvoke(messages, config=config) + except Exception as e: + raise_for_provider_http_error(e) + raise + + if not isinstance(response, AIMessage): + raise AgentRuntimeError( + code=AgentRuntimeErrorCode.LLM_INVALID_RESPONSE, + title=f"Structured-output LLM returned {type(response).__name__} invalid response.", + detail=( + "The language model returned an unexpected response type." + "If you are using a BYOM configuration, verify your model deployment.", + ), + category=UiPathErrorCategory.SYSTEM, + ) + + return {"messages": [response]} + + return conversational_output_node diff --git a/src/uipath_langchain/agent/react/router_conversational.py b/src/uipath_langchain/agent/react/router_conversational.py index ea3fdf611..84dafe504 100644 --- a/src/uipath_langchain/agent/react/router_conversational.py +++ b/src/uipath_langchain/agent/react/router_conversational.py @@ -21,31 +21,47 @@ logger = logging.getLogger(__name__) -def create_route_agent_conversational(valid_targets: Container[str] | None = None): +def create_route_agent_conversational( + valid_targets: Container[str] | None = None, + with_generate_output_node: bool = False, +): """Create a routing function for conversational agents. It routes between agent and tool calls until - the agent response has no tool calls, then it routes to the USER_MESSAGE_WAIT node which does an interrupt. + the agent response has no tool calls, then it routes to either the + GENERATE_CONVERSATIONAL_OUTPUT node (when the agent declares custom output + fields) or directly to TERMINATE. Args: - valid_targets: Allowed routing destinations + valid_targets: Allowed routing destinations. + with_generate_output_node: When True, route AGENT-without-tool-calls to the + GENERATE_CONVERSATIONAL_OUTPUT node so the structured output can + be extracted before TERMINATE. When False, route straight to + TERMINATE. Returns: Routing function for LangGraph conditional edges """ def route_agent_conversational( state: AgentGraphState, - ) -> str | Literal[AgentGraphNode.TERMINATE] | Literal[AgentGraphNode.AGENT]: - """Route after agent + ) -> ( + str + | Literal[AgentGraphNode.TERMINATE] + | Literal[AgentGraphNode.AGENT] + | Literal[AgentGraphNode.GENERATE_CONVERSATIONAL_OUTPUT] + ): + """Route after agent. Routing logic: - 3. If tool calls, route to specific tool nodes (return list of tool names) - 4. If no tool calls, route to user message wait node - - Returns: - - str: Tool node name for sequential execution - - AgentGraphNode.USER_MESSAGE_WAIT: When there are no tool calls + - If the latest AIMessage has tool calls + - If pending tools, route to the next pending tool node. + - Otherwise: route to AGENT as all tool calls completed. + - Otherwise: + - If schema declares custom output fields: route to + GENERATE_CONVERSATIONAL_OUTPUT to generate the output fields. + - Otherwise: route straight to TERMINATE. Raises: - AgentNodeRoutingException: When encountering unexpected state (empty messages, non-AIMessage, or excessive completions) + AgentRuntimeError: ROUTING_ERROR when state has no AIMessage, or + when a routed tool name is not in `valid_targets`. """ last_message = find_latest_ai_message(state.messages) if last_message is None: @@ -77,6 +93,10 @@ def route_agent_conversational( return current_tool_name else: - return AgentGraphNode.TERMINATE + return ( + AgentGraphNode.GENERATE_CONVERSATIONAL_OUTPUT + if with_generate_output_node + else AgentGraphNode.TERMINATE + ) return route_agent_conversational diff --git a/src/uipath_langchain/agent/react/terminate_node.py b/src/uipath_langchain/agent/react/terminate_node.py index 092bf72a1..c4ff94077 100644 --- a/src/uipath_langchain/agent/react/terminate_node.py +++ b/src/uipath_langchain/agent/react/terminate_node.py @@ -6,13 +6,19 @@ from langchain_core.messages import AIMessage from pydantic import BaseModel, ValidationError -from uipath.agent.react import END_EXECUTION_TOOL, RAISE_ERROR_TOOL +from uipath.agent.react import ( + END_EXECUTION_TOOL, + RAISE_ERROR_TOOL, + SET_CONVERSATIONAL_OUTPUT_TOOL, +) from uipath.core.chat import UiPathConversationMessageData from uipath.runtime.errors import UiPathErrorCategory from ...runtime.messages import UiPathChatMessagesMapper from ..exceptions import AgentRuntimeError, AgentRuntimeErrorCode +from .constants import UIPATH_CONVERSATIONAL_AGENT_RESPONSE_MESSAGES_FIELD from .types import AgentGraphState +from .utils import has_custom_conversational_output_fields def _handle_end_execution( @@ -50,7 +56,15 @@ def _handle_raise_error(args: dict[str, Any]) -> NoReturn: def _handle_end_conversational( state: AgentGraphState, response_schema: type[BaseModel] | None ) -> dict[str, Any]: - """Handle conversational agent termination by returning converted messages.""" + """Handle conversational agent termination by returning converted messages. + + When the agent's output schema declares custom fields beyond + `uipath__agent_response_messages`, the GENERATE_CONVERSATIONAL_OUTPUT node + runs immediately before TERMINATE and produces an AIMessage carrying a + `set_conversational_output` tool call. The custom field values are + extracted from that tool call's args and merged with the converted + message history for the final output. + """ if state.inner_state.initial_message_count is None: raise AgentRuntimeError( code=AgentRuntimeErrorCode.STATE_ERROR, @@ -67,8 +81,51 @@ def _handle_end_conversational( category=UiPathErrorCategory.SYSTEM, ) + # Handle custom structured output fields if declared in the response schema, which are expected to be + # set via the set_conversational_output tool call in the last AIMessage from the GENERATE_CONVERSATIONAL_OUTPUT node. + has_custom_output = has_custom_conversational_output_fields(response_schema) + custom_output_fields: dict[str, Any] = {} + if has_custom_output: + last_ai_message = state.messages[-1] if state.messages else None + if not isinstance(last_ai_message, AIMessage): + raise AgentRuntimeError( + code=AgentRuntimeErrorCode.ROUTING_ERROR, + title=f"Expected last message to be AIMessage, got {type(last_ai_message).__name__}.", + detail=( + "Custom output fields expected, but last message is not an AIMessage." + ), + category=UiPathErrorCategory.SYSTEM, + ) + set_output_call = next( + ( + tc + for tc in (last_ai_message.tool_calls or []) + if tc["name"] == SET_CONVERSATIONAL_OUTPUT_TOOL.name + ), + None, + ) + if set_output_call is None: + raise AgentRuntimeError( + code=AgentRuntimeErrorCode.OUTPUT_VALIDATION_ERROR, + title="Expected set_conversational_output tool call for last AIMessage.", + detail=( + "Custom output fields expected, but last AIMessage does not contain a " + "set_conversational_output tool call." + ), + category=UiPathErrorCategory.SYSTEM, + ) + custom_output_fields = dict(set_output_call["args"]) + initial_count = state.inner_state.initial_message_count - new_messages = state.messages[initial_count:] + # When custom output fields are present, the last AIMessage is the + # GENERATE_CONVERSATIONAL_OUTPUT node's framework-internal tool call — + # it carries no user-visible content and must not appear in the + # converted `uipath__agent_response_messages` payload. + new_messages = ( + state.messages[initial_count:-1] + if has_custom_output + else state.messages[initial_count:] + ) converted_messages: list[UiPathConversationMessageData] = [] @@ -81,15 +138,26 @@ def _handle_end_conversational( ) ) - output = { - "uipath__agent_response_messages": [ + output: dict[str, Any] = { + **custom_output_fields, + UIPATH_CONVERSATIONAL_AGENT_RESPONSE_MESSAGES_FIELD: [ msg.model_dump(by_alias=True) for msg in converted_messages - ] + ], } - validated = response_schema.model_validate(output) + try: + validated = response_schema.model_validate(output) + except ValidationError as e: + raise AgentRuntimeError( + code=AgentRuntimeErrorCode.OUTPUT_VALIDATION_ERROR, + title="Conversational agent output did not match the expected schema", + detail=( + "The conversational agent's output does not satisfy the configured output schema. " + f"Verify the output, and adjust the schema or the prompt. Details:\n{e}" + ), + category=UiPathErrorCategory.USER, + ) from e # Dump with exclude_none to prevent UiPathConversation... fields with None values from being outputted (e.g. UiPathConversationContentPartData.isTranscript). - # May need to revisit if other output fields are added for conversational agents, where we want nulls outputted. return validated.model_dump(by_alias=True, exclude_none=True) diff --git a/src/uipath_langchain/agent/react/tools/tools.py b/src/uipath_langchain/agent/react/tools/tools.py index 710cdd9fe..cdb722124 100644 --- a/src/uipath_langchain/agent/react/tools/tools.py +++ b/src/uipath_langchain/agent/react/tools/tools.py @@ -7,8 +7,11 @@ from uipath.agent.react import ( END_EXECUTION_TOOL, RAISE_ERROR_TOOL, + SET_CONVERSATIONAL_OUTPUT_TOOL, ) +from ..utils import build_conversational_output_args_schema + def create_end_execution_tool( agent_output_schema: type[BaseModel] | None = None, @@ -41,6 +44,24 @@ async def raise_error_fn(**kwargs: Any) -> dict[str, Any]: ) +def create_set_conversational_output_tool( + agent_output_schema: type[BaseModel], +) -> StructuredTool: + """Called by conversational-agents at the end of the loop when custom-output + fields are declared. Never executed — args are extracted and intercepted during termination.""" + input_schema = build_conversational_output_args_schema(agent_output_schema) + + async def set_conversational_output_fn(**kwargs: Any) -> dict[str, Any]: + return kwargs + + return StructuredTool( + name=SET_CONVERSATIONAL_OUTPUT_TOOL.name, + description=SET_CONVERSATIONAL_OUTPUT_TOOL.description, + args_schema=input_schema, + coroutine=set_conversational_output_fn, + ) + + def create_flow_control_tools( agent_output_schema: type[BaseModel] | None = None, ) -> list[BaseTool]: diff --git a/src/uipath_langchain/agent/react/types.py b/src/uipath_langchain/agent/react/types.py index 2801bc51c..4dda11f6d 100644 --- a/src/uipath_langchain/agent/react/types.py +++ b/src/uipath_langchain/agent/react/types.py @@ -4,7 +4,11 @@ from langchain_core.messages import AnyMessage from langgraph.graph.message import add_messages from pydantic import BaseModel, Field, model_validator -from uipath.agent.react import END_EXECUTION_TOOL, RAISE_ERROR_TOOL +from uipath.agent.react import ( + END_EXECUTION_TOOL, + RAISE_ERROR_TOOL, + SET_CONVERSATIONAL_OUTPUT_TOOL, +) from uipath.platform.attachments import Attachment from uipath_langchain.agent.react.reducers import ( @@ -12,7 +16,11 @@ merge_objects, ) -FLOW_CONTROL_TOOLS = [END_EXECUTION_TOOL.name, RAISE_ERROR_TOOL.name] +FLOW_CONTROL_TOOLS = [ + END_EXECUTION_TOOL.name, + RAISE_ERROR_TOOL.name, + SET_CONVERSATIONAL_OUTPUT_TOOL.name, +] class InnerAgentGraphState(BaseModel): @@ -56,6 +64,7 @@ class AgentGraphNode(StrEnum): AGENT = "agent" LLM = "llm" TOOLS = "tools" + GENERATE_CONVERSATIONAL_OUTPUT = "generate-conversational-output" TERMINATE = "terminate" GUARDED_TERMINATE = "guarded-terminate" MEMORY_RECALL = "memory_recall" diff --git a/src/uipath_langchain/agent/react/utils.py b/src/uipath_langchain/agent/react/utils.py index 26552b0cc..e2fa6d2c7 100644 --- a/src/uipath_langchain/agent/react/utils.py +++ b/src/uipath_langchain/agent/react/utils.py @@ -4,6 +4,7 @@ from langchain_core.messages import AIMessage, AnyMessage, BaseMessage, ToolMessage from pydantic import BaseModel +from pydantic import create_model as pydantic_create_model from uipath.agent.react import END_EXECUTION_TOOL from uipath.runtime.errors import UiPathErrorCategory @@ -11,6 +12,9 @@ AgentRuntimeError, AgentRuntimeErrorCode, ) +from uipath_langchain.agent.react.constants import ( + UIPATH_CONVERSATIONAL_AGENT_RESPONSE_MESSAGES_FIELD, +) from uipath_langchain.agent.react.jsonschema_pydantic_converter import create_model from uipath_langchain.agent.react.types import ( AgentGraphState, @@ -38,6 +42,43 @@ def resolve_output_model( return END_EXECUTION_TOOL.args_schema +def has_custom_conversational_output_fields( + output_schema: type[BaseModel] | None, +) -> bool: + """Return True iff the schema declares fields beyond `uipath__agent_response_messages`. + + Used to decide whether a conversational agent needs the + `GENERATE_CONVERSATIONAL_OUTPUT` node inserted between AGENT and TERMINATE. + """ + if output_schema is None: + return False + return any( + name != UIPATH_CONVERSATIONAL_AGENT_RESPONSE_MESSAGES_FIELD + for name in output_schema.model_fields + ) + + +def build_conversational_output_args_schema( + schema: type[BaseModel], +) -> type[BaseModel]: + """Strip `uipath__agent_response_messages` from the schema. + + That field is populated from message history at termination — not by the + LLM. Stripping it produces the args schema the LLM should actually fill + via `set_conversational_output`. + """ + custom_fields: dict[str, Any] = { + name: (field.annotation, field) + for name, field in schema.model_fields.items() + if name != UIPATH_CONVERSATIONAL_AGENT_RESPONSE_MESSAGES_FIELD + } + return pydantic_create_model( + f"{schema.__name__}SetConversationalOutputArgs", + __base__=BaseModel, + **custom_fields, + ) + + def extract_input_data_from_state( state: BaseModel | dict[str, Any], input_model: type[BaseModel], diff --git a/src/uipath_langchain/agent/tools/internal_tools/analyze_files_tool.py b/src/uipath_langchain/agent/tools/internal_tools/analyze_files_tool.py index 1bf76f387..66fe455a8 100644 --- a/src/uipath_langchain/agent/tools/internal_tools/analyze_files_tool.py +++ b/src/uipath_langchain/agent/tools/internal_tools/analyze_files_tool.py @@ -15,7 +15,6 @@ ) from langchain_core.runnables.config import RunnableConfig, var_child_runnable_config from langchain_core.tools import StructuredTool -from langgraph.constants import TAG_NOSTREAM from opentelemetry import trace as otel_trace from uipath.agent.models.agent import ( AgentInternalToolResourceConfig, @@ -48,7 +47,10 @@ from uipath_langchain.agent.tools.structured_tool_with_argument_properties import ( StructuredToolWithArgumentProperties, ) -from uipath_langchain.agent.tools.utils import sanitize_tool_name +from uipath_langchain.agent.tools.utils import ( + config_without_streaming, + sanitize_tool_name, +) from uipath_langchain.chat.helpers import ( append_content_blocks_to_message, extract_text_content, @@ -139,15 +141,6 @@ def _llm_call_attachments_payload(files: list[FileInfo]) -> str | None: return json.dumps([att.model_dump(by_alias=True) for att in attachments]) -def _config_without_streaming(config: RunnableConfig | None) -> RunnableConfig: - """Tag config with TAG_NOSTREAM so LangGraph's StreamMessagesHandler skips - this LLM call — prevents its response from leaking into the conversation - stream as a visible content part.""" - new_config = cast(RunnableConfig, dict(config) if config else {}) - new_config["tags"] = [*(new_config.get("tags") or []), TAG_NOSTREAM] - return new_config - - def _config_with_llm_call_attachments( config: RunnableConfig | None, files: list[FileInfo] ) -> RunnableConfig | None: @@ -256,8 +249,8 @@ def create_analyze_file_tool( input_model = create_model(resource.input_schema) output_model = create_model(resource.output_schema) - # Disable streaming so for conversational loops, the internal LLM call doesn't leak - # AIMessageChunk events into the graph stream. + # Explicitly disable streaming - for conversational, no streaming is needed as this + # internal tool-call does not produce streamed conversation events. non_streaming_llm = llm.model_copy(update={"disable_streaming": True}) @mockable( @@ -337,7 +330,7 @@ async def tool_fn(**kwargs: Any): cast(AnyMessage, human_message_with_files), ] config = var_child_runnable_config.get(None) - config = _config_without_streaming(config) + config = config_without_streaming(config) config = _config_with_llm_call_attachments(config, files) result = await non_streaming_llm.ainvoke(messages, config=config) diff --git a/src/uipath_langchain/agent/tools/utils.py b/src/uipath_langchain/agent/tools/utils.py index da6609611..e26c87c21 100644 --- a/src/uipath_langchain/agent/tools/utils.py +++ b/src/uipath_langchain/agent/tools/utils.py @@ -1,8 +1,10 @@ """Tool-related utility functions.""" import re -from typing import Any +from typing import Any, cast +from langchain_core.runnables.config import RunnableConfig +from langgraph.constants import TAG_NOSTREAM from uipath.agent.models.agent import TaskTitle, TextBuilderTaskTitle from uipath.agent.utils.text_tokens import build_string_from_tokens @@ -45,6 +47,19 @@ def sanitize_dict_for_serialization(args: dict[str, Any]) -> dict[str, Any]: return converted_args +def config_without_streaming(config: RunnableConfig | None) -> RunnableConfig: + """Return a RunnableConfig with `TAG_NOSTREAM` appended to its tags, so + LangGraph's StreamMessagesHandler skips this LLM call. + + Used for internal LLM calls (e.g. structured-output extraction, + attachment analysis) to prevent their responses from leaking into + the conversation stream for conversational agents. + """ + new_config = cast(RunnableConfig, dict(config) if config else {}) + new_config["tags"] = [*(new_config.get("tags") or []), TAG_NOSTREAM] + return new_config + + def resolve_task_title( task_title: TaskTitle | str | None, agent_input: dict[str, Any], diff --git a/tests/agent/react/test_conversational_output_node.py b/tests/agent/react/test_conversational_output_node.py new file mode 100644 index 000000000..15fae2136 --- /dev/null +++ b/tests/agent/react/test_conversational_output_node.py @@ -0,0 +1,143 @@ +"""Tests for the GENERATE_CONVERSATIONAL_OUTPUT node.""" + +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest +from langchain_core.messages import AIMessage, HumanMessage, SystemMessage +from langgraph.constants import TAG_NOSTREAM +from pydantic import BaseModel, Field +from uipath.agent.react import SET_CONVERSATIONAL_OUTPUT_TOOL + +from uipath_langchain.agent.react.conversational_output_node import ( + create_conversational_output_node, +) +from uipath_langchain.agent.react.types import AgentGraphState, InnerAgentGraphState + + +class _OutputSchema(BaseModel): + uipath__agent_response_messages: list = Field(default_factory=list) + handoff_target: str = "none" + ready_for_handoff: bool = False + + +def _make_state(messages: list[Any]) -> AgentGraphState: + return AgentGraphState( + messages=messages, + inner_state=InnerAgentGraphState(initial_message_count=1), + ) + + +def _make_mock_model() -> tuple[MagicMock, MagicMock, MagicMock]: + """Return a chat model whose + `model_copy(...).bind_tools(...).ainvoke(...)` chain records the messages + and returns a canned AIMessage with the set_conversational_output tool + call. + + The TAG_NOSTREAM tag is supplied at call time via the `config=` kwarg + (mirroring the analyze-files-tool pattern), so `bind_tools` returns the + runnable that's invoked directly. + + Returns (model, non_streaming, bound). + """ + response = AIMessage( + content="", + tool_calls=[ + { + "name": SET_CONVERSATIONAL_OUTPUT_TOOL.name, + "args": {"handoff_target": "billing", "ready_for_handoff": True}, + "id": "call_1", + } + ], + ) + + bound = MagicMock() + bound.ainvoke = AsyncMock(return_value=response) + + non_streaming = MagicMock() + non_streaming.bind_tools = MagicMock(return_value=bound) + + model = MagicMock() + model.model_copy = MagicMock(return_value=non_streaming) + return model, non_streaming, bound + + +class TestCreateConversationalOutputNode: + @pytest.mark.asyncio + async def test_returns_response_with_tool_call(self): + model, _non_streaming, _bound = _make_mock_model() + node = create_conversational_output_node(model, _OutputSchema) + + state = _make_state([SystemMessage(content="sys"), HumanMessage(content="hi")]) + result = await node(state) + + assert len(result["messages"]) == 1 + ai = result["messages"][0] + assert isinstance(ai, AIMessage) + assert ai.tool_calls[0]["name"] == SET_CONVERSATIONAL_OUTPUT_TOOL.name + assert ai.tool_calls[0]["args"]["handoff_target"] == "billing" + + @pytest.mark.asyncio + async def test_appends_human_instruction_for_llm_call(self): + """The framework instruction is appended as a HumanMessage for the LLM + call, but never returned to state.""" + model, _non_streaming, bound = _make_mock_model() + node = create_conversational_output_node(model, _OutputSchema) + + agent_reply = AIMessage(content="here's my reply") + state = _make_state( + [ + SystemMessage(content="sys"), + HumanMessage(content="hi"), + agent_reply, + ] + ) + result = await node(state) + + # The LLM was invoked with state.messages PLUS one extra HumanMessage. + bound.ainvoke.assert_awaited_once() + invoked_messages = bound.ainvoke.await_args.args[0] + assert len(invoked_messages) == len(state.messages) + 1 + assert isinstance(invoked_messages[-1], HumanMessage) + assert "set_conversational_output" in invoked_messages[-1].content + + # The instruction was NOT persisted into the returned messages. + assert len(result["messages"]) == 1 + + @pytest.mark.asyncio + async def test_binds_only_set_conversational_output_tool(self): + model, non_streaming, _bound = _make_mock_model() + create_conversational_output_node(model, _OutputSchema) + + non_streaming.bind_tools.assert_called_once() + tools_arg = non_streaming.bind_tools.call_args.args[0] + assert len(tools_arg) == 1 + assert tools_arg[0].name == SET_CONVERSATIONAL_OUTPUT_TOOL.name + + @pytest.mark.asyncio + async def test_ainvoke_config_includes_tag_nostream(self): + """TAG_NOSTREAM is added to the per-call config (mirroring the + analyze-files-tool pattern) — not via .with_config at construction.""" + model, _non_streaming, bound = _make_mock_model() + node = create_conversational_output_node(model, _OutputSchema) + + state = _make_state([SystemMessage(content="sys"), HumanMessage(content="hi")]) + await node(state) + + bound.ainvoke.assert_awaited_once() + config_arg = bound.ainvoke.await_args.kwargs.get("config") + assert config_arg is not None + assert TAG_NOSTREAM in config_arg.get("tags", []) + + @pytest.mark.asyncio + async def test_disables_streaming_on_internal_llm(self): + """The node copies the model with `disable_streaming=True` so the + underlying provider call doesn't stream — same pattern used by + analyze-files.""" + model, _non_streaming, _bound = _make_mock_model() + create_conversational_output_node(model, _OutputSchema) + + model.model_copy.assert_called_once() + update_arg = model.model_copy.call_args.kwargs.get("update") + assert update_arg is not None + assert update_arg.get("disable_streaming") is True diff --git a/tests/agent/react/test_create_agent.py b/tests/agent/react/test_create_agent.py index fb611c8d6..c62f0c0f9 100644 --- a/tests/agent/react/test_create_agent.py +++ b/tests/agent/react/test_create_agent.py @@ -9,6 +9,7 @@ from langchain_core.runnables.graph import Edge from langchain_core.tools import BaseTool from langgraph.graph import StateGraph +from pydantic import BaseModel, Field from uipath_langchain.agent.react.agent import create_agent from uipath_langchain.agent.react.init_node import create_init_node @@ -293,3 +294,74 @@ def test_conversational_agent_without_tools( mock_route_agent.assert_not_called() mock_route_agent_conversational.assert_called_once() mock_create_flow_control_tools.assert_not_called() + + +class _ConversationalOutputSchemaWithCustomFields(BaseModel): + """A conversational agent's `outputSchema` that exercises the new node.""" + + uipath__agent_response_messages: list = Field(default_factory=list) + handoff_target: str = "none" + ready_for_handoff: bool = False + + +class _ConversationalOutputSchemaNoCustomFields(BaseModel): + uipath__agent_response_messages: list = Field(default_factory=list) + + +class TestCreateAgentGenerateConversationalOutput: + """Topology assertions for the GENERATE_CONVERSATIONAL_OUTPUT node.""" + + @pytest.fixture + def mock_model(self): + return _make_mock_model() + + @pytest.fixture + def mock_tool_a(self): + return _make_mock_tool(mock_tool_a_name) + + @pytest.fixture + def messages(self): + return [SystemMessage(content="You are a helpful assistant.")] + + def test_node_added_when_conversational_with_custom_output( + self, mock_model, mock_tool_a, messages + ): + result: StateGraph[Any] = create_agent( + mock_model, + [mock_tool_a], + messages, + output_schema=_ConversationalOutputSchemaWithCustomFields, + config=AgentGraphConfig(is_conversational=True), + ) + graph = result.compile().get_graph() + assert AgentGraphNode.GENERATE_CONVERSATIONAL_OUTPUT in graph.nodes + assert Edge( + AgentGraphNode.GENERATE_CONVERSATIONAL_OUTPUT, + AgentGraphNode.TERMINATE, + ) in set(graph.edges) + + def test_node_absent_when_conversational_without_custom_output( + self, mock_model, mock_tool_a, messages + ): + result: StateGraph[Any] = create_agent( + mock_model, + [mock_tool_a], + messages, + output_schema=_ConversationalOutputSchemaNoCustomFields, + config=AgentGraphConfig(is_conversational=True), + ) + graph = result.compile().get_graph() + assert AgentGraphNode.GENERATE_CONVERSATIONAL_OUTPUT not in graph.nodes + + def test_node_absent_for_autonomous_agents(self, mock_model, mock_tool_a, messages): + """Even when output_schema has custom fields, autonomous agents do not + get the new node — it is conversational-only.""" + result: StateGraph[Any] = create_agent( + mock_model, + [mock_tool_a], + messages, + output_schema=_ConversationalOutputSchemaWithCustomFields, + config=AgentGraphConfig(is_conversational=False), + ) + graph = result.compile().get_graph() + assert AgentGraphNode.GENERATE_CONVERSATIONAL_OUTPUT not in graph.nodes diff --git a/tests/agent/react/test_flow_control_tools.py b/tests/agent/react/test_flow_control_tools.py new file mode 100644 index 000000000..ef649d664 --- /dev/null +++ b/tests/agent/react/test_flow_control_tools.py @@ -0,0 +1,44 @@ +"""Tests for flow-control tool factories in agent.react.tools.tools.""" + +from pydantic import BaseModel, Field +from uipath.agent.react import SET_CONVERSATIONAL_OUTPUT_TOOL + +from uipath_langchain.agent.react.tools.tools import ( + create_set_conversational_output_tool, +) + + +class _FullOutputSchema(BaseModel): + uipath__agent_response_messages: list = Field(default_factory=list) + handoff_target: str = "none" + ready_for_handoff: bool = False + urgency: str | None = None + + +class TestCreateSetConversationalOutputTool: + def test_tool_name_matches_registry(self): + tool = create_set_conversational_output_tool(_FullOutputSchema) + assert tool.name == SET_CONVERSATIONAL_OUTPUT_TOOL.name + + def test_tool_description_matches_registry(self): + tool = create_set_conversational_output_tool(_FullOutputSchema) + assert tool.description == SET_CONVERSATIONAL_OUTPUT_TOOL.description + + def test_args_schema_strips_response_messages_field(self): + tool = create_set_conversational_output_tool(_FullOutputSchema) + fields = tool.args_schema.model_fields + assert "uipath__agent_response_messages" not in fields + assert "handoff_target" in fields + assert "ready_for_handoff" in fields + assert "urgency" in fields + + def test_args_schema_accepts_partial_payload(self): + """Validates the "N/A"-style placeholder workflow — the schema must + accept partial payloads where optional fields are omitted.""" + tool = create_set_conversational_output_tool(_FullOutputSchema) + validated = tool.args_schema.model_validate( + {"handoff_target": "billing", "ready_for_handoff": True} + ) + assert validated.handoff_target == "billing" + assert validated.ready_for_handoff is True + assert validated.urgency is None diff --git a/tests/agent/react/test_router_conversational.py b/tests/agent/react/test_router_conversational.py index 3a511a9ae..629986ea9 100644 --- a/tests/agent/react/test_router_conversational.py +++ b/tests/agent/react/test_router_conversational.py @@ -295,3 +295,66 @@ def test_default_valid_targets_skips_guard(self): state = AgentGraphState(messages=[HumanMessage(content="query"), ai_message]) assert route_func(state) == "unwired_tool" + + +class TestRouteAgentConversationalCustomOutput: + """Routing for conversational agents with custom output fields. + + When the schema declares fields beyond `uipath__agent_response_messages`, + a GENERATE_CONVERSATIONAL_OUTPUT node sits between AGENT and TERMINATE. + """ + + def test_routes_to_generate_conversational_output_when_custom_output(self): + """No tool calls + has_custom_output=True → GENERATE_CONVERSATIONAL_OUTPUT.""" + route_func = create_route_agent_conversational( + valid_targets=[ + AgentGraphNode.TERMINATE, + AgentGraphNode.GENERATE_CONVERSATIONAL_OUTPUT, + ], + with_generate_output_node=True, + ) + ai_message = AIMessage(content="here is my reply", tool_calls=[]) + state = AgentGraphState(messages=[HumanMessage(content="hi"), ai_message]) + + assert route_func(state) == AgentGraphNode.GENERATE_CONVERSATIONAL_OUTPUT + + def test_routes_to_terminate_when_no_custom_output(self): + """No tool calls + has_custom_output=False → TERMINATE (existing path).""" + route_func = create_route_agent_conversational( + valid_targets=[AgentGraphNode.TERMINATE], + with_generate_output_node=False, + ) + ai_message = AIMessage(content="here is my reply", tool_calls=[]) + state = AgentGraphState(messages=[HumanMessage(content="hi"), ai_message]) + + assert route_func(state) == AgentGraphNode.TERMINATE + + def test_has_custom_output_does_not_affect_tool_routing(self): + """has_custom_output=True must not change tool-routing behavior — the + new branch only fires when the AIMessage has no tool calls.""" + route_func = create_route_agent_conversational( + valid_targets=[ + "real_tool", + AgentGraphNode.TERMINATE, + AgentGraphNode.GENERATE_CONVERSATIONAL_OUTPUT, + ], + with_generate_output_node=True, + ) + ai_message = AIMessage( + content="calling tool", + tool_calls=[{"name": "real_tool", "args": {}, "id": "call_1"}], + ) + state = AgentGraphState(messages=[HumanMessage(content="query"), ai_message]) + + assert route_func(state) == "real_tool" + + def test_has_custom_output_default_false_preserves_legacy_routing(self): + """The new parameter defaults to False so existing callers that don't + opt in keep routing to TERMINATE on AGENT-without-tool-calls.""" + route_func = create_route_agent_conversational( + valid_targets=[AgentGraphNode.TERMINATE] + ) + ai_message = AIMessage(content="reply", tool_calls=[]) + state = AgentGraphState(messages=[HumanMessage(content="hi"), ai_message]) + + assert route_func(state) == AgentGraphNode.TERMINATE diff --git a/tests/agent/react/test_terminate_node.py b/tests/agent/react/test_terminate_node.py index 59cbb5f7a..6746d7064 100644 --- a/tests/agent/react/test_terminate_node.py +++ b/tests/agent/react/test_terminate_node.py @@ -5,7 +5,11 @@ import pytest from langchain_core.messages import AIMessage, HumanMessage, ToolMessage from pydantic import BaseModel -from uipath.agent.react import END_EXECUTION_TOOL, RAISE_ERROR_TOOL +from uipath.agent.react import ( + END_EXECUTION_TOOL, + RAISE_ERROR_TOOL, + SET_CONVERSATIONAL_OUTPUT_TOOL, +) from uipath.core.chat import UiPathConversationMessageData from uipath.runtime.errors import UiPathErrorCategory @@ -173,38 +177,83 @@ class ResponseSchema(BaseModel): assert messages[0]["toolCalls"][0]["name"] == "test_tool" assert messages[0]["toolCalls"][0]["input"] == {"param": "value"} - def test_conversational_ignores_end_execution_tool(self): - """Conversational mode should ignore END_EXECUTION tool calls.""" + def test_conversational_extracts_custom_output_fields(self): + """When the response schema has custom fields, the terminate node + reads them from the last AIMessage's set_conversational_output tool + call and merges them with the response_messages.""" class ResponseSchema(BaseModel): uipath__agent_response_messages: list[UiPathConversationMessageData] + handoff_target: str + ready_for_handoff: bool terminate_node = create_terminate_node( response_schema=ResponseSchema, is_conversational=True ) - ai_message = AIMessage( - content="Done", + + agent_reply = AIMessage(content="Sure, I'll route you to billing.") + set_output_msg = AIMessage( + content="", tool_calls=[ { - "name": END_EXECUTION_TOOL.name, - "args": {"result": "completed"}, + "name": SET_CONVERSATIONAL_OUTPUT_TOOL.name, + "args": { + "handoff_target": "billing", + "ready_for_handoff": True, + }, "id": "call_1", } ], ) state = MockAgentGraphState( - messages=[HumanMessage(content="Initial"), ai_message], + messages=[ + HumanMessage(content="I have a billing issue"), + agent_reply, + set_output_msg, + ], inner_state=MockInnerState(initial_message_count=1), ) - # Should process normally, not treat as special result = terminate_node(state) - assert "uipath__agent_response_messages" in result + assert result["handoff_target"] == "billing" + assert result["ready_for_handoff"] is True + # The conversational reply is preserved; the set_conversational_output + # AIMessage is dropped by the flow-control filter in the converter. messages = result["uipath__agent_response_messages"] assert len(messages) == 1 - assert messages[0]["role"] == "assistant" - assert "Done" in str(messages[0]["contentParts"][0]["data"]) + assert "Sure, I'll route you to billing." in str( + messages[0]["contentParts"][0]["data"] + ) + + def test_conversational_custom_output_raises_when_no_tool_call(self): + """When the schema has custom fields but the last AIMessage doesn't + carry a set_conversational_output tool call, terminate raises a clear + OUTPUT_VALIDATION_ERROR.""" + + class ResponseSchema(BaseModel): + uipath__agent_response_messages: list[UiPathConversationMessageData] + handoff_target: str + + terminate_node = create_terminate_node( + response_schema=ResponseSchema, is_conversational=True + ) + + # Last AIMessage has no tool calls — GENERATE_CONVERSATIONAL_OUTPUT + # failed to call the tool (model behavior failure). + agent_reply = AIMessage(content="Some reply") + state = MockAgentGraphState( + messages=[HumanMessage(content="hi"), agent_reply], + inner_state=MockInnerState(initial_message_count=1), + ) + + with pytest.raises(AgentRuntimeError) as exc_info: + terminate_node(state) + + assert exc_info.value.error_info.code == AgentRuntimeError.full_code( + AgentRuntimeErrorCode.OUTPUT_VALIDATION_ERROR + ) + assert "set_conversational_output" in exc_info.value.error_info.title class TestTerminateNodeNonConversational: diff --git a/tests/agent/react/test_utils.py b/tests/agent/react/test_utils.py index 44ae796ce..9b1ca6ed3 100644 --- a/tests/agent/react/test_utils.py +++ b/tests/agent/react/test_utils.py @@ -2,15 +2,18 @@ import pytest from langchain_core.messages import AIMessage, AnyMessage, HumanMessage, ToolMessage +from pydantic import BaseModel, Field from uipath_langchain.agent.exceptions import ( AgentRuntimeError, AgentRuntimeErrorCode, ) from uipath_langchain.agent.react.utils import ( + build_conversational_output_args_schema, count_consecutive_thinking_messages, extract_current_tool_call_index, find_latest_ai_message, + has_custom_conversational_output_fields, ) @@ -401,3 +404,74 @@ def test_unexpected_message_type(self): assert exc_info.value.error_info.code == AgentRuntimeError.full_code( AgentRuntimeErrorCode.STATE_ERROR ) + + +class TestHasCustomConversationalOutputFields: + """Tests for has_custom_conversational_output_fields helper.""" + + def test_returns_false_for_none_schema(self): + assert has_custom_conversational_output_fields(None) is False + + def test_returns_false_when_only_response_messages_field(self): + class ResponseOnly(BaseModel): + uipath__agent_response_messages: list = Field(default_factory=list) + + assert has_custom_conversational_output_fields(ResponseOnly) is False + + def test_returns_true_when_extra_field_exists(self): + class WithExtras(BaseModel): + uipath__agent_response_messages: list = Field(default_factory=list) + handoff_target: str = "none" + + assert has_custom_conversational_output_fields(WithExtras) is True + + def test_returns_true_when_no_response_messages_field(self): + """Defensive: a schema with only custom fields still requires the new + node (so set_conversational_output can populate them).""" + + class CustomOnly(BaseModel): + web_searched: str = "no" + + assert has_custom_conversational_output_fields(CustomOnly) is True + + +class TestBuildConversationalOutputArgsSchema: + """Tests for build_conversational_output_args_schema helper.""" + + def test_strips_response_messages_field(self): + class FullSchema(BaseModel): + uipath__agent_response_messages: list = Field(default_factory=list) + handoff_target: str = "none" + ready_for_handoff: bool = False + + args_schema = build_conversational_output_args_schema(FullSchema) + assert "uipath__agent_response_messages" not in args_schema.model_fields + assert "handoff_target" in args_schema.model_fields + assert "ready_for_handoff" in args_schema.model_fields + + def test_preserves_field_types_and_defaults(self): + class FullSchema(BaseModel): + uipath__agent_response_messages: list = Field(default_factory=list) + urgency: str = "low" + + args_schema = build_conversational_output_args_schema(FullSchema) + instance = args_schema.model_validate({"urgency": "high"}) + assert instance.urgency == "high" + + def test_no_response_messages_field_passes_through(self): + """When the field isn't present, the result mirrors the input schema.""" + + class CustomOnly(BaseModel): + web_searched: str = "no" + + args_schema = build_conversational_output_args_schema(CustomOnly) + assert "web_searched" in args_schema.model_fields + + def test_generated_schema_name_is_distinct(self): + class FullSchema(BaseModel): + uipath__agent_response_messages: list = Field(default_factory=list) + web_searched: str = "no" + + args_schema = build_conversational_output_args_schema(FullSchema) + assert args_schema.__name__ != FullSchema.__name__ + assert "SetConversationalOutputArgs" in args_schema.__name__ diff --git a/tests/agent/tools/internal_tools/test_analyze_files_tool.py b/tests/agent/tools/internal_tools/test_analyze_files_tool.py index 6b376ec28..ea8f05ade 100644 --- a/tests/agent/tools/internal_tools/test_analyze_files_tool.py +++ b/tests/agent/tools/internal_tools/test_analyze_files_tool.py @@ -27,7 +27,6 @@ ANALYZE_FILES_SYSTEM_MESSAGE, LLM_CALL_ATTACHMENTS_METADATA_KEY, _config_with_llm_call_attachments, - _config_without_streaming, _is_pii_scope_for_files, _resolve_job_attachment_arguments, create_analyze_file_tool, @@ -1017,31 +1016,6 @@ def test_is_case_sensitive(self) -> None: ) -class TestConfigWithoutStreaming: - """Tests for _config_without_streaming — ensures TAG_NOSTREAM is injected.""" - - def test_adds_nostream_tag_to_empty_config(self) -> None: - result = _config_without_streaming(None) - assert TAG_NOSTREAM in result["tags"] - - def test_adds_nostream_tag_to_existing_config(self) -> None: - config: RunnableConfig = {"tags": ["existing_tag"]} - result = _config_without_streaming(config) - assert "existing_tag" in result["tags"] - assert TAG_NOSTREAM in result["tags"] - - def test_preserves_other_config_keys(self) -> None: - config: RunnableConfig = {"metadata": {"key": "value"}, "tags": ["t"]} - result = _config_without_streaming(config) - assert result["metadata"] == {"key": "value"} - assert TAG_NOSTREAM in result["tags"] - - def test_handles_config_without_tags(self) -> None: - config: RunnableConfig = {"metadata": {"key": "value"}} - result = _config_without_streaming(config) - assert result["tags"] == [TAG_NOSTREAM] - - class TestConfigWithLlmCallAttachments: """The attachments payload travels to the llmCall span via langchain config metadata.""" diff --git a/tests/agent/tools/test_utils.py b/tests/agent/tools/test_utils.py index 02418513a..f98153f24 100644 --- a/tests/agent/tools/test_utils.py +++ b/tests/agent/tools/test_utils.py @@ -1,8 +1,11 @@ """Tests for tools/utils.py module.""" +from langchain_core.runnables.config import RunnableConfig +from langgraph.constants import TAG_NOSTREAM from pydantic import BaseModel from uipath_langchain.agent.tools.utils import ( + config_without_streaming, sanitize_dict_for_serialization, sanitize_tool_name, ) @@ -122,3 +125,29 @@ def __init__(self): result = sanitize_dict_for_serialization(input_dict) assert result["obj"] is obj + + +class TestConfigWithoutStreaming: + """Ensures TAG_NOSTREAM is injected into the runnable config so LangGraph + skips the LLM call's tokens in the messages stream.""" + + def test_adds_nostream_tag_to_empty_config(self) -> None: + result = config_without_streaming(None) + assert TAG_NOSTREAM in result["tags"] + + def test_adds_nostream_tag_to_existing_config(self) -> None: + config: RunnableConfig = {"tags": ["existing_tag"]} + result = config_without_streaming(config) + assert "existing_tag" in result["tags"] + assert TAG_NOSTREAM in result["tags"] + + def test_preserves_other_config_keys(self) -> None: + config: RunnableConfig = {"metadata": {"key": "value"}, "tags": ["t"]} + result = config_without_streaming(config) + assert result["metadata"] == {"key": "value"} + assert TAG_NOSTREAM in result["tags"] + + def test_handles_config_without_tags(self) -> None: + config: RunnableConfig = {"metadata": {"key": "value"}} + result = config_without_streaming(config) + assert result["tags"] == [TAG_NOSTREAM] From e69fa94c703c739ee2de5e45c43af71bd6804dbc Mon Sep 17 00:00:00 2001 From: Maxwell Du <60411452+maxduu@users.noreply.github.com> Date: Wed, 1 Jul 2026 18:01:07 -0400 Subject: [PATCH 2/4] fix: ai message logic --- src/uipath_langchain/agent/react/agent.py | 1 + .../agent/react/router_conversational.py | 2 +- .../agent/react/terminate_node.py | 48 ++++++----------- src/uipath_langchain/agent/react/types.py | 4 ++ tests/agent/react/test_create_agent.py | 52 +++++++++++++++++-- tests/agent/react/test_terminate_node.py | 45 +++++++++++++--- 6 files changed, 107 insertions(+), 45 deletions(-) diff --git a/src/uipath_langchain/agent/react/agent.py b/src/uipath_langchain/agent/react/agent.py index 7e2342521..36c0b1741 100644 --- a/src/uipath_langchain/agent/react/agent.py +++ b/src/uipath_langchain/agent/react/agent.py @@ -134,6 +134,7 @@ def create_agent( with_conversational_output_node = ( config.is_conversational + and config.conversational_outputs_enabled and has_custom_conversational_output_fields(output_schema) ) diff --git a/src/uipath_langchain/agent/react/router_conversational.py b/src/uipath_langchain/agent/react/router_conversational.py index 84dafe504..54ef48937 100644 --- a/src/uipath_langchain/agent/react/router_conversational.py +++ b/src/uipath_langchain/agent/react/router_conversational.py @@ -55,7 +55,7 @@ def route_agent_conversational( - If pending tools, route to the next pending tool node. - Otherwise: route to AGENT as all tool calls completed. - Otherwise: - - If schema declares custom output fields: route to + - If schema declares custom output fields: route to GENERATE_CONVERSATIONAL_OUTPUT to generate the output fields. - Otherwise: route straight to TERMINATE. diff --git a/src/uipath_langchain/agent/react/terminate_node.py b/src/uipath_langchain/agent/react/terminate_node.py index c4ff94077..80a784c60 100644 --- a/src/uipath_langchain/agent/react/terminate_node.py +++ b/src/uipath_langchain/agent/react/terminate_node.py @@ -18,7 +18,6 @@ from ..exceptions import AgentRuntimeError, AgentRuntimeErrorCode from .constants import UIPATH_CONVERSATIONAL_AGENT_RESPONSE_MESSAGES_FIELD from .types import AgentGraphState -from .utils import has_custom_conversational_output_fields def _handle_end_execution( @@ -81,22 +80,11 @@ def _handle_end_conversational( category=UiPathErrorCategory.SYSTEM, ) - # Handle custom structured output fields if declared in the response schema, which are expected to be - # set via the set_conversational_output tool call in the last AIMessage from the GENERATE_CONVERSATIONAL_OUTPUT node. - has_custom_output = has_custom_conversational_output_fields(response_schema) - custom_output_fields: dict[str, Any] = {} - if has_custom_output: - last_ai_message = state.messages[-1] if state.messages else None - if not isinstance(last_ai_message, AIMessage): - raise AgentRuntimeError( - code=AgentRuntimeErrorCode.ROUTING_ERROR, - title=f"Expected last message to be AIMessage, got {type(last_ai_message).__name__}.", - detail=( - "Custom output fields expected, but last message is not an AIMessage." - ), - category=UiPathErrorCategory.SYSTEM, - ) - set_output_call = next( + # Extract structured output fields from the last AIMessage's + # `set_conversational_output` tool call, when present. + last_ai_message = state.messages[-1] if state.messages else None + set_output_call = ( + next( ( tc for tc in (last_ai_message.tool_calls or []) @@ -104,26 +92,20 @@ def _handle_end_conversational( ), None, ) - if set_output_call is None: - raise AgentRuntimeError( - code=AgentRuntimeErrorCode.OUTPUT_VALIDATION_ERROR, - title="Expected set_conversational_output tool call for last AIMessage.", - detail=( - "Custom output fields expected, but last AIMessage does not contain a " - "set_conversational_output tool call." - ), - category=UiPathErrorCategory.SYSTEM, - ) - custom_output_fields = dict(set_output_call["args"]) + if isinstance(last_ai_message, AIMessage) + else None + ) + custom_output_fields: dict[str, Any] = ( + dict(set_output_call["args"]) if set_output_call is not None else {} + ) initial_count = state.inner_state.initial_message_count - # When custom output fields are present, the last AIMessage is the - # GENERATE_CONVERSATIONAL_OUTPUT node's framework-internal tool call — - # it carries no user-visible content and must not appear in the - # converted `uipath__agent_response_messages` payload. + # Drop the GENERATE_CONVERSATIONAL_OUTPUT AIMessage from the converted + # response payload — it carries no user-visible content, only the + # framework-internal set_conversational_output tool call. new_messages = ( state.messages[initial_count:-1] - if has_custom_output + if set_output_call is not None else state.messages[initial_count:] ) diff --git a/src/uipath_langchain/agent/react/types.py b/src/uipath_langchain/agent/react/types.py index 4dda11f6d..8f840060e 100644 --- a/src/uipath_langchain/agent/react/types.py +++ b/src/uipath_langchain/agent/react/types.py @@ -132,3 +132,7 @@ class AgentGraphConfig(BaseModel): default=False, description="If set, the LLM will guarantee schema validation of the tool calls.", ) + conversational_outputs_enabled: bool = Field( + default=False, + description="Whether generation of custom-defined outputs for conversational agents is enabled.", + ) diff --git a/tests/agent/react/test_create_agent.py b/tests/agent/react/test_create_agent.py index c62f0c0f9..df73c8990 100644 --- a/tests/agent/react/test_create_agent.py +++ b/tests/agent/react/test_create_agent.py @@ -323,7 +323,7 @@ def mock_tool_a(self): def messages(self): return [SystemMessage(content="You are a helpful assistant.")] - def test_node_added_when_conversational_with_custom_output( + def test_node_added_when_conversational_with_custom_output_and_storage_version_supported( self, mock_model, mock_tool_a, messages ): result: StateGraph[Any] = create_agent( @@ -331,7 +331,10 @@ def test_node_added_when_conversational_with_custom_output( [mock_tool_a], messages, output_schema=_ConversationalOutputSchemaWithCustomFields, - config=AgentGraphConfig(is_conversational=True), + config=AgentGraphConfig( + is_conversational=True, + conversational_outputs_enabled=True, + ), ) graph = result.compile().get_graph() assert AgentGraphNode.GENERATE_CONVERSATIONAL_OUTPUT in graph.nodes @@ -348,7 +351,10 @@ def test_node_absent_when_conversational_without_custom_output( [mock_tool_a], messages, output_schema=_ConversationalOutputSchemaNoCustomFields, - config=AgentGraphConfig(is_conversational=True), + config=AgentGraphConfig( + is_conversational=True, + conversational_outputs_enabled=True, + ), ) graph = result.compile().get_graph() assert AgentGraphNode.GENERATE_CONVERSATIONAL_OUTPUT not in graph.nodes @@ -361,7 +367,45 @@ def test_node_absent_for_autonomous_agents(self, mock_model, mock_tool_a, messag [mock_tool_a], messages, output_schema=_ConversationalOutputSchemaWithCustomFields, - config=AgentGraphConfig(is_conversational=False), + config=AgentGraphConfig( + is_conversational=False, + conversational_outputs_enabled=True, + ), + ) + graph = result.compile().get_graph() + assert AgentGraphNode.GENERATE_CONVERSATIONAL_OUTPUT not in graph.nodes + + def test_node_absent_when_storage_version_unsupported( + self, mock_model, mock_tool_a, messages + ): + """Legacy agents (storageVersion < 51.0.0) shipped a default + outputSchema that must not trigger the new node — even when the + schema has fields beyond `uipath__agent_response_messages`.""" + result: StateGraph[Any] = create_agent( + mock_model, + [mock_tool_a], + messages, + output_schema=_ConversationalOutputSchemaWithCustomFields, + config=AgentGraphConfig( + is_conversational=True, + conversational_outputs_enabled=False, + ), + ) + graph = result.compile().get_graph() + assert AgentGraphNode.GENERATE_CONVERSATIONAL_OUTPUT not in graph.nodes + + def test_node_absent_when_storage_version_flag_defaults_to_false( + self, mock_model, mock_tool_a, messages + ): + """Callers that don't opt in via the new flag get legacy behavior — + the field defaults to False so direct create_agent() usage without + an explicit opt-in is safe.""" + result: StateGraph[Any] = create_agent( + mock_model, + [mock_tool_a], + messages, + output_schema=_ConversationalOutputSchemaWithCustomFields, + config=AgentGraphConfig(is_conversational=True), ) graph = result.compile().get_graph() assert AgentGraphNode.GENERATE_CONVERSATIONAL_OUTPUT not in graph.nodes diff --git a/tests/agent/react/test_terminate_node.py b/tests/agent/react/test_terminate_node.py index 6746d7064..6aa16cf40 100644 --- a/tests/agent/react/test_terminate_node.py +++ b/tests/agent/react/test_terminate_node.py @@ -226,10 +226,14 @@ class ResponseSchema(BaseModel): messages[0]["contentParts"][0]["data"] ) - def test_conversational_custom_output_raises_when_no_tool_call(self): - """When the schema has custom fields but the last AIMessage doesn't - carry a set_conversational_output tool call, terminate raises a clear - OUTPUT_VALIDATION_ERROR.""" + def test_conversational_missing_set_output_call_falls_through_to_schema_validation( + self, + ) -> None: + """When the schema has required custom fields but the last AIMessage + doesn't carry a set_conversational_output tool call, terminate does + NOT raise a routing error — it best-effort passes an empty custom + dict through, and the surrounding schema validation raises a clear + OUTPUT_VALIDATION_ERROR referencing the missing required field.""" class ResponseSchema(BaseModel): uipath__agent_response_messages: list[UiPathConversationMessageData] @@ -239,8 +243,8 @@ class ResponseSchema(BaseModel): response_schema=ResponseSchema, is_conversational=True ) - # Last AIMessage has no tool calls — GENERATE_CONVERSATIONAL_OUTPUT - # failed to call the tool (model behavior failure). + # Last AIMessage has no tool calls — feature could be off, or the + # extraction model didn't call the tool. agent_reply = AIMessage(content="Some reply") state = MockAgentGraphState( messages=[HumanMessage(content="hi"), agent_reply], @@ -250,10 +254,37 @@ class ResponseSchema(BaseModel): with pytest.raises(AgentRuntimeError) as exc_info: terminate_node(state) + # The error is a schema-validation error surfaced by Pydantic, not + # a routing error — the message references the missing required + # field on the response schema, not `set_conversational_output`. assert exc_info.value.error_info.code == AgentRuntimeError.full_code( AgentRuntimeErrorCode.OUTPUT_VALIDATION_ERROR ) - assert "set_conversational_output" in exc_info.value.error_info.title + assert "handoff_target" in exc_info.value.error_info.detail + + def test_conversational_no_custom_fields_no_set_output_call_succeeds(self): + """When the schema has no custom fields and no set_conversational_output + tool call was made, terminate succeeds without extracting anything.""" + + class ResponseSchema(BaseModel): + uipath__agent_response_messages: list[UiPathConversationMessageData] + + terminate_node = create_terminate_node( + response_schema=ResponseSchema, is_conversational=True + ) + agent_reply = AIMessage(content="Some reply") + state = MockAgentGraphState( + messages=[HumanMessage(content="hi"), agent_reply], + inner_state=MockInnerState(initial_message_count=1), + ) + + result = terminate_node(state) + + # The final reply is preserved — no last-message slicing since no + # set_conversational_output call was found. + messages = result["uipath__agent_response_messages"] + assert len(messages) == 1 + assert "Some reply" in str(messages[0]["contentParts"][0]["data"]) class TestTerminateNodeNonConversational: From bd654a42aee37bd0ebb003513f6a00dfd08c4bbd Mon Sep 17 00:00:00 2001 From: Maxwell Du <60411452+maxduu@users.noreply.github.com> Date: Thu, 2 Jul 2026 13:18:22 -0400 Subject: [PATCH 3/4] fix: simplify logic for enabled outputs --- src/uipath_langchain/agent/react/agent.py | 1 - .../agent/react/terminate_node.py | 2 +- src/uipath_langchain/agent/react/types.py | 4 -- tests/agent/react/test_create_agent.py | 52 ++----------------- 4 files changed, 5 insertions(+), 54 deletions(-) diff --git a/src/uipath_langchain/agent/react/agent.py b/src/uipath_langchain/agent/react/agent.py index 36c0b1741..7e2342521 100644 --- a/src/uipath_langchain/agent/react/agent.py +++ b/src/uipath_langchain/agent/react/agent.py @@ -134,7 +134,6 @@ def create_agent( with_conversational_output_node = ( config.is_conversational - and config.conversational_outputs_enabled and has_custom_conversational_output_fields(output_schema) ) diff --git a/src/uipath_langchain/agent/react/terminate_node.py b/src/uipath_langchain/agent/react/terminate_node.py index 80a784c60..d5365b7a9 100644 --- a/src/uipath_langchain/agent/react/terminate_node.py +++ b/src/uipath_langchain/agent/react/terminate_node.py @@ -80,7 +80,7 @@ def _handle_end_conversational( category=UiPathErrorCategory.SYSTEM, ) - # Extract structured output fields from the last AIMessage's + # Extract structured output fields from the last AIMessage's # `set_conversational_output` tool call, when present. last_ai_message = state.messages[-1] if state.messages else None set_output_call = ( diff --git a/src/uipath_langchain/agent/react/types.py b/src/uipath_langchain/agent/react/types.py index 8f840060e..4dda11f6d 100644 --- a/src/uipath_langchain/agent/react/types.py +++ b/src/uipath_langchain/agent/react/types.py @@ -132,7 +132,3 @@ class AgentGraphConfig(BaseModel): default=False, description="If set, the LLM will guarantee schema validation of the tool calls.", ) - conversational_outputs_enabled: bool = Field( - default=False, - description="Whether generation of custom-defined outputs for conversational agents is enabled.", - ) diff --git a/tests/agent/react/test_create_agent.py b/tests/agent/react/test_create_agent.py index df73c8990..c62f0c0f9 100644 --- a/tests/agent/react/test_create_agent.py +++ b/tests/agent/react/test_create_agent.py @@ -323,7 +323,7 @@ def mock_tool_a(self): def messages(self): return [SystemMessage(content="You are a helpful assistant.")] - def test_node_added_when_conversational_with_custom_output_and_storage_version_supported( + def test_node_added_when_conversational_with_custom_output( self, mock_model, mock_tool_a, messages ): result: StateGraph[Any] = create_agent( @@ -331,10 +331,7 @@ def test_node_added_when_conversational_with_custom_output_and_storage_version_s [mock_tool_a], messages, output_schema=_ConversationalOutputSchemaWithCustomFields, - config=AgentGraphConfig( - is_conversational=True, - conversational_outputs_enabled=True, - ), + config=AgentGraphConfig(is_conversational=True), ) graph = result.compile().get_graph() assert AgentGraphNode.GENERATE_CONVERSATIONAL_OUTPUT in graph.nodes @@ -351,10 +348,7 @@ def test_node_absent_when_conversational_without_custom_output( [mock_tool_a], messages, output_schema=_ConversationalOutputSchemaNoCustomFields, - config=AgentGraphConfig( - is_conversational=True, - conversational_outputs_enabled=True, - ), + config=AgentGraphConfig(is_conversational=True), ) graph = result.compile().get_graph() assert AgentGraphNode.GENERATE_CONVERSATIONAL_OUTPUT not in graph.nodes @@ -367,45 +361,7 @@ def test_node_absent_for_autonomous_agents(self, mock_model, mock_tool_a, messag [mock_tool_a], messages, output_schema=_ConversationalOutputSchemaWithCustomFields, - config=AgentGraphConfig( - is_conversational=False, - conversational_outputs_enabled=True, - ), - ) - graph = result.compile().get_graph() - assert AgentGraphNode.GENERATE_CONVERSATIONAL_OUTPUT not in graph.nodes - - def test_node_absent_when_storage_version_unsupported( - self, mock_model, mock_tool_a, messages - ): - """Legacy agents (storageVersion < 51.0.0) shipped a default - outputSchema that must not trigger the new node — even when the - schema has fields beyond `uipath__agent_response_messages`.""" - result: StateGraph[Any] = create_agent( - mock_model, - [mock_tool_a], - messages, - output_schema=_ConversationalOutputSchemaWithCustomFields, - config=AgentGraphConfig( - is_conversational=True, - conversational_outputs_enabled=False, - ), - ) - graph = result.compile().get_graph() - assert AgentGraphNode.GENERATE_CONVERSATIONAL_OUTPUT not in graph.nodes - - def test_node_absent_when_storage_version_flag_defaults_to_false( - self, mock_model, mock_tool_a, messages - ): - """Callers that don't opt in via the new flag get legacy behavior — - the field defaults to False so direct create_agent() usage without - an explicit opt-in is safe.""" - result: StateGraph[Any] = create_agent( - mock_model, - [mock_tool_a], - messages, - output_schema=_ConversationalOutputSchemaWithCustomFields, - config=AgentGraphConfig(is_conversational=True), + config=AgentGraphConfig(is_conversational=False), ) graph = result.compile().get_graph() assert AgentGraphNode.GENERATE_CONVERSATIONAL_OUTPUT not in graph.nodes From d9e119a1dca3bcb1aa9c2b3eed1c23611c1f601c Mon Sep 17 00:00:00 2001 From: Maxwell Du <60411452+maxduu@users.noreply.github.com> Date: Thu, 2 Jul 2026 13:22:28 -0400 Subject: [PATCH 4/4] chore: update comment --- .../agent/react/conversational_output_node.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/src/uipath_langchain/agent/react/conversational_output_node.py b/src/uipath_langchain/agent/react/conversational_output_node.py index e167a67ce..abdadc55f 100644 --- a/src/uipath_langchain/agent/react/conversational_output_node.py +++ b/src/uipath_langchain/agent/react/conversational_output_node.py @@ -3,13 +3,7 @@ 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 reliably extract the structured -output for the turn — decoupling conversational quality from schema -compliance. - -The LLM call is tagged with `TAG_NOSTREAM` so its tokens / events never -reach the chat-UI message stream. TERMINATE then reads the tool call's -args from `state.messages[-1]`. +tool bound and `tool_choice="any"` to extract the structured output for the turn. """ from typing import TypeVar