From 5ee2ee0e9dc285e8925eb37276472f00b09aedd9 Mon Sep 17 00:00:00 2001 From: Cristian Pufu Date: Sun, 1 Mar 2026 16:59:18 +0200 Subject: [PATCH] fix: pydantic ai proper schema and events payload --- README.md | 4 +- packages/uipath-pydantic-ai/pyproject.toml | 2 +- .../samples/graph-flow/README.md | 2 +- .../samples/multi-agent/README.md | 2 +- .../samples/programmatic-handoff/README.md | 2 +- .../samples/quickstart-agent/README.md | 4 +- .../samples/quickstart-agent/main.py | 4 +- .../samples/quickstart-agent/pyproject.toml | 2 +- .../src/uipath_pydantic_ai/chat/openai.py | 10 +- .../src/uipath_pydantic_ai/runtime/factory.py | 1 + .../src/uipath_pydantic_ai/runtime/runtime.py | 222 +++++--- .../src/uipath_pydantic_ai/runtime/schema.py | 110 ++-- .../uipath-pydantic-ai/tests/test_runtime.py | 492 +++++++++++++++++- .../uipath-pydantic-ai/tests/test_schema.py | 128 ++++- packages/uipath-pydantic-ai/uv.lock | 2 +- 15 files changed, 855 insertions(+), 132 deletions(-) diff --git a/README.md b/README.md index 6f1fa8b4..b64977e0 100644 --- a/README.md +++ b/README.md @@ -13,6 +13,7 @@ All packages extend the [UiPath Python SDK](https://github.com/UiPath/uipath-pyt | [LlamaIndex](https://www.llamaindex.ai/) | [![PyPI](https://img.shields.io/pypi/v/uipath-llamaindex)](https://pypi.org/project/uipath-llamaindex/) | [![Downloads](https://img.shields.io/pypi/dm/uipath-llamaindex.svg)](https://pypi.org/project/uipath-llamaindex/) | [README](packages/uipath-llamaindex/README.md) · [Docs](https://uipath.github.io/uipath-python/llamaindex/quick_start/) · [Samples](packages/uipath-llamaindex/samples/) | | [Microsoft Agent Framework](https://github.com/microsoft/agent-framework) | [![PyPI](https://img.shields.io/pypi/v/uipath-agent-framework)](https://pypi.org/project/uipath-agent-framework/) | [![Downloads](https://img.shields.io/pypi/dm/uipath-agent-framework.svg)](https://pypi.org/project/uipath-agent-framework/) | [README](packages/uipath-agent-framework/README.md) · [Samples](packages/uipath-agent-framework/samples/) | | [OpenAI Agents](https://github.com/openai/openai-agents-python) | [![PyPI](https://img.shields.io/pypi/v/uipath-openai-agents)](https://pypi.org/project/uipath-openai-agents/) | [![Downloads](https://img.shields.io/pypi/dm/uipath-openai-agents.svg)](https://pypi.org/project/uipath-openai-agents/) | [README](packages/uipath-openai-agents/README.md) · [Docs](https://uipath.github.io/uipath-python/openai-agents/quick_start/) · [Samples](packages/uipath-openai-agents/samples/) | +| [PydanticAI](https://github.com/pydantic/pydantic-ai) | [![PyPI](https://img.shields.io/pypi/v/uipath-pydantic-ai)](https://pypi.org/project/uipath-pydantic-ai/) | [![Downloads](https://img.shields.io/pypi/dm/uipath-pydantic-ai.svg)](https://pypi.org/project/uipath-pydantic-ai/) | [README](packages/uipath-pydantic-ai/README.md) · [Samples](packages/uipath-pydantic-ai/samples/) | ## Structure @@ -25,7 +26,8 @@ uipath-integrations-python/ ├── uipath-llamaindex/ # LlamaIndex runtime ├── uipath-openai-agents/ # OpenAI Agents runtime ├── uipath-google-adk/ # Google ADK runtime - └── uipath-agent-framework/ # Microsoft Agent Framework runtime + ├── uipath-agent-framework/ # Microsoft Agent Framework runtime + └── uipath-pydantic-ai/ # PydanticAI runtime ``` ## Development diff --git a/packages/uipath-pydantic-ai/pyproject.toml b/packages/uipath-pydantic-ai/pyproject.toml index 625050c6..0a844c0d 100644 --- a/packages/uipath-pydantic-ai/pyproject.toml +++ b/packages/uipath-pydantic-ai/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "uipath-pydantic-ai" -version = "0.0.1" +version = "0.0.2" description = "Python SDK that enables developers to build and deploy PydanticAI agents to the UiPath Cloud Platform" readme = "README.md" requires-python = ">=3.11" diff --git a/packages/uipath-pydantic-ai/samples/graph-flow/README.md b/packages/uipath-pydantic-ai/samples/graph-flow/README.md index 596ad244..e7bc786d 100644 --- a/packages/uipath-pydantic-ai/samples/graph-flow/README.md +++ b/packages/uipath-pydantic-ai/samples/graph-flow/README.md @@ -34,7 +34,7 @@ flowchart TB ## Run ``` -uipath run agent '{"messages": "Write a blog post about the future of AI agents"}' +uipath run agent '{"messages": [{"contentParts": [{"data": {"inline": "Write a blog post about the future of AI agents"}}], "role": "user"}]}' ``` ## Debug diff --git a/packages/uipath-pydantic-ai/samples/multi-agent/README.md b/packages/uipath-pydantic-ai/samples/multi-agent/README.md index af0fb46c..00d55a35 100644 --- a/packages/uipath-pydantic-ai/samples/multi-agent/README.md +++ b/packages/uipath-pydantic-ai/samples/multi-agent/README.md @@ -19,7 +19,7 @@ flowchart TB ## Run ``` -uipath run agent '{"messages": "Write a report about machine learning"}' +uipath run agent '{"messages": [{"contentParts": [{"data": {"inline": "Write a report about machine learning"}}], "role": "user"}]}' ``` ## Debug diff --git a/packages/uipath-pydantic-ai/samples/programmatic-handoff/README.md b/packages/uipath-pydantic-ai/samples/programmatic-handoff/README.md index 588a5955..dd3e2db0 100644 --- a/packages/uipath-pydantic-ai/samples/programmatic-handoff/README.md +++ b/packages/uipath-pydantic-ai/samples/programmatic-handoff/README.md @@ -24,7 +24,7 @@ flowchart TB ## Run ``` -uipath run agent '{"messages": "I was charged twice for my subscription last month"}' +uipath run agent '{"messages": [{"contentParts": [{"data": {"inline": "I was charged twice for my subscription last month"}}], "role": "user"}]}' ``` ## Debug diff --git a/packages/uipath-pydantic-ai/samples/quickstart-agent/README.md b/packages/uipath-pydantic-ai/samples/quickstart-agent/README.md index 82ef8ca3..8709f2ef 100644 --- a/packages/uipath-pydantic-ai/samples/quickstart-agent/README.md +++ b/packages/uipath-pydantic-ai/samples/quickstart-agent/README.md @@ -19,11 +19,11 @@ flowchart TB ## Run ``` -uipath run agent '{"messages": "What is the weather in San Francisco?"}' +uipath run agent '{"messages": [{"contentParts": [{"data": {"inline": "What is the weather in San Francisco?"}}], "role": "user"}]}' ``` ## Debug ``` -uipath dev web +uipath dev ``` diff --git a/packages/uipath-pydantic-ai/samples/quickstart-agent/main.py b/packages/uipath-pydantic-ai/samples/quickstart-agent/main.py index 07278bf3..2d7f1901 100644 --- a/packages/uipath-pydantic-ai/samples/quickstart-agent/main.py +++ b/packages/uipath-pydantic-ai/samples/quickstart-agent/main.py @@ -1,10 +1,10 @@ import httpx -from pydantic_ai import Agent +from pydantic_ai import Agent, RunContext from uipath_pydantic_ai.chat import UiPathChatOpenAI -def get_weather(ctx, location: str) -> str: +def get_weather(ctx: RunContext[None], location: str) -> str: """Get the current weather for a location using the Open-Meteo API. Args: diff --git a/packages/uipath-pydantic-ai/samples/quickstart-agent/pyproject.toml b/packages/uipath-pydantic-ai/samples/quickstart-agent/pyproject.toml index 260fed9f..15b0701c 100644 --- a/packages/uipath-pydantic-ai/samples/quickstart-agent/pyproject.toml +++ b/packages/uipath-pydantic-ai/samples/quickstart-agent/pyproject.toml @@ -8,7 +8,7 @@ requires-python = ">=3.11" dependencies = [ "uipath", "uipath-pydantic-ai", - "pydantic-ai>=0.1.0", + "pydantic-ai>=1.63.0", ] [dependency-groups] diff --git a/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/chat/openai.py b/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/chat/openai.py index cbd81f0f..54fb39db 100644 --- a/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/chat/openai.py +++ b/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/chat/openai.py @@ -5,7 +5,7 @@ import httpx from openai import AsyncOpenAI, OpenAI -from pydantic_ai.models.openai import OpenAIModel +from pydantic_ai.models.openai import OpenAIChatModel from pydantic_ai.providers.openai import OpenAIProvider from uipath._utils._ssl_context import get_httpx_client_kwargs from uipath.utils import EndpointManager @@ -74,7 +74,7 @@ class UiPathChatOpenAI: This client wraps the OpenAI SDK and configures it to use UiPath's LLM Gateway endpoints with proper authentication and headers. - Returns PydanticAI-compatible OpenAIModel instances. + Returns PydanticAI-compatible OpenAIChatModel instances. Example: ```python @@ -173,7 +173,7 @@ def __init__( ) # Create PydanticAI-compatible model - self._model = OpenAIModel( + self._model = OpenAIChatModel( self._model_name, provider=OpenAIProvider(openai_client=self._async_client), ) @@ -225,8 +225,8 @@ def _build_base_url(self) -> str: raise ValueError("UIPATH_URL environment variable is required") @property - def model(self) -> OpenAIModel: - """Get the PydanticAI-compatible OpenAIModel.""" + def model(self) -> OpenAIChatModel: + """Get the PydanticAI-compatible OpenAIChatModel.""" return self._model @property diff --git a/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/runtime/factory.py b/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/runtime/factory.py index 12063e10..c415baef 100644 --- a/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/runtime/factory.py +++ b/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/runtime/factory.py @@ -58,6 +58,7 @@ def _setup_instrumentation(self) -> None: self.context.trace_manager.add_span_processor( OpenInferenceSpanProcessor() ) + Agent.instrument_all() except Exception: pass diff --git a/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/runtime/runtime.py b/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/runtime/runtime.py index 638c6a6c..a47d212f 100644 --- a/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/runtime/runtime.py +++ b/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/runtime/runtime.py @@ -63,84 +63,146 @@ async def stream( input: dict[str, Any] | None = None, options: UiPathStreamOptions | None = None, ) -> AsyncGenerator[UiPathRuntimeEvent, None]: - """Stream agent execution events in real-time. + """Stream agent execution events in real-time.""" + from pydantic_graph import End - Uses agent.iter() for node-level iteration, emitting - UiPathRuntimeStateEvent with node_name and phase (STARTED/COMPLETED) - to drive graph node highlighting in the UI. - - Node mapping to graph nodes (from schema.py): - - ModelRequestNode -> agent node (STARTED/COMPLETED) - - CallToolsNode -> {agent}_tools node (STARTED/COMPLETED) - """ try: user_prompt, deps = self._prepare_input(input) agent_name = self.agent.name or "agent" - tools_node = f"{agent_name}_tools" + tools_node_name = f"{agent_name}_tools" has_tools = any( ts.tools for ts in self.agent.toolsets if isinstance(ts, FunctionToolset) ) - # Signal agent node STARTED - yield UiPathRuntimeStateEvent( - payload={}, - node_name=agent_name, - phase=UiPathRuntimeStatePhase.STARTED, - ) - async with self.agent.iter(user_prompt, deps=deps) as agent_run: - async for node in agent_run: + node = agent_run.next_node + + while not isinstance(node, End): if Agent.is_model_request_node(node): - # LLM call: highlight agent node yield UiPathRuntimeStateEvent( - payload={}, + payload=self._model_request_payload(node), node_name=agent_name, phase=UiPathRuntimeStatePhase.STARTED, ) + model_node = node + node = await agent_run.next(node) + yield UiPathRuntimeMessageEvent( - payload=json.loads(serialize_json(node.request)), + payload=json.loads(serialize_json(model_node.request)), metadata={"event_name": "model_request"}, ) yield UiPathRuntimeStateEvent( - payload={}, + payload=self._model_response_payload(node), node_name=agent_name, phase=UiPathRuntimeStatePhase.COMPLETED, ) - elif Agent.is_call_tools_node(node) and has_tools: - # Tool execution: highlight tools node - yield UiPathRuntimeStateEvent( - payload={}, - node_name=tools_node, - phase=UiPathRuntimeStatePhase.STARTED, - ) + elif Agent.is_call_tools_node(node): + tool_calls = node.model_response.tool_calls if has_tools else [] - yield UiPathRuntimeStateEvent( - payload={}, - node_name=tools_node, - phase=UiPathRuntimeStatePhase.COMPLETED, - ) + if tool_calls: + yield UiPathRuntimeStateEvent( + payload={ + "tool_calls": [ + {"tool_name": tc.tool_name} for tc in tool_calls + ], + }, + node_name=tools_node_name, + phase=UiPathRuntimeStatePhase.STARTED, + ) + + node = await agent_run.next(node) + + if tool_calls: + yield UiPathRuntimeStateEvent( + payload=self._tool_results_payload(node), + node_name=tools_node_name, + phase=UiPathRuntimeStatePhase.COMPLETED, + ) + + else: + node = await agent_run.next(node) - # Capture result inside the async with block assert agent_run.result is not None final_output = agent_run.result.output - # Signal agent node COMPLETED - yield UiPathRuntimeStateEvent( - payload={}, - node_name=agent_name, - phase=UiPathRuntimeStatePhase.COMPLETED, - ) - yield self._create_success_result(final_output) except Exception as e: raise self._create_runtime_error(e) from e + @staticmethod + def _model_request_payload(node: Any) -> dict[str, Any]: + """Build payload for a ModelRequestNode STARTED event.""" + payload: dict[str, Any] = {} + try: + parts = node.request.parts if node.request else [] + prompt_parts = [ + p.content + for p in parts + if hasattr(p, "content") and isinstance(p.content, str) + ] + if prompt_parts: + payload["prompt"] = prompt_parts[-1] + except Exception: + pass + return payload + + @staticmethod + def _model_response_payload(next_node: Any) -> dict[str, Any]: + """Build payload for a ModelRequestNode COMPLETED event. + + After agent_run.next() the returned node is the CallToolsNode + which carries the model_response with usage data. + """ + payload: dict[str, Any] = {} + try: + response = next_node.model_response + if response.model_name: + payload["model_name"] = response.model_name + usage = response.usage + if usage: + payload["usage"] = { + "input_tokens": usage.input_tokens, + "output_tokens": usage.output_tokens, + } + except Exception: + pass + return payload + + @staticmethod + def _tool_results_payload(next_node: Any) -> dict[str, Any]: + """Build payload for a CallToolsNode COMPLETED event. + + After agent_run.next() the returned node is a ModelRequestNode + whose request.parts contain ToolReturnPart objects with results. + """ + from pydantic_ai.messages import ToolReturnPart + + payload: dict[str, Any] = {} + try: + parts = next_node.request.parts if next_node.request else [] + results = [] + for p in parts: + if not isinstance(p, ToolReturnPart): + continue + result: dict[str, Any] = {"tool_name": p.tool_name} + content = p.content + if isinstance(content, (str, dict, list)): + result["content"] = content + else: + result["content"] = str(content) + results.append(result) + if results: + payload["tool_results"] = results + except Exception: + pass + return payload + def _prepare_input( self, input: dict[str, Any] | None ) -> tuple[str, BaseModel | None]: @@ -172,34 +234,78 @@ def _prepare_input( return self._extract_messages(input), None def _extract_messages(self, input: dict[str, Any]) -> str: - """Extract a user prompt string from the 'messages' field.""" - messages = input.get("messages", "") - if not isinstance(messages, (str, list)): + """Extract a user prompt string from the 'messages' field. + + Expects UiPath conversation format: + [{"role": "user", "contentParts": [{"data": {"inline": "..."}}]}] + """ + messages = input.get("messages") + if not isinstance(messages, list) or not messages: return "" - if isinstance(messages, list): - parts = [] - for msg in messages: - if isinstance(msg, dict): - parts.append(str(msg.get("content", ""))) - else: - parts.append(str(msg)) - return " ".join(parts) + # Extract text from the last user message + for msg in reversed(messages): + if not isinstance(msg, dict): + continue + role = msg.get("role", "") + if role and role != "user": + continue + text = self._extract_text_from_content_parts(msg) + if text: + return text + + # Fallback: extract from last message regardless of role + if isinstance(messages[-1], dict): + text = self._extract_text_from_content_parts(messages[-1]) + if text: + return text + + return "" + + @staticmethod + def _extract_text_from_content_parts(msg: dict[str, Any]) -> str: + """Extract text from a message's contentParts in UiPath conversation format.""" + content_parts = msg.get("contentParts") + if not isinstance(content_parts, list): + return "" - return messages + texts: list[str] = [] + for cp in content_parts: + if not isinstance(cp, dict): + continue + data = cp.get("data") + if isinstance(data, dict) and "inline" in data: + inline = data["inline"] + if isinstance(inline, str) and inline: + texts.append(inline) + return "".join(texts) def _create_success_result(self, output: Any) -> UiPathRuntimeResult: """Create result for successful completion.""" - serialized_output = json.loads(serialize_json(output)) - - if not isinstance(serialized_output, dict): - serialized_output = {"result": serialized_output} + if self.agent.output_type is str and isinstance(output, str): + serialized_output = self._wrap_as_conversation_message(output) + else: + serialized_output = json.loads(serialize_json(output)) + if not isinstance(serialized_output, dict): + serialized_output = {"result": serialized_output} return UiPathRuntimeResult( output=serialized_output, status=UiPathRuntimeStatus.SUCCESSFUL, ) + @staticmethod + def _wrap_as_conversation_message(text: str) -> dict[str, Any]: + """Wrap a plain string as a UiPath conversation message output.""" + return { + "messages": [ + { + "role": "assistant", + "contentParts": [{"data": {"inline": text}}], + } + ] + } + def _create_runtime_error(self, e: Exception) -> UiPathPydanticAIRuntimeError: """Map exceptions to appropriate runtime errors.""" if isinstance(e, UiPathPydanticAIRuntimeError): diff --git a/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/runtime/schema.py b/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/runtime/schema.py index 999bce6a..9d2912b9 100644 --- a/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/runtime/schema.py +++ b/packages/uipath-pydantic-ai/src/uipath_pydantic_ai/runtime/schema.py @@ -37,6 +37,21 @@ def _get_agent_name(agent: Agent[Any, Any]) -> str: return agent.name or "agent" +def _get_model_name(agent: Agent[Any, Any]) -> str | None: + """Extract the model name from a PydanticAI Agent.""" + try: + model = agent.model + if model is None: + return None + if isinstance(model, str): + return model + if hasattr(model, "model_name"): + return model.model_name + except Exception: + pass + return None + + def _get_tool_names(agent: Agent[Any, Any]) -> list[str]: """Get the list of function tool names from an agent.""" tool_names: list[str] = [] @@ -117,7 +132,7 @@ def get_entrypoints_schema(agent: Agent[Any, Any]) -> dict[str, Any]: Output schema: - If agent has output_type (Pydantic model): output IS that model's schema - - Otherwise: generic result fallback + - Otherwise: UiPath conversation message format """ # --- Input schema --- deps_type = get_deps_type(agent) @@ -127,21 +142,8 @@ def get_entrypoints_schema(agent: Agent[Any, Any]) -> dict[str, Any]: input_schema = None if input_schema is None: - # Conversational fallback: messages-only input - input_schema = { - "type": "object", - "properties": { - "messages": { - "anyOf": [ - {"type": "string"}, - {"type": "array", "items": {"type": "object"}}, - ], - "title": "Messages", - "description": "User messages to send to the agent", - }, - }, - "required": ["messages"], - } + # Conversational fallback: UiPath conversation message format + input_schema = _default_messages_schema() # --- Output schema --- output_type = _get_output_type(agent) @@ -151,22 +153,8 @@ def get_entrypoints_schema(agent: Agent[Any, Any]) -> dict[str, Any]: output_schema = _extract_schema_from_model(output_type) if output_schema is None: - # Generic result fallback - output_schema = { - "type": "object", - "properties": { - "result": { - "title": "Result", - "description": "The agent's response", - "anyOf": [ - {"type": "string"}, - {"type": "object"}, - {"type": "array", "items": {"type": "object"}}, - ], - } - }, - "required": ["result"], - } + # Conversational fallback: UiPath conversation message format + output_schema = _default_messages_schema() return {"input": input_schema, "output": output_schema} @@ -189,13 +177,17 @@ def _add_agent_and_tools(current_agent: Agent[Any, Any]) -> None: return visited.add(agent_name) + model_name = _get_model_name(current_agent) + node_type = "model" if model_name else "node" + metadata = {"model_name": model_name} if model_name else None + nodes.append( UiPathRuntimeNode( id=agent_name, name=agent_name, - type="node", + type=node_type, subgraph=None, - metadata=None, + metadata=metadata, ) ) @@ -276,6 +268,56 @@ def _add_agent_and_tools(current_agent: Agent[Any, Any]) -> None: return UiPathRuntimeGraph(nodes=nodes, edges=edges) +def _conversation_message_item_schema() -> dict[str, Any]: + """Minimal message schema: role and contentParts required, contentParts items only need data.inline.""" + return { + "type": "object", + "properties": { + "role": {"type": "string"}, + "contentParts": { + "type": "array", + "items": { + "type": "object", + "properties": { + "mimeType": {"type": "string"}, + "data": { + "type": "object", + "properties": { + "inline": {}, + }, + "required": ["inline"], + }, + "citations": { + "type": "array", + "items": {"type": "object"}, + }, + }, + "required": ["data"], + }, + }, + "toolCalls": {"type": "array", "items": {"type": "object"}}, + "interrupts": {"type": "array", "items": {"type": "object"}}, + }, + "required": ["role", "contentParts"], + } + + +def _default_messages_schema() -> dict[str, Any]: + """Default messages schema using UiPath conversation message format.""" + return { + "type": "object", + "properties": { + "messages": { + "type": "array", + "items": _conversation_message_item_schema(), + "title": "Messages", + "description": "UiPath conversation messages", + } + }, + "required": ["messages"], + } + + __all__ = [ "get_entrypoints_schema", "get_agent_schema", diff --git a/packages/uipath-pydantic-ai/tests/test_runtime.py b/packages/uipath-pydantic-ai/tests/test_runtime.py index 9eb1b45f..11bedda2 100644 --- a/packages/uipath-pydantic-ai/tests/test_runtime.py +++ b/packages/uipath-pydantic-ai/tests/test_runtime.py @@ -1,11 +1,34 @@ """Tests for PydanticAI runtime execution.""" +from typing import Any + import pytest +from pydantic import BaseModel +from pydantic_ai import Agent from uipath_pydantic_ai.runtime.errors import ( UiPathPydanticAIErrorCode, UiPathPydanticAIRuntimeError, ) +from uipath_pydantic_ai.runtime.runtime import UiPathPydanticAIRuntime + +# ============= HELPERS ============= + + +def _uipath_message(text: str, role: str = "user") -> dict[str, Any]: + """Build a UiPath conversation message dict.""" + return { + "role": role, + "contentParts": [{"data": {"inline": text}}], + } + + +def _uipath_input(text: str) -> dict[str, Any]: + """Build a UiPath input dict with a single user message.""" + return {"messages": [_uipath_message(text)]} + + +# ============= ERROR TESTS ============= def test_error_handling(): @@ -51,21 +74,64 @@ def test_error_codes(): ) -def test_runtime_input_preparation(): - """Test that runtime correctly prepares agent input (conversational mode).""" - from pydantic_ai import Agent +# ============= INPUT PREPARATION TESTS (UiPath format) ============= - from uipath_pydantic_ai.runtime.runtime import UiPathPydanticAIRuntime +def test_runtime_input_uipath_message(): + """Test that runtime extracts text from UiPath conversation messages.""" agent = Agent("test", name="test_agent") runtime = UiPathPydanticAIRuntime(agent=agent) - # Test string messages -> (prompt, None deps) - prompt, deps = runtime._prepare_input({"messages": "Hello"}) + prompt, deps = runtime._prepare_input(_uipath_input("Hello")) assert prompt == "Hello" assert deps is None - # Test empty input + +def test_runtime_input_multiple_messages_takes_last_user(): + """Test that runtime picks the last user message from a conversation.""" + agent = Agent("test", name="test_agent") + runtime = UiPathPydanticAIRuntime(agent=agent) + + prompt, deps = runtime._prepare_input( + { + "messages": [ + _uipath_message("First question"), + _uipath_message("I am the assistant", role="assistant"), + _uipath_message("Second question"), + ] + } + ) + assert prompt == "Second question" + assert deps is None + + +def test_runtime_input_multipart_content(): + """Test that runtime concatenates multiple content parts.""" + agent = Agent("test", name="test_agent") + runtime = UiPathPydanticAIRuntime(agent=agent) + + prompt, deps = runtime._prepare_input( + { + "messages": [ + { + "role": "user", + "contentParts": [ + {"data": {"inline": "Hello "}}, + {"data": {"inline": "World"}}, + ], + } + ] + } + ) + assert prompt == "Hello World" + assert deps is None + + +def test_runtime_input_empty(): + """Test empty input returns empty prompt.""" + agent = Agent("test", name="test_agent") + runtime = UiPathPydanticAIRuntime(agent=agent) + prompt, deps = runtime._prepare_input(None) assert prompt == "" assert deps is None @@ -74,26 +140,44 @@ def test_runtime_input_preparation(): assert prompt == "" assert deps is None - # Test non-string/non-list messages + +def test_runtime_input_empty_messages(): + """Test empty messages array returns empty prompt.""" + agent = Agent("test", name="test_agent") + runtime = UiPathPydanticAIRuntime(agent=agent) + + prompt, deps = runtime._prepare_input({"messages": []}) + assert prompt == "" + assert deps is None + + +def test_runtime_input_non_list_messages(): + """Test non-list messages returns empty prompt.""" + agent = Agent("test", name="test_agent") + runtime = UiPathPydanticAIRuntime(agent=agent) + prompt, deps = runtime._prepare_input({"messages": 123}) assert prompt == "" assert deps is None - # Test list input + +def test_runtime_input_falls_back_to_last_message(): + """Test that when no user message exists, falls back to last message.""" + agent = Agent("test", name="test_agent") + runtime = UiPathPydanticAIRuntime(agent=agent) + prompt, deps = runtime._prepare_input( - {"messages": [{"content": "Hello"}, {"content": "World"}]} + {"messages": [_uipath_message("Assistant reply", role="assistant")]} ) - assert "Hello" in prompt - assert "World" in prompt + assert prompt == "Assistant reply" assert deps is None +# ============= STRUCTURED INPUT TESTS ============= + + def test_runtime_structured_input_preparation(): """Test that runtime correctly prepares structured deps input.""" - from pydantic import BaseModel - from pydantic_ai import Agent - - from uipath_pydantic_ai.runtime.runtime import UiPathPydanticAIRuntime class MyInput(BaseModel): query: str @@ -112,10 +196,6 @@ class MyInput(BaseModel): def test_runtime_structured_input_with_messages(): """Test that deps model with a 'messages' field uses it as prompt.""" - from pydantic import BaseModel - from pydantic_ai import Agent - - from uipath_pydantic_ai.runtime.runtime import UiPathPydanticAIRuntime class ReviewInput(BaseModel): messages: str @@ -130,3 +210,375 @@ class ReviewInput(BaseModel): assert prompt == "Review this code" assert isinstance(deps, ReviewInput) assert deps.review_type == "security" + + +# ============= GET_SCHEMA TESTS ============= + + +@pytest.mark.asyncio +async def test_get_schema_conversational_agent(): + """Test that get_schema() returns UiPath conversation message format for a plain agent.""" + agent = Agent("test", name="test_agent") + runtime = UiPathPydanticAIRuntime(agent=agent, runtime_id="test", entrypoint="test") + + schema = await runtime.get_schema() + + # Input should use UiPath conversation messages + assert "messages" in schema.input["properties"] + messages_prop = schema.input["properties"]["messages"] + assert messages_prop["type"] == "array" + assert messages_prop["items"]["properties"]["role"]["type"] == "string" + assert "contentParts" in messages_prop["items"]["properties"] + + # Output should also use UiPath conversation messages + assert "messages" in schema.output["properties"] + out_messages = schema.output["properties"]["messages"] + assert out_messages["type"] == "array" + + +@pytest.mark.asyncio +async def test_get_schema_structured_output_agent(): + """Test that get_schema() returns structured output when output_type is set.""" + + class MyOutput(BaseModel): + answer: str + confidence: float + + agent = Agent("test", name="test_agent", output_type=MyOutput) + runtime = UiPathPydanticAIRuntime(agent=agent, runtime_id="test", entrypoint="test") + + schema = await runtime.get_schema() + + # Input should still be UiPath conversation messages + assert "messages" in schema.input["properties"] + + # Output should be the structured model + assert "answer" in schema.output["properties"] + assert "confidence" in schema.output["properties"] + + +@pytest.mark.asyncio +async def test_get_schema_structured_input_agent(): + """Test that get_schema() returns structured input when deps_type is set.""" + + class MyInput(BaseModel): + query: str + + agent = Agent("test", name="test_agent", deps_type=MyInput) + runtime = UiPathPydanticAIRuntime(agent=agent, runtime_id="test", entrypoint="test") + + schema = await runtime.get_schema() + + # Input should be the structured model + assert "query" in schema.input["properties"] + + # Output should be UiPath conversation messages + assert "messages" in schema.output["properties"] + + +# ============= E2E TESTS WITH MOCKED LLM ============= + + +@pytest.mark.asyncio +async def test_e2e_execute_with_uipath_messages(): + """E2E test: execute a conversational agent with UiPath message input.""" + from pydantic_ai.models.test import TestModel + + agent = Agent(TestModel(), name="test_agent") + runtime = UiPathPydanticAIRuntime(agent=agent, runtime_id="test", entrypoint="test") + + result = await runtime.execute(input=_uipath_input("What is 2+2?")) + + assert result.status.value == "successful" + assert isinstance(result.output, dict) + assert "messages" in result.output + messages = result.output["messages"] + assert len(messages) == 1 + assert messages[0]["role"] == "assistant" + assert isinstance(messages[0]["contentParts"][0]["data"]["inline"], str) + + +@pytest.mark.asyncio +async def test_e2e_execute_with_multi_turn_messages(): + """E2E test: execute with multi-turn conversation, picks last user message.""" + from pydantic_ai.models.test import TestModel + + agent = Agent(TestModel(), name="test_agent") + runtime = UiPathPydanticAIRuntime(agent=agent, runtime_id="test", entrypoint="test") + + result = await runtime.execute( + input={ + "messages": [ + _uipath_message("Hello"), + _uipath_message("I'm an assistant", role="assistant"), + _uipath_message("What is the weather?"), + ] + } + ) + + assert result.status.value == "successful" + assert isinstance(result.output, dict) + assert "messages" in result.output + assert result.output["messages"][0]["role"] == "assistant" + + +@pytest.mark.asyncio +async def test_e2e_execute_structured_output(): + """E2E test: execute agent with structured output_type.""" + from pydantic_ai.models.test import TestModel + + class CityInfo(BaseModel): + name: str + population: int + + agent = Agent( + TestModel(custom_output_args={"name": "Paris", "population": 2161000}), + name="city_agent", + output_type=CityInfo, + ) + runtime = UiPathPydanticAIRuntime(agent=agent, runtime_id="test", entrypoint="test") + + result = await runtime.execute(input=_uipath_input("Tell me about Paris")) + + assert result.status.value == "successful" + assert isinstance(result.output, dict) + assert result.output["name"] == "Paris" + assert result.output["population"] == 2161000 + + +@pytest.mark.asyncio +async def test_e2e_execute_with_tools(): + """E2E test: execute agent with tools using UiPath message input.""" + from pydantic_ai.models.test import TestModel + + def get_answer(ctx, question: str) -> str: + """Answer a question. + + Args: + ctx: The agent context. + question: The question to answer. + + Returns: + The answer. + """ + return "42" + + agent = Agent( + TestModel(), + name="tool_agent", + tools=[get_answer], + ) + runtime = UiPathPydanticAIRuntime(agent=agent, runtime_id="test", entrypoint="test") + + result = await runtime.execute(input=_uipath_input("What is the meaning of life?")) + + assert result.status.value == "successful" + assert isinstance(result.output, dict) + assert "messages" in result.output + assert result.output["messages"][0]["role"] == "assistant" + + +@pytest.mark.asyncio +async def test_e2e_execute_structured_deps(): + """E2E test: execute agent with structured deps_type.""" + from pydantic_ai.models.test import TestModel + + class ResearchInput(BaseModel): + topic: str + max_sources: int = 3 + + agent = Agent( + TestModel(), + name="research_agent", + deps_type=ResearchInput, + ) + runtime = UiPathPydanticAIRuntime(agent=agent, runtime_id="test", entrypoint="test") + + result = await runtime.execute( + input={"topic": "quantum computing", "max_sources": 5} + ) + + assert result.status.value == "successful" + assert isinstance(result.output, dict) + assert "messages" in result.output + assert result.output["messages"][0]["role"] == "assistant" + + +@pytest.mark.asyncio +async def test_e2e_stream_with_uipath_messages(): + """E2E test: stream a conversational agent with UiPath message input.""" + from pydantic_ai.models.test import TestModel + from uipath.runtime import UiPathRuntimeResult + + agent = Agent(TestModel(), name="test_agent") + runtime = UiPathPydanticAIRuntime(agent=agent, runtime_id="test", entrypoint="test") + + events = [] + async for event in runtime.stream(input=_uipath_input("Tell me a joke")): + events.append(event) + + # Last event should be the result + assert len(events) > 0 + result = events[-1] + assert isinstance(result, UiPathRuntimeResult) + assert result.status.value == "successful" + + +@pytest.mark.asyncio +async def test_e2e_stream_events_no_duplicates(): + """E2E test: verify stream events have no duplicates and proper STARTED/COMPLETED pairing.""" + from pydantic_ai.models.test import TestModel + from uipath.runtime.events import ( + UiPathRuntimeStateEvent, + UiPathRuntimeStatePhase, + ) + + agent = Agent(TestModel(), name="my_agent") + runtime = UiPathPydanticAIRuntime(agent=agent, runtime_id="test", entrypoint="test") + + state_events: list[tuple[str, UiPathRuntimeStatePhase]] = [] + async for event in runtime.stream(input=_uipath_input("Hello")): + if isinstance(event, UiPathRuntimeStateEvent) and event.node_name: + state_events.append((event.node_name, event.phase)) + + # Should have exactly one STARTED/COMPLETED pair for the agent node + agent_events = [e for e in state_events if e[0] == "my_agent"] + assert len(agent_events) >= 2 # At least one STARTED + COMPLETED + + # Every STARTED must be followed by a COMPLETED for the same node + open_nodes: dict[str, int] = {} + for node_name, phase in state_events: + if phase == UiPathRuntimeStatePhase.STARTED: + open_nodes[node_name] = open_nodes.get(node_name, 0) + 1 + elif phase == UiPathRuntimeStatePhase.COMPLETED: + assert open_nodes.get(node_name, 0) > 0, ( + f"COMPLETED without STARTED for {node_name}" + ) + open_nodes[node_name] -= 1 + + # All nodes should be closed (no dangling STARTED without COMPLETED) + for node_name, count in open_nodes.items(): + assert count == 0, f"Dangling STARTED for {node_name}" + + # No consecutive duplicate events + for i in range(1, len(state_events)): + assert state_events[i] != state_events[i - 1], ( + f"Consecutive duplicate event: {state_events[i]}" + ) + + +@pytest.mark.asyncio +async def test_e2e_stream_with_tools_events(): + """E2E test: verify tools node events when agent has tools.""" + from pydantic_ai.models.test import TestModel + from uipath.runtime.events import ( + UiPathRuntimeStateEvent, + UiPathRuntimeStatePhase, + ) + + def my_tool(ctx, query: str) -> str: + """Search for something. + + Args: + ctx: The agent context. + query: The search query. + + Returns: + Search results. + """ + return f"Result for {query}" + + agent = Agent(TestModel(), name="tool_agent", tools=[my_tool]) + runtime = UiPathPydanticAIRuntime(agent=agent, runtime_id="test", entrypoint="test") + + state_events = [] + async for event in runtime.stream(input=_uipath_input("Search for cats")): + if isinstance(event, UiPathRuntimeStateEvent): + state_events.append((event.node_name, event.phase)) + + # Should have agent events + agent_events = [(n, p) for n, p in state_events if n == "tool_agent"] + assert len(agent_events) >= 2 + + # Should have tools events (TestModel calls tools by default) + tools_events = [(n, p) for n, p in state_events if n == "tool_agent_tools"] + assert len(tools_events) >= 2 # At least one STARTED + COMPLETED pair + + # Tools events should be properly paired + for i in range(0, len(tools_events), 2): + assert tools_events[i][1] == UiPathRuntimeStatePhase.STARTED + assert tools_events[i + 1][1] == UiPathRuntimeStatePhase.COMPLETED + + # Agent events should also be properly paired + for i in range(0, len(agent_events), 2): + assert agent_events[i][1] == UiPathRuntimeStatePhase.STARTED + assert agent_events[i + 1][1] == UiPathRuntimeStatePhase.COMPLETED + + +@pytest.mark.asyncio +async def test_e2e_stream_event_payloads(): + """E2E test: verify state events carry meaningful payloads.""" + from pydantic_ai.models.test import TestModel + from uipath.runtime.events import ( + UiPathRuntimeStateEvent, + UiPathRuntimeStatePhase, + ) + + def my_tool(ctx, query: str) -> str: + """Search for something. + + Args: + ctx: The agent context. + query: The search query. + + Returns: + Search results. + """ + return f"Result for {query}" + + agent = Agent(TestModel(), name="payload_agent", tools=[my_tool]) + runtime = UiPathPydanticAIRuntime(agent=agent, runtime_id="test", entrypoint="test") + + state_events: list[UiPathRuntimeStateEvent] = [] + async for event in runtime.stream(input=_uipath_input("Search for cats")): + if isinstance(event, UiPathRuntimeStateEvent): + state_events.append(event) + + # Agent COMPLETED events should have model_name and usage + agent_completed = [ + e + for e in state_events + if e.node_name == "payload_agent" + and e.phase == UiPathRuntimeStatePhase.COMPLETED + ] + assert len(agent_completed) >= 1 + for event in agent_completed: + assert "model_name" in event.payload or "usage" in event.payload + + # Tools STARTED events should have tool_calls with tool_name + tools_started = [ + e + for e in state_events + if e.node_name == "payload_agent_tools" + and e.phase == UiPathRuntimeStatePhase.STARTED + ] + assert len(tools_started) >= 1 + for event in tools_started: + assert "tool_calls" in event.payload + assert len(event.payload["tool_calls"]) > 0 + assert "tool_name" in event.payload["tool_calls"][0] + + # Tools COMPLETED events should have tool_results with content + tools_completed = [ + e + for e in state_events + if e.node_name == "payload_agent_tools" + and e.phase == UiPathRuntimeStatePhase.COMPLETED + ] + assert len(tools_completed) >= 1 + for event in tools_completed: + assert "tool_results" in event.payload + assert len(event.payload["tool_results"]) > 0 + result = event.payload["tool_results"][0] + assert "tool_name" in result + assert "content" in result diff --git a/packages/uipath-pydantic-ai/tests/test_schema.py b/packages/uipath-pydantic-ai/tests/test_schema.py index c4d1bd86..702404bf 100644 --- a/packages/uipath-pydantic-ai/tests/test_schema.py +++ b/packages/uipath-pydantic-ai/tests/test_schema.py @@ -1,5 +1,7 @@ """Tests for PydanticAI schema extraction and graph building.""" +from typing import Any + from pydantic import BaseModel from pydantic_ai import Agent @@ -211,11 +213,14 @@ def test_schema_with_output_type(): def test_schema_fallback_without_output_type(): - """Test that schema falls back to defaults when no output_type.""" + """Test that schema falls back to UiPath conversation messages when no output_type.""" schema = get_entrypoints_schema(agent_plain) assert "messages" in schema["input"]["properties"] - assert "result" in schema["output"]["properties"] + # Output should be UiPath conversation messages, not a generic "result" + assert "messages" in schema["output"]["properties"] + output_messages = schema["output"]["properties"]["messages"] + assert output_messages["type"] == "array" def test_schema_with_plain_agent(): @@ -223,7 +228,7 @@ def test_schema_with_plain_agent(): schema = get_entrypoints_schema(agent_plain) assert "messages" in schema["input"]["properties"] - assert "result" in schema["output"]["properties"] + assert "messages" in schema["output"]["properties"] # ============= GRAPH TESTS ============= @@ -262,7 +267,7 @@ def test_graph_node_types(): assert node_types["__start__"] == "__start__" assert node_types["__end__"] == "__end__" - assert node_types["tools_agent"] == "node" + assert node_types["tools_agent"] == "model" assert node_types["tools_agent_tools"] == "tool" @@ -286,6 +291,28 @@ def test_graph_tool_edges(): assert ("tools_agent_tools", "tools_agent", None) in edges +def test_graph_agent_model_metadata(): + """Test that agent nodes with a model have model_name metadata.""" + graph = get_agent_schema(agent_with_tools) + + node_map = {node.id: node for node in graph.nodes} + agent_node = node_map["tools_agent"] + assert agent_node.type == "model" + assert agent_node.metadata is not None + assert agent_node.metadata["model_name"] == "test" + + +def test_graph_agent_without_model(): + """Test that agent nodes without a model use type=node.""" + no_model_agent = Agent(None, name="no_model", defer_model_check=True) + graph = get_agent_schema(no_model_agent) + + node_map = {node.id: node for node in graph.nodes} + agent_node = node_map["no_model"] + assert agent_node.type == "node" + assert agent_node.metadata is None + + def test_graph_tools_metadata(): """Test that tools nodes have correct metadata.""" graph = get_agent_schema(agent_with_tools) @@ -313,3 +340,96 @@ def test_graph_no_subgraphs(): for node in graph.nodes: assert node.subgraph is None + + +# ============= UIPATH CONVERSATION MESSAGE FORMAT TESTS ============= + + +def _validate_conversation_message_schema(schema: dict[str, Any]) -> None: + """Helper to validate that a schema matches the UiPath conversation message format.""" + assert schema["type"] == "object" + assert "messages" in schema["properties"] + assert "messages" in schema["required"] + + messages_prop = schema["properties"]["messages"] + assert messages_prop["type"] == "array" + assert "items" in messages_prop + + item = messages_prop["items"] + assert item["type"] == "object" + assert "role" in item["properties"] + assert item["properties"]["role"]["type"] == "string" + assert "contentParts" in item["properties"] + assert "role" in item["required"] + assert "contentParts" in item["required"] + + content_parts = item["properties"]["contentParts"] + assert content_parts["type"] == "array" + cp_item = content_parts["items"] + assert cp_item["type"] == "object" + assert "data" in cp_item["properties"] + assert "data" in cp_item["required"] + + data_prop = cp_item["properties"]["data"] + assert data_prop["type"] == "object" + assert "inline" in data_prop["properties"] + assert "inline" in data_prop["required"] + + +def test_default_input_schema_is_uipath_messages(): + """Test that conversational agents get UiPath conversation message input schema.""" + schema = get_entrypoints_schema(agent_plain) + _validate_conversation_message_schema(schema["input"]) + + +def test_default_output_schema_is_uipath_messages(): + """Test that conversational agents get UiPath conversation message output schema.""" + schema = get_entrypoints_schema(agent_plain) + _validate_conversation_message_schema(schema["output"]) + + +def test_output_only_agent_has_uipath_messages_input(): + """Test agent with output_type still uses UiPath messages for input.""" + schema = get_entrypoints_schema(agent_with_output) + _validate_conversation_message_schema(schema["input"]) + # Output should be the structured model, not messages + assert "original_text" in schema["output"]["properties"] + + +def test_deps_only_agent_has_uipath_messages_output(): + """Test agent with deps_type still uses UiPath messages for output.""" + schema = get_entrypoints_schema(agent_with_deps) + _validate_conversation_message_schema(schema["output"]) + # Input should be the structured model, not messages + assert "text" in schema["input"]["properties"] + + +def test_message_schema_has_optional_fields(): + """Test that the message schema includes optional fields like toolCalls and interrupts.""" + schema = get_entrypoints_schema(agent_plain) + item = schema["input"]["properties"]["messages"]["items"] + assert "toolCalls" in item["properties"] + assert "interrupts" in item["properties"] + assert "mimeType" in item["properties"]["contentParts"]["items"]["properties"] + assert "citations" in item["properties"]["contentParts"]["items"]["properties"] + + +def test_message_schema_structure(): + """Test that the default message schema has the correct UiPath structure.""" + schema = get_entrypoints_schema(agent_plain) + input_schema = schema["input"] + + messages_prop = input_schema["properties"]["messages"] + assert messages_prop["title"] == "Messages" + assert messages_prop["description"] == "UiPath conversation messages" + + item = messages_prop["items"] + # toolCalls and interrupts should be arrays of objects + assert item["properties"]["toolCalls"] == { + "type": "array", + "items": {"type": "object"}, + } + assert item["properties"]["interrupts"] == { + "type": "array", + "items": {"type": "object"}, + } diff --git a/packages/uipath-pydantic-ai/uv.lock b/packages/uipath-pydantic-ai/uv.lock index 732ab7a2..5ee75dbe 100644 --- a/packages/uipath-pydantic-ai/uv.lock +++ b/packages/uipath-pydantic-ai/uv.lock @@ -3656,7 +3656,7 @@ wheels = [ [[package]] name = "uipath-pydantic-ai" -version = "0.0.1" +version = "0.0.2" source = { editable = "." } dependencies = [ { name = "openinference-instrumentation-pydantic-ai" },