Fix a typing issue I introduced

Change-Id: I93d2e4accede9e93bcc67bb5b74c9b371ae440a7

Author: mikeas1

examples/google_adk/calendar_agent/adk_agent_executor.py

@@ -1,33 +1,32 @@
 import asyncio
 import logging
-
-from collections.abc import AsyncGenerator
 from collections import namedtuple
+from collections.abc import AsyncGenerator
 from typing import Dict
+from urllib.parse import parse_qs, urlparse
+
+from google.adk import Runner
+from google.adk.auth import AuthConfig
+from google.adk.events import Event
+from google.genai import types
 
-from a2a.server.agent_execution.context import RequestContext
 from a2a.server.agent_execution import AgentExecutor
+from a2a.server.agent_execution.context import RequestContext
 from a2a.server.events.event_queue import EventQueue
+from a2a.server.tasks import TaskUpdater
 from a2a.types import (
     AgentCard,
     FilePart,
     FileWithBytes,
     FileWithUri,
     Part,
-    UnsupportedOperationError,
     TaskState,
     TextPart,
+    UnsupportedOperationError,
 )
-from a2a.utils import TaskUpdater
 from a2a.utils.errors import ServerError
 from a2a.utils.message import new_agent_text_message
 
-from google.adk import Runner
-from google.adk.events import Event
-from google.adk.auth import AuthConfig
-from google.genai import types
-from urllib.parse import urlparse, parse_qs
-
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 

src/a2a/server/tasks/__init__.py

@@ -1,7 +1,7 @@
 from a2a.server.tasks.inmemory_task_store import InMemoryTaskStore
+from a2a.server.tasks.result_aggregator import ResultAggregator
 from a2a.server.tasks.task_manager import TaskManager
 from a2a.server.tasks.task_store import TaskStore
-from a2a.server.tasks.result_aggregator import ResultAggregator
-
+from a2a.server.tasks.task_updater import TaskUpdater
 
-__all__ = ['InMemoryTaskStore', 'TaskManager', 'TaskStore', 'ResultAggregator']
+__all__ = ['InMemoryTaskStore', 'TaskManager', 'TaskStore', 'ResultAggregator', 'TaskUpdater']

src/a2a/server/tasks/task_updater.py

@@ -0,0 +1,85 @@
+import uuid
+
+from a2a.server.events import EventQueue
+from a2a.types import (Artifact, Message, Part, Role, TaskArtifactUpdateEvent,
+                       TaskState, TaskStatus, TaskStatusUpdateEvent)
+
+
+class TaskUpdater:
+    """Helper class for publishing updates to a task."""
+
+    def __init__(self, event_queue: EventQueue, task_id: str, context_id: str):
+        self.event_queue = event_queue
+        self.task_id = task_id
+        self.context_id = context_id
+
+    def update_status(
+        self, state: TaskState, message: Message | None = None, final=False
+    ):
+        """Update the status of the task."""
+        self.event_queue.enqueue_event(
+            TaskStatusUpdateEvent(
+                taskId=self.task_id,
+                contextId=self.context_id,
+                final=final,
+                status=TaskStatus(
+                    state=state,
+                    message=message,
+                ),
+            )
+        )
+
+    def add_artifact(
+        self,
+        parts: list[Part],
+        artifact_id=str(uuid.uuid4()),
+        name: str | None = None,
+    ):
+        """Add an artifact to the task."""
+        self.event_queue.enqueue_event(
+            TaskArtifactUpdateEvent(
+                taskId=self.task_id,
+                contextId=self.context_id,
+                artifact=Artifact(
+                    artifactId=artifact_id,
+                    name=name,
+                    parts=parts,
+                ),
+            )
+        )
+
+    def complete(self, message: Message | None = None):
+        """Mark the task as completed."""
+        self.update_status(
+            TaskState.completed,
+            message=message,
+            final=True,
+        )
+
+    def submit(self, message: Message | None = None):
+        """Mark the task as submitted."""
+        self.update_status(
+            TaskState.submitted,
+            message=message,
+        )
+
+    def start_work(self, message: Message | None = None):
+        """Mark the task as working."""
+        self.update_status(
+            TaskState.working,
+            message=message,
+        )
+
+    def new_agent_message(
+        self, parts: list[Part], final=False, metadata=None
+    ) -> Message:
+        """Create a new message for the task."""
+        return Message(
+            role=Role.agent,
+            taskId=self.task_id,
+            contextId=self.context_id,
+            messageId=str(uuid.uuid4()),
+            metadata=metadata,
+            final=final,
+            parts=parts,
+        )

src/a2a/utils/__init__.py

@@ -9,8 +9,7 @@
     get_text_parts,
     new_agent_text_message,
 )
-from a2a.utils.task import TaskUpdater, new_task
-
+from a2a.utils.task import new_task
 
 __all__ = [
     'append_artifact_to_task',
@@ -21,5 +20,4 @@
     'new_agent_text_message',
     'new_task',
     'new_text_artifact',
-    'TaskUpdater',
 ]

src/a2a/utils/task.py

@@ -1,16 +1,8 @@
 import uuid
 
-from a2a.types import (
-    Artifact,
-    Message,
-    Part,
-    Role,
-    Task,
-    TaskArtifactUpdateEvent,
-    TaskState,
-    TaskStatus,
-    TaskStatusUpdateEvent,
-)
+from a2a.types import (Artifact, Message, Part, Role, Task,
+                       TaskArtifactUpdateEvent, TaskState, TaskStatus,
+                       TaskStatusUpdateEvent)
 
 
 def new_task(request: Message) -> Task:
@@ -22,83 +14,3 @@ def new_task(request: Message) -> Task:
         ),
         history=[request],
     )
-
-
-class TaskUpdater:
-    """Helper class for publishing updates to a task."""
-
-    def __init__(self, event_queue: EventQueue, task_id: str, context_id: str):
-        self.event_queue = event_queue
-        self.task_id = task_id
-        self.context_id = context_id
-
-    def update_status(
-        self, state: TaskState, message: Message | None = None, final=False
-    ):
-        """Update the status of the task."""
-        self.event_queue.enqueue_event(
-            TaskStatusUpdateEvent(
-                taskId=self.task_id,
-                contextId=self.context_id,
-                final=final,
-                status=TaskStatus(
-                    state=state,
-                    message=message,
-                ),
-            )
-        )
-
-    def add_artifact(
-        self,
-        parts: list[Part],
-        artifact_id=str(uuid.uuid4()),
-        name: str | None = None,
-    ):
-        """Add an artifact to the task."""
-        self.event_queue.enqueue_event(
-            TaskArtifactUpdateEvent(
-                taskId=self.task_id,
-                contextId=self.context_id,
-                artifact=Artifact(
-                    artifactId=artifact_id,
-                    name=name,
-                    parts=parts,
-                ),
-            )
-        )
-
-    def complete(self, message: Message | None = None):
-        """Mark the task as completed."""
-        self.update_status(
-            TaskState.completed,
-            message=message,
-            final=True,
-        )
-
-    def submit(self, message: Message | None = None):
-        """Mark the task as submitted."""
-        self.update_status(
-            TaskState.submitted,
-            message=message,
-        )
-
-    def start_work(self, message: Message | None = None):
-        """Mark the task as working."""
-        self.update_status(
-            TaskState.working,
-            message=message,
-        )
-
-    def new_agent_message(
-        self, parts: list[Part], final=False, metadata=None
-    ) -> Message:
-        """Create a new message for the task."""
-        return Message(
-            role=Role.agent,
-            taskId=self.task_id,
-            contextId=self.context_id,
-            messageId=str(uuid.uuid4()),
-            metadata=metadata,
-            final=final,
-            parts=parts,
-        )