fix(sdk): yield types

Signed-off-by: Radek Ježek <radek.jezek@ibm.com>

Author: jezekra1

apps/beeai-cli/src/beeai_cli/commands/agent.py

@@ -320,7 +320,11 @@ async def _run_agent(
                                     if handle_input is None:
                                         raise ValueError("Agent requires input but no input handler provided")
 
-                                    console.print("\n[bold]Agent requires your input[/bold]\n")
+                                    text = ""
+                                    for part in task_status.message.parts if task_status.message else []:
+                                        if isinstance(part.root, TextPart):
+                                            text = part.root.text
+                                    console.print(f"\n[bold]Agent requires your input[/bold]: {text}\n")
                                     user_input = handle_input()
 
                                     # Send the user input back to the agent

apps/beeai-sdk/src/beeai_sdk/a2a/types.py

@@ -3,10 +3,35 @@
 import uuid
 from typing import Literal, TypeAlias
 
-from a2a.types import Artifact, Message, Part, Role, TaskStatus, TextPart
+from a2a.types import (
+    Artifact,
+    DataPart,
+    FilePart,
+    Message,
+    Part,
+    Role,
+    TaskArtifactUpdateEvent,
+    TaskStatus,
+    TaskStatusUpdateEvent,
+    TextPart,
+)
 from pydantic import Field, model_validator
 
-RunYield: TypeAlias = Message | Part | TaskStatus | Artifact | str | None | dict | Exception
+RunYield: TypeAlias = (
+    Message
+    | Part
+    | TaskStatus
+    | Artifact
+    | TextPart
+    | FilePart
+    | DataPart
+    | TaskStatusUpdateEvent
+    | TaskArtifactUpdateEvent
+    | str
+    | None
+    | dict
+    | Exception
+)
 RunYieldResume: TypeAlias = Message | None