fix: use public API guards, fix enum comparison, run ruff format

- Use Agent.is_call_tools_node() / Agent.is_model_request_node() for
  type narrowing instead of private module imports
- Remove all imports from pydantic_ai._agent_graph
- Compare phase against UiPathRuntimeStatePhase.STARTED enum directly
- Run ruff format on all files

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

Author: cristipufu

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

@@ -7,8 +7,7 @@
 
 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 pydantic_ai.messages import ToolReturnPart
 from uipath.core.chat.content import (
     UiPathConversationContentPartChunkEvent,
     UiPathConversationContentPartEndEvent,
@@ -154,12 +153,12 @@ async def stream(
                                 ),
                             )
 
-                        assert isinstance(next_node, CallToolsNode)
-                        yield UiPathRuntimeStateEvent(
-                            payload=self._model_response_payload(next_node),
-                            node_name=agent_name,
-                            phase=UiPathRuntimeStatePhase.COMPLETED,
-                        )
+                        if Agent.is_call_tools_node(next_node):
+                            yield UiPathRuntimeStateEvent(
+                                payload=self._model_response_payload(next_node),
+                                node_name=agent_name,
+                                phase=UiPathRuntimeStatePhase.COMPLETED,
+                            )
                         node = next_node
 
                     elif Agent.is_call_tools_node(node):
@@ -178,7 +177,7 @@ async def stream(
 
                         next_node = await agent_run.next(node)
 
-                        if tool_calls and isinstance(next_node, ModelRequestNode):
+                        if tool_calls and Agent.is_model_request_node(next_node):
                             yield UiPathRuntimeStateEvent(
                                 payload=self._tool_results_payload(next_node),
                                 node_name=tools_node_name,
@@ -204,9 +203,7 @@ def _get_timestamp() -> str:
         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]:
+    def _model_request_payload(node: Any) -> dict[str, Any]:
         """Build payload for a ModelRequestNode STARTED event."""
         payload: dict[str, Any] = {}
         try:
@@ -223,17 +220,15 @@ def _model_request_payload(
         return payload
 
     @staticmethod
-    def _model_response_payload(
-        next_node: CallToolsNode[Any, Any],
-    ) -> dict[str, Any]:
+    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: ModelResponse = next_node.model_response
+            response = next_node.model_response
             if response.model_name:
                 payload["model_name"] = response.model_name
             usage = response.usage
@@ -247,9 +242,7 @@ def _model_response_payload(
         return payload
 
     @staticmethod
-    def _tool_results_payload(
-        next_node: ModelRequestNode[Any, Any],
-    ) -> dict[str, Any]:
+    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

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

@@ -286,11 +286,11 @@ def _conversation_message_item_schema() -> dict[str, Any]:
                                 "inline": {},
                             },
                             "required": ["inline"],
-                        }
+                        },
                     },
                     "required": ["data"],
                 },
-            }
+            },
         },
         "required": ["role", "contentParts"],
     }

packages/uipath-pydantic-ai/tests/test_runtime.py

@@ -721,19 +721,18 @@ def my_tool(ctx, query: str) -> str:
     assert len(ends) >= 1
 
     # Text chunks should exist
-    chunks = [
-        e
-        for e in msg_events
-        if e.content_part and e.content_part.chunk
-    ]
+    chunks = [e for e in msg_events if e.content_part and e.content_part.chunk]
     assert len(chunks) >= 1
 
 
 @pytest.mark.asyncio
 async def test_stream_tool_only_turn_skips_message_events():
     """Model turns that produce only tool calls (no text) should not emit message events."""
     from pydantic_ai.models.test import TestModel
-    from uipath.runtime.events import UiPathRuntimeStateEvent
+    from uipath.runtime.events import (
+        UiPathRuntimeStateEvent,
+        UiPathRuntimeStatePhase,
+    )
 
     def my_tool(ctx, query: str) -> str:
         """A tool.
@@ -761,9 +760,9 @@ def my_tool(ctx, query: str) -> str:
 
     # Should have multiple model turns via state events (tool turn + final turn)
     agent_started = [
-        e for e in state_events
-        if e.node_name == "skip_agent"
-        and e.phase.value == "started"
+        e
+        for e in state_events
+        if e.node_name == "skip_agent" and e.phase == UiPathRuntimeStatePhase.STARTED
     ]
     assert len(agent_started) >= 2  # at least 2 model request turns