|
16 | 16 | async def _handle_stream( |
17 | 17 | stream: Any, current_task_id: str | None |
18 | 18 | ) -> str | None: |
19 | | - async for event, task in stream: |
20 | | - if not task: |
21 | | - continue |
| 19 | + async for event in stream: |
22 | 20 | if not current_task_id: |
23 | | - current_task_id = task.id |
24 | | - |
25 | | - if event: |
26 | | - if event.HasField('status_update'): |
27 | | - state_name = TaskState.Name(event.status_update.status.state) |
28 | | - print(f'TaskStatusUpdate [state={state_name}]:', end=' ') |
29 | | - if event.status_update.status.HasField('message'): |
30 | | - for part in event.status_update.status.message.parts: |
31 | | - if part.text: |
32 | | - print(part.text, end=' ') |
33 | | - print() |
34 | | - |
35 | | - if ( |
36 | | - event.status_update.status.state |
37 | | - == TaskState.TASK_STATE_COMPLETED |
38 | | - ): |
39 | | - current_task_id = None |
40 | | - print('--- Task Completed ---') |
41 | | - |
42 | | - elif event.HasField('artifact_update'): |
43 | | - print( |
44 | | - f'TaskArtifactUpdate [name={event.artifact_update.artifact.name}]:', |
45 | | - end=' ', |
46 | | - ) |
47 | | - for part in event.artifact_update.artifact.parts: |
| 21 | + if event.HasField('task'): |
| 22 | + current_task_id = event.task.id |
| 23 | + print(f'Task [state={TaskState.Name(event.task.status.state)}]') |
| 24 | + else: |
| 25 | + raise ValueError('No task found in the first event') |
| 26 | + |
| 27 | + if event.HasField('status_update'): |
| 28 | + state_name = TaskState.Name(event.status_update.status.state) |
| 29 | + print(f'TaskStatusUpdate [state={state_name}]:', end=' ') |
| 30 | + if event.status_update.status.HasField('message'): |
| 31 | + for part in event.status_update.status.message.parts: |
48 | 32 | if part.text: |
49 | 33 | print(part.text, end=' ') |
50 | | - print() |
51 | | - |
| 34 | + print() |
| 35 | + if ( |
| 36 | + event.status_update.status.state |
| 37 | + == TaskState.TASK_STATE_COMPLETED |
| 38 | + ): |
| 39 | + current_task_id = None |
| 40 | + print('--- Task Completed ---') |
| 41 | + elif event.HasField('artifact_update'): |
| 42 | + print( |
| 43 | + f'TaskArtifactUpdate [name={event.artifact_update.artifact.name}]:', |
| 44 | + end=' ', |
| 45 | + ) |
| 46 | + for part in event.artifact_update.artifact.parts: |
| 47 | + if part.text: |
| 48 | + print(part.text, end=' ') |
| 49 | + print() |
52 | 50 | return current_task_id |
53 | 51 |
|
54 | 52 |
|
|
0 commit comments