fix: pydantic-ai stream LLM tokens as UiPathConversationMessageEvent

Replace raw dict payloads with proper UiPathConversationMessageEvent objects
following the START -> CHUNK(s) -> END lifecycle. Use node.stream() +
stream_text(delta=True) for real-time token streaming to the Chat UI.

- Stream individual LLM tokens via pydantic-ai's ModelRequestNode.stream()
- Emit proper message lifecycle events (start, content chunks, end)
- Add strong typing for node parameters (ModelRequestNode, CallToolsNode)
- Simplify message schema (remove unused toolCalls/interrupts/citations)
- Add 6 e2e streaming tests with mocked LLM
- Bump version to 0.0.3

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: cristipufu

packages/uipath-pydantic-ai/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "uipath-pydantic-ai"
-version = "0.0.2"
+version = "0.0.3"
 description = "Python SDK that enables developers to build and deploy PydanticAI agents to the UiPath Cloud Platform"
 readme = "README.md"
 requires-python = ">=3.11"

packages/uipath-pydantic-ai/samples/quickstart-agent/agent.mermaid

@@ -0,0 +1,9 @@
+flowchart TB
+  __start__(__start__)
+  weather_agent(weather_agent)
+  weather_agent_tools(tools)
+  __end__(__end__)
+  weather_agent --> weather_agent_tools
+  weather_agent_tools --> weather_agent
+  __start__ --> |input|weather_agent
+  weather_agent --> |output|__end__

packages/uipath-pydantic-ai/src/uipath_pydantic_ai/runtime/runtime.py

@@ -1,11 +1,25 @@
 """Runtime class for executing PydanticAI Agents within the UiPath framework."""
 
 import json
+from datetime import datetime, timezone
 from typing import Any, AsyncGenerator
 from uuid import uuid4
 
 from pydantic import BaseModel
 from pydantic_ai import Agent, FunctionToolset
+from pydantic_ai._agent_graph import CallToolsNode, ModelRequestNode
+from pydantic_ai.messages import ModelResponse, ToolReturnPart
+from uipath.core.chat.content import (
+    UiPathConversationContentPartChunkEvent,
+    UiPathConversationContentPartEndEvent,
+    UiPathConversationContentPartEvent,
+    UiPathConversationContentPartStartEvent,
+)
+from uipath.core.chat.message import (
+    UiPathConversationMessageEndEvent,
+    UiPathConversationMessageEvent,
+    UiPathConversationMessageStartEvent,
+)
 from uipath.core.serialization import serialize_json
 from uipath.runtime import (
     UiPathExecuteOptions,
@@ -88,18 +102,65 @@ async def stream(
                         )
 
                         model_node = node
-                        node = await agent_run.next(node)
-
-                        yield UiPathRuntimeMessageEvent(
-                            payload=json.loads(serialize_json(model_node.request)),
-                            metadata={"event_name": "model_request"},
-                        )
+                        message_id = str(uuid4())
+                        content_part_id = f"chunk-{message_id}-0"
+                        has_text = False
+
+                        async with model_node.stream(agent_run.ctx) as stream:
+                            async for text_chunk in stream.stream_text(
+                                delta=True, debounce_by=None
+                            ):
+                                if not has_text:
+                                    has_text = True
+                                    yield UiPathRuntimeMessageEvent(
+                                        payload=UiPathConversationMessageEvent(
+                                            message_id=message_id,
+                                            start=UiPathConversationMessageStartEvent(
+                                                role="assistant",
+                                                timestamp=self._get_timestamp(),
+                                            ),
+                                            content_part=UiPathConversationContentPartEvent(
+                                                content_part_id=content_part_id,
+                                                start=UiPathConversationContentPartStartEvent(
+                                                    mime_type="text/plain",
+                                                ),
+                                            ),
+                                        ),
+                                    )
+
+                                yield UiPathRuntimeMessageEvent(
+                                    payload=UiPathConversationMessageEvent(
+                                        message_id=message_id,
+                                        content_part=UiPathConversationContentPartEvent(
+                                            content_part_id=content_part_id,
+                                            chunk=UiPathConversationContentPartChunkEvent(
+                                                data=text_chunk,
+                                            ),
+                                        ),
+                                    ),
+                                )
+
+                        next_node = await agent_run.next(model_node)
+
+                        if has_text:
+                            yield UiPathRuntimeMessageEvent(
+                                payload=UiPathConversationMessageEvent(
+                                    message_id=message_id,
+                                    end=UiPathConversationMessageEndEvent(),
+                                    content_part=UiPathConversationContentPartEvent(
+                                        content_part_id=content_part_id,
+                                        end=UiPathConversationContentPartEndEvent(),
+                                    ),
+                                ),
+                            )
 
+                        assert isinstance(next_node, CallToolsNode)
                         yield UiPathRuntimeStateEvent(
-                            payload=self._model_response_payload(node),
+                            payload=self._model_response_payload(next_node),
                             node_name=agent_name,
                             phase=UiPathRuntimeStatePhase.COMPLETED,
                         )
+                        node = next_node
 
                     elif Agent.is_call_tools_node(node):
                         tool_calls = node.model_response.tool_calls if has_tools else []
@@ -115,14 +176,15 @@ async def stream(
                                 phase=UiPathRuntimeStatePhase.STARTED,
                             )
 
-                        node = await agent_run.next(node)
+                        next_node = await agent_run.next(node)
 
-                        if tool_calls:
+                        if tool_calls and isinstance(next_node, ModelRequestNode):
                             yield UiPathRuntimeStateEvent(
-                                payload=self._tool_results_payload(node),
+                                payload=self._tool_results_payload(next_node),
                                 node_name=tools_node_name,
                                 phase=UiPathRuntimeStatePhase.COMPLETED,
                             )
+                        node = next_node
 
                     else:
                         node = await agent_run.next(node)
@@ -136,7 +198,15 @@ async def stream(
             raise self._create_runtime_error(e) from e
 
     @staticmethod
-    def _model_request_payload(node: Any) -> dict[str, Any]:
+    def _get_timestamp() -> str:
+        """Get current UTC timestamp in ISO 8601 format."""
+        now = datetime.now(timezone.utc)
+        return now.strftime("%Y-%m-%dT%H:%M:%S.") + f"{now.microsecond // 1000:03d}Z"
+
+    @staticmethod
+    def _model_request_payload(
+        node: ModelRequestNode[Any, Any],
+    ) -> dict[str, Any]:
         """Build payload for a ModelRequestNode STARTED event."""
         payload: dict[str, Any] = {}
         try:
@@ -153,15 +223,17 @@ def _model_request_payload(node: Any) -> dict[str, Any]:
         return payload
 
     @staticmethod
-    def _model_response_payload(next_node: Any) -> dict[str, Any]:
+    def _model_response_payload(
+        next_node: CallToolsNode[Any, 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
+            response: ModelResponse = next_node.model_response
             if response.model_name:
                 payload["model_name"] = response.model_name
             usage = response.usage
@@ -175,14 +247,14 @@ def _model_response_payload(next_node: Any) -> dict[str, Any]:
         return payload
 
     @staticmethod
-    def _tool_results_payload(next_node: Any) -> dict[str, Any]:
+    def _tool_results_payload(
+        next_node: ModelRequestNode[Any, 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 []

packages/uipath-pydantic-ai/src/uipath_pydantic_ai/runtime/schema.py

@@ -286,17 +286,11 @@ def _conversation_message_item_schema() -> dict[str, Any]:
                                 "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"],
     }