|
6 | 6 | HTTP yields. |
7 | 7 |
|
8 | 8 | Text and thinking tokens stream as deltas inside coalesced streaming |
9 | | -contexts. Tool requests and tool results are emitted as full |
10 | | -``adk.messages.create(...)`` calls (Option A — matches the async LangGraph |
11 | | -helper's convention). To stream tool-call argument tokens, see the sync |
12 | | -converter at ``agentex.lib.adk._modules._pydantic_ai_sync`` which yields |
13 | | -``ToolRequestDelta`` events. |
| 9 | +contexts. Tool requests and tool results are posted as open+close pairs |
| 10 | +on a streaming context (the unified surface persists ``initial_content`` |
| 11 | +when a context is closed without deltas). This matches the ``auto_send`` |
| 12 | +convention used by all other async/Temporal harnesses. |
14 | 13 |
|
15 | 14 | Tracing is opt-in via a ``tracing_handler`` parameter — see |
16 | 15 | ``create_pydantic_ai_tracing_handler`` in |
|
19 | 18 |
|
20 | 19 | from __future__ import annotations |
21 | 20 |
|
22 | | -from typing import TYPE_CHECKING, Any |
| 21 | +from typing import TYPE_CHECKING |
23 | 22 |
|
24 | 23 | if TYPE_CHECKING: |
25 | 24 | from agentex.lib.adk._modules._pydantic_ai_tracing import ( |
@@ -49,230 +48,14 @@ async def stream_pydantic_ai_events( |
49 | 48 | more text) return only the final text segment, matching the |
50 | 49 | ``stream_langgraph_events`` convention. |
51 | 50 | """ |
52 | | - # Lazy imports so pydantic-ai isn't required at module load time. |
53 | | - import json |
54 | | - |
55 | | - from pydantic_ai.messages import ( |
56 | | - TextPart, |
57 | | - PartEndEvent, |
58 | | - ThinkingPart, |
59 | | - ToolCallPart, |
60 | | - TextPartDelta, |
61 | | - PartDeltaEvent, |
62 | | - PartStartEvent, |
63 | | - ThinkingPartDelta, |
64 | | - FunctionToolResultEvent, |
| 51 | + from agentex.lib.core.harness.emitter import UnifiedEmitter |
| 52 | + from agentex.lib.adk._modules._pydantic_ai_turn import PydanticAITurn |
| 53 | + |
| 54 | + turn = PydanticAITurn(stream, model=None, tracing_handler=tracing_handler) |
| 55 | + emitter = UnifiedEmitter( |
| 56 | + task_id=task_id, |
| 57 | + trace_id=None, |
| 58 | + parent_span_id=None, |
65 | 59 | ) |
66 | | - |
67 | | - from agentex.lib import adk |
68 | | - from agentex.types.text_content import TextContent |
69 | | - from agentex.types.reasoning_content import ReasoningContent |
70 | | - from agentex.types.task_message_delta import TextDelta |
71 | | - from agentex.types.task_message_update import StreamTaskMessageDelta |
72 | | - from agentex.types.tool_request_content import ToolRequestContent |
73 | | - from agentex.types.tool_response_content import ToolResponseContent |
74 | | - from agentex.types.reasoning_content_delta import ReasoningContentDelta |
75 | | - |
76 | | - text_context = None |
77 | | - reasoning_context = None |
78 | | - final_text = "" |
79 | | - |
80 | | - # Per Pydantic-AI part-index bookkeeping. Part indices restart at 0 on |
81 | | - # each new model response, so we overwrite on PartStartEvent. |
82 | | - part_kind: dict[int, str] = {} |
83 | | - tool_call_info: dict[int, tuple[str, str]] = {} |
84 | | - |
85 | | - async def _close_text(): |
86 | | - nonlocal text_context |
87 | | - if text_context: |
88 | | - await text_context.close() |
89 | | - text_context = None |
90 | | - |
91 | | - async def _close_reasoning(): |
92 | | - nonlocal reasoning_context |
93 | | - if reasoning_context: |
94 | | - await reasoning_context.close() |
95 | | - reasoning_context = None |
96 | | - |
97 | | - try: |
98 | | - async for event in stream: |
99 | | - if isinstance(event, PartStartEvent): |
100 | | - if isinstance(event.part, TextPart): |
101 | | - await _close_reasoning() |
102 | | - await _close_text() |
103 | | - |
104 | | - final_text = "" |
105 | | - text_context = await adk.streaming.streaming_task_message_context( |
106 | | - task_id=task_id, |
107 | | - initial_content=TextContent( |
108 | | - author="agent", |
109 | | - content="", |
110 | | - format="markdown", |
111 | | - ), |
112 | | - ).__aenter__() |
113 | | - part_kind[event.index] = "text" |
114 | | - |
115 | | - # Pydantic AI puts the first streaming chunk in |
116 | | - # PartStartEvent.part.content; surface it as a Delta so it |
117 | | - # actually renders (Start.content is initialization, not body). |
118 | | - if event.part.content: |
119 | | - final_text += event.part.content |
120 | | - await text_context.stream_update( |
121 | | - StreamTaskMessageDelta( |
122 | | - parent_task_message=text_context.task_message, |
123 | | - delta=TextDelta(type="text", text_delta=event.part.content), |
124 | | - type="delta", |
125 | | - ) |
126 | | - ) |
127 | | - |
128 | | - elif isinstance(event.part, ThinkingPart): |
129 | | - await _close_text() |
130 | | - await _close_reasoning() |
131 | | - |
132 | | - reasoning_context = await adk.streaming.streaming_task_message_context( |
133 | | - task_id=task_id, |
134 | | - initial_content=ReasoningContent( |
135 | | - author="agent", |
136 | | - summary=[], |
137 | | - content=[], |
138 | | - type="reasoning", |
139 | | - style="active", |
140 | | - ), |
141 | | - ).__aenter__() |
142 | | - part_kind[event.index] = "reasoning" |
143 | | - |
144 | | - if event.part.content: |
145 | | - await reasoning_context.stream_update( |
146 | | - StreamTaskMessageDelta( |
147 | | - parent_task_message=reasoning_context.task_message, |
148 | | - delta=ReasoningContentDelta( |
149 | | - type="reasoning_content", |
150 | | - content_index=0, |
151 | | - content_delta=event.part.content, |
152 | | - ), |
153 | | - type="delta", |
154 | | - ) |
155 | | - ) |
156 | | - |
157 | | - elif isinstance(event.part, ToolCallPart): |
158 | | - await _close_text() |
159 | | - await _close_reasoning() |
160 | | - tool_call_info[event.index] = ( |
161 | | - event.part.tool_call_id, |
162 | | - event.part.tool_name, |
163 | | - ) |
164 | | - part_kind[event.index] = "tool_call" |
165 | | - |
166 | | - elif isinstance(event, PartDeltaEvent): |
167 | | - kind = part_kind.get(event.index) |
168 | | - if kind == "text" and isinstance(event.delta, TextPartDelta) and text_context: |
169 | | - final_text += event.delta.content_delta |
170 | | - await text_context.stream_update( |
171 | | - StreamTaskMessageDelta( |
172 | | - parent_task_message=text_context.task_message, |
173 | | - delta=TextDelta(type="text", text_delta=event.delta.content_delta), |
174 | | - type="delta", |
175 | | - ) |
176 | | - ) |
177 | | - elif ( |
178 | | - kind == "reasoning" |
179 | | - and isinstance(event.delta, ThinkingPartDelta) |
180 | | - and reasoning_context |
181 | | - and event.delta.content_delta |
182 | | - ): |
183 | | - await reasoning_context.stream_update( |
184 | | - StreamTaskMessageDelta( |
185 | | - parent_task_message=reasoning_context.task_message, |
186 | | - delta=ReasoningContentDelta( |
187 | | - type="reasoning_content", |
188 | | - content_index=0, |
189 | | - content_delta=event.delta.content_delta, |
190 | | - ), |
191 | | - type="delta", |
192 | | - ) |
193 | | - ) |
194 | | - # Tool-call arg deltas: Pydantic AI accumulates them; we |
195 | | - # surface the final args on PartEndEvent below (Option A). |
196 | | - |
197 | | - elif isinstance(event, PartEndEvent): |
198 | | - kind = part_kind.get(event.index) |
199 | | - if kind == "text": |
200 | | - await _close_text() |
201 | | - elif kind == "reasoning": |
202 | | - await _close_reasoning() |
203 | | - elif kind == "tool_call" and isinstance(event.part, ToolCallPart): |
204 | | - tool_call_id, tool_name = tool_call_info.get(event.index, ("", "")) |
205 | | - args = event.part.args |
206 | | - if isinstance(args, str): |
207 | | - try: |
208 | | - args = json.loads(args) if args else {} |
209 | | - except json.JSONDecodeError: |
210 | | - args = {"_raw": args} |
211 | | - elif args is None: |
212 | | - args = {} |
213 | | - await adk.messages.create( |
214 | | - task_id=task_id, |
215 | | - content=ToolRequestContent( |
216 | | - tool_call_id=tool_call_id, |
217 | | - name=tool_name, |
218 | | - arguments=args, |
219 | | - author="agent", |
220 | | - ), |
221 | | - ) |
222 | | - if tracing_handler is not None and tool_call_id: |
223 | | - await tracing_handler.on_tool_start( |
224 | | - tool_call_id=tool_call_id, |
225 | | - tool_name=tool_name, |
226 | | - arguments=args, |
227 | | - ) |
228 | | - |
229 | | - elif isinstance(event, FunctionToolResultEvent): |
230 | | - await _close_text() |
231 | | - await _close_reasoning() |
232 | | - |
233 | | - result = event.part |
234 | | - tool_call_id = result.tool_call_id |
235 | | - tool_name = getattr(result, "tool_name", "") or "" |
236 | | - # Preserve structure for dicts / lists / Pydantic models so the |
237 | | - # UI can render them as JSON, not as Python repr. Matches the |
238 | | - # sync converter's ``_tool_return_content`` helper exactly — |
239 | | - # ``str(content)`` on a dict produces ``"{'k': 'v'}"`` which is |
240 | | - # invalid JSON and unreadable in the UI. |
241 | | - content = getattr(result, "content", None) |
242 | | - content_payload: Any |
243 | | - if content is None: |
244 | | - content_payload = str(result) |
245 | | - elif isinstance(content, (str, int, float, bool, list, dict)): |
246 | | - content_payload = content |
247 | | - elif hasattr(content, "model_dump"): |
248 | | - try: |
249 | | - content_payload = content.model_dump() |
250 | | - except Exception: |
251 | | - content_payload = str(content) |
252 | | - else: |
253 | | - content_payload = str(content) |
254 | | - await adk.messages.create( |
255 | | - task_id=task_id, |
256 | | - content=ToolResponseContent( |
257 | | - tool_call_id=tool_call_id, |
258 | | - name=tool_name, |
259 | | - content=content_payload, |
260 | | - author="agent", |
261 | | - ), |
262 | | - ) |
263 | | - if tracing_handler is not None and tool_call_id: |
264 | | - await tracing_handler.on_tool_end( |
265 | | - tool_call_id=tool_call_id, |
266 | | - result=content_payload, |
267 | | - ) |
268 | | - |
269 | | - # FunctionToolCallEvent / FinalResultEvent / AgentRunResultEvent |
270 | | - # are intentionally ignored — same as the sync converter. |
271 | | - |
272 | | - finally: |
273 | | - if text_context: |
274 | | - await text_context.close() |
275 | | - if reasoning_context: |
276 | | - await reasoning_context.close() |
277 | | - |
278 | | - return final_text |
| 60 | + result = await emitter.auto_send_turn(turn) |
| 61 | + return result.final_text |
0 commit comments