sorted import block

Author: michael-chou359

examples/tutorials/10_async/10_temporal/110_pydantic_ai/project/agent.py

@@ -17,13 +17,13 @@
 
 from __future__ import annotations
 
-from collections.abc import AsyncIterable
 from datetime import datetime
+from collections.abc import AsyncIterable
 
 from pydantic import BaseModel
 from pydantic_ai import Agent, RunContext
-from pydantic_ai.durable_exec.temporal import TemporalAgent
 from pydantic_ai.messages import AgentStreamEvent
+from pydantic_ai.durable_exec.temporal import TemporalAgent
 
 from project.tools import get_weather
 from agentex.lib.adk import stream_pydantic_ai_events
@@ -61,9 +61,7 @@ def _build_base_agent() -> Agent[TaskDeps, str]:
     agent: Agent[TaskDeps, str] = Agent(
         MODEL_NAME,
         deps_type=TaskDeps,
-        system_prompt=SYSTEM_PROMPT.format(
-            timestamp=datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-        ),
+        system_prompt=SYSTEM_PROMPT.format(timestamp=datetime.now().strftime("%Y-%m-%d %H:%M:%S")),
     )
     agent.tool_plain(get_weather)
     return agent

examples/tutorials/10_async/10_temporal/110_pydantic_ai/project/workflow.py

@@ -9,12 +9,13 @@
 
 from __future__ import annotations
 
-import json
 import os
+import json
 
 from temporalio import workflow
 
 from agentex.lib import adk
+from project.agent import TaskDeps, temporal_agent
 from agentex.lib.types.acp import SendEventParams, CreateTaskParams
 from agentex.lib.types.tracing import SGPTracingProcessorConfig
 from agentex.lib.utils.logging import make_logger
@@ -26,8 +27,6 @@
     add_tracing_processor_config,
 )
 
-from project.agent import TaskDeps, temporal_agent
-
 add_tracing_processor_config(
     SGPTracingProcessorConfig(
         sgp_api_key=os.environ.get("SGP_API_KEY", ""),

src/agentex/lib/adk/_modules/_pydantic_ai_async.py

@@ -31,15 +31,15 @@ async def stream_pydantic_ai_events(stream, task_id: str) -> str:
     import json
 
     from pydantic_ai.messages import (
-        FunctionToolResultEvent,
-        PartDeltaEvent,
-        PartEndEvent,
-        PartStartEvent,
         TextPart,
-        TextPartDelta,
+        PartEndEvent,
         ThinkingPart,
-        ThinkingPartDelta,
         ToolCallPart,
+        TextPartDelta,
+        PartDeltaEvent,
+        PartStartEvent,
+        ThinkingPartDelta,
+        FunctionToolResultEvent,
     )
 
     from agentex.lib import adk

src/agentex/lib/adk/_modules/_pydantic_ai_sync.py

@@ -23,36 +23,36 @@ async def handle_message_send(params):
 import json
 from typing import Any, AsyncIterator
 
+from pydantic_ai.run import AgentRunResultEvent
 from pydantic_ai.messages import (
-    FinalResultEvent,
-    FunctionToolCallEvent,
-    FunctionToolResultEvent,
-    PartDeltaEvent,
-    PartEndEvent,
-    PartStartEvent,
     TextPart,
-    TextPartDelta,
+    PartEndEvent,
     ThinkingPart,
-    ThinkingPartDelta,
     ToolCallPart,
-    ToolCallPartDelta,
+    TextPartDelta,
+    PartDeltaEvent,
+    PartStartEvent,
     ToolReturnPart,
+    FinalResultEvent,
+    ThinkingPartDelta,
+    ToolCallPartDelta,
+    FunctionToolCallEvent,
+    FunctionToolResultEvent,
 )
-from pydantic_ai.run import AgentRunResultEvent
 
 from agentex.lib.utils.logging import make_logger
-from agentex.types.reasoning_content_delta import ReasoningContentDelta
-from agentex.types.task_message_content import TextContent
 from agentex.types.task_message_delta import TextDelta
+from agentex.types.tool_request_delta import ToolRequestDelta
 from agentex.types.task_message_update import (
-    StreamTaskMessageDelta,
     StreamTaskMessageDone,
     StreamTaskMessageFull,
+    StreamTaskMessageDelta,
     StreamTaskMessageStart,
 )
+from agentex.types.task_message_content import TextContent
 from agentex.types.tool_request_content import ToolRequestContent
-from agentex.types.tool_request_delta import ToolRequestDelta
 from agentex.types.tool_response_content import ToolResponseContent
+from agentex.types.reasoning_content_delta import ReasoningContentDelta
 
 logger = make_logger(__name__)
 
@@ -98,9 +98,7 @@ def _tool_return_content(result: ToolReturnPart | Any) -> Any:
 
 async def convert_pydantic_ai_to_agentex_events(
     stream_response: AsyncIterator[Any],
-) -> AsyncIterator[
-    StreamTaskMessageStart | StreamTaskMessageDelta | StreamTaskMessageFull | StreamTaskMessageDone
-]:
+) -> AsyncIterator[StreamTaskMessageStart | StreamTaskMessageDelta | StreamTaskMessageFull | StreamTaskMessageDone]:
     """Convert a Pydantic AI agent event stream into Agentex stream events.
 
     Mapping:

tests/lib/adk/test_pydantic_ai_async.py

@@ -10,34 +10,33 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass, field
 from typing import Any, AsyncIterator
+from dataclasses import field, dataclass
 
 import pytest
 from pydantic_ai.messages import (
-    FunctionToolResultEvent,
-    PartDeltaEvent,
-    PartEndEvent,
-    PartStartEvent,
-    RetryPromptPart,
     TextPart,
-    TextPartDelta,
+    PartEndEvent,
     ThinkingPart,
-    ThinkingPartDelta,
     ToolCallPart,
+    TextPartDelta,
+    PartDeltaEvent,
+    PartStartEvent,
     ToolReturnPart,
+    RetryPromptPart,
+    ThinkingPartDelta,
+    FunctionToolResultEvent,
 )
 
-from agentex.lib.adk._modules._pydantic_ai_async import stream_pydantic_ai_events
-from agentex.types.reasoning_content import ReasoningContent
-from agentex.types.reasoning_content_delta import ReasoningContentDelta
 from agentex.types.task_message import TaskMessage
+from agentex.types.text_content import TextContent
+from agentex.types.reasoning_content import ReasoningContent
 from agentex.types.task_message_delta import TextDelta
 from agentex.types.task_message_update import StreamTaskMessageDelta
-from agentex.types.text_content import TextContent
 from agentex.types.tool_request_content import ToolRequestContent
 from agentex.types.tool_response_content import ToolResponseContent
-
+from agentex.types.reasoning_content_delta import ReasoningContentDelta
+from agentex.lib.adk._modules._pydantic_ai_async import stream_pydantic_ai_events
 
 TASK_ID = "task_test"
 
@@ -70,9 +69,7 @@ async def __aexit__(self, exc_type, exc_val, exc_tb) -> bool:
 
     async def stream_update(self, update: StreamTaskMessageDelta) -> None:
         if self.closed:
-            raise AssertionError(
-                "stream_update called after close — helper closed the wrong context"
-            )
+            raise AssertionError("stream_update called after close — helper closed the wrong context")
         self.updates.append(update)
 
     async def close(self) -> None:
@@ -85,9 +82,7 @@ class FakeStreamingModule:
     def __init__(self) -> None:
         self.contexts: list[FakeContext] = []
 
-    def streaming_task_message_context(
-        self, *, task_id: str, initial_content: Any
-    ) -> FakeContext:
+    def streaming_task_message_context(self, *, task_id: str, initial_content: Any) -> FakeContext:
         tm = TaskMessage(
             id=f"m{len(self.contexts) + 1}",
             task_id=task_id,
@@ -255,9 +250,7 @@ async def test_empty_thinking_delta_is_skipped(
         await stream_pydantic_ai_events(_aiter(events), TASK_ID)
 
         ctx = streaming.contexts[0]
-        assert _reasoning_deltas(ctx) == [], (
-            "Empty ThinkingPartDelta must not publish a zero-length reasoning delta"
-        )
+        assert _reasoning_deltas(ctx) == [], "Empty ThinkingPartDelta must not publish a zero-length reasoning delta"
         assert ctx.closed is True
 
 
@@ -274,9 +267,7 @@ async def test_tool_call_emits_full_tool_request_message_on_part_end(
             ),
             PartEndEvent(
                 index=1,
-                part=ToolCallPart(
-                    tool_name="get_weather", args='{"city":"Paris"}', tool_call_id="c1"
-                ),
+                part=ToolCallPart(tool_name="get_weather", args='{"city":"Paris"}', tool_call_id="c1"),
             ),
         ]
         await stream_pydantic_ai_events(_aiter(events), TASK_ID)
@@ -299,15 +290,11 @@ async def test_tool_call_with_dict_args_passes_through(
         events = [
             PartStartEvent(
                 index=0,
-                part=ToolCallPart(
-                    tool_name="search", args={"q": "weather"}, tool_call_id="c"
-                ),
+                part=ToolCallPart(tool_name="search", args={"q": "weather"}, tool_call_id="c"),
             ),
             PartEndEvent(
                 index=0,
-                part=ToolCallPart(
-                    tool_name="search", args={"q": "weather"}, tool_call_id="c"
-                ),
+                part=ToolCallPart(tool_name="search", args={"q": "weather"}, tool_call_id="c"),
             ),
         ]
         await stream_pydantic_ai_events(_aiter(events), TASK_ID)
@@ -366,9 +353,7 @@ async def test_tool_return_emits_full_tool_response_message(
         _, messages = fake_adk
         events = [
             FunctionToolResultEvent(
-                part=ToolReturnPart(
-                    tool_name="get_weather", content="Sunny, 72F", tool_call_id="c1"
-                ),
+                part=ToolReturnPart(tool_name="get_weather", content="Sunny, 72F", tool_call_id="c1"),
             ),
         ]
         await stream_pydantic_ai_events(_aiter(events), TASK_ID)
@@ -387,9 +372,7 @@ async def test_tool_return_with_non_string_content_stringifies(
         _, messages = fake_adk
         events = [
             FunctionToolResultEvent(
-                part=ToolReturnPart(
-                    tool_name="t", content={"temp": 72, "sky": "clear"}, tool_call_id="c"
-                ),
+                part=ToolReturnPart(tool_name="t", content={"temp": 72, "sky": "clear"}, tool_call_id="c"),
             ),
         ]
         await stream_pydantic_ai_events(_aiter(events), TASK_ID)
@@ -447,9 +430,7 @@ async def test_text_then_tool_then_text_uses_separate_contexts_in_order(
                 part=ToolCallPart(tool_name="get_weather", args="{}", tool_call_id="c1"),
             ),
             FunctionToolResultEvent(
-                part=ToolReturnPart(
-                    tool_name="get_weather", content="Sunny", tool_call_id="c1"
-                ),
+                part=ToolReturnPart(tool_name="get_weather", content="Sunny", tool_call_id="c1"),
             ),
             # Second model response: more text.
             PartStartEvent(index=0, part=TextPart(content="")),
@@ -458,9 +439,7 @@ async def test_text_then_tool_then_text_uses_separate_contexts_in_order(
         ]
         final = await stream_pydantic_ai_events(_aiter(events), TASK_ID)
 
-        assert len(streaming.contexts) == 2, (
-            "One context per text part — tool calls don't open streaming contexts"
-        )
+        assert len(streaming.contexts) == 2, "One context per text part — tool calls don't open streaming contexts"
         assert all(ctx.closed for ctx in streaming.contexts)
         assert _text_deltas(streaming.contexts[0]) == ["Looking up..."]
         assert _text_deltas(streaming.contexts[1]) == ["It's sunny."]