wip refactoring

Author: guglielmo-san

src/a2a/client/__init__.py

@@ -22,7 +22,6 @@
     A2AClientTimeoutError,
     AgentCardResolutionError,
 )
-from a2a.client.helpers import create_text_message_object
 from a2a.client.interceptors import ClientCallInterceptor
 
 
@@ -41,6 +40,5 @@
     'CredentialService',
     'InMemoryContextCredentialStore',
     'create_client',
-    'create_text_message_object',
     'minimal_agent_card',
 ]

src/a2a/client/helpers.py

@@ -1,11 +1,10 @@
 """Helper functions for the A2A client."""
 
 from typing import Any
-from uuid import uuid4
 
 from google.protobuf.json_format import ParseDict
 
-from a2a.types.a2a_pb2 import AgentCard, Message, Part, Role
+from a2a.types.a2a_pb2 import AgentCard
 
 
 def parse_agent_card(agent_card_data: dict[str, Any]) -> AgentCard:
@@ -111,20 +110,3 @@ def _handle_security_compatibility(agent_card_data: dict[str, Any]) -> None:
                 new_scheme_wrapper = {mapped_name: scheme.copy()}
                 scheme.clear()
                 scheme.update(new_scheme_wrapper)
-
-
-def create_text_message_object(
-    role: Role = Role.ROLE_USER, content: str = ''
-) -> Message:
-    """Create a Message object containing a single text Part.
-
-    Args:
-        role: The role of the message sender (user or agent). Defaults to Role.ROLE_USER.
-        content: The text content of the message. Defaults to an empty string.
-
-    Returns:
-        A `Message` object with a new UUID message_id.
-    """
-    return Message(
-        role=role, parts=[Part(text=content)], message_id=str(uuid4())
-    )

src/a2a/helpers/__init__.py

@@ -0,0 +1 @@
+"""Helper functions for the A2A Python SDK."""

src/a2a/utils/agent_card.py src/a2a/helpers/agent_card.pysrc/a2a/utils/agent_card.py renamed to src/a2a/helpers/agent_card.py

@@ -0,0 +1,209 @@
+"""Unified helper functions for creating and handling A2A types."""
+
+import uuid
+from collections.abc import Sequence
+from typing import Any
+
+from google.protobuf.json_format import MessageToDict
+
+from a2a.types.a2a_pb2 import (
+    Artifact,
+    Message,
+    Part,
+    Role,
+    StreamResponse,
+    Task,
+    TaskArtifactUpdateEvent,
+    TaskState,
+    TaskStatus,
+    TaskStatusUpdateEvent,
+)
+
+
+# --- Message Helpers ---
+
+
+def new_message(
+    parts: list[Part],
+    role: Role = Role.ROLE_AGENT,
+    context_id: str | None = None,
+    task_id: str | None = None,
+) -> Message:
+    """Creates a new message containing a list of Parts."""
+    return Message(
+        role=role,
+        parts=parts,
+        message_id=str(uuid.uuid4()),
+        task_id=task_id,
+        context_id=context_id,
+    )
+
+
+def new_text_message(
+    text: str,
+    context_id: str | None = None,
+    task_id: str | None = None,
+    role: Role = Role.ROLE_AGENT,
+) -> Message:
+    """Creates a new message containing a single text Part."""
+    return new_message(
+        parts=[Part(text=text)],
+        role=role,
+        task_id=task_id,
+        context_id=context_id,
+    )
+
+
+def get_message_text(message: Message, delimiter: str = '\n') -> str:
+    """Extracts and joins all text content from a Message's parts."""
+    return delimiter.join(get_text_parts(message.parts))
+
+
+# --- Artifact Helpers ---
+
+
+def new_artifact(
+    parts: list[Part],
+    name: str,
+    description: str | None = None,
+) -> Artifact:
+    """Creates a new Artifact object."""
+    return Artifact(
+        artifact_id=str(uuid.uuid4()),
+        parts=parts,
+        name=name,
+        description=description,
+    )
+
+
+def new_text_artifact(
+    name: str,
+    text: str,
+    description: str | None = None,
+) -> Artifact:
+    """Creates a new Artifact object containing only a single text Part."""
+    return new_artifact(
+        [Part(text=text)],
+        name,
+        description,
+    )
+
+
+def get_artifact_text(artifact: Artifact, delimiter: str = '\n') -> str:
+    """Extracts and joins all text content from an Artifact's parts."""
+    return delimiter.join(get_text_parts(artifact.parts))
+
+
+# --- Task Helpers ---
+
+
+def new_task_from_request(request: Message) -> Task:
+    """Creates a new Task object from an initial user message."""
+    if not request.role:
+        raise TypeError('Message role cannot be None')
+    if not request.parts:
+        raise ValueError('Message parts cannot be empty')
+    for part in request.parts:
+        if part.HasField('text') and not part.text:
+            raise ValueError('Message.text cannot be empty')
+
+    return Task(
+        status=TaskStatus(state=TaskState.TASK_STATE_SUBMITTED),
+        id=request.task_id or str(uuid.uuid4()),
+        context_id=request.context_id or str(uuid.uuid4()),
+        history=[request],
+    )
+
+
+def new_task(
+    task_id: str,
+    context_id: str,
+    state: TaskState,
+    artifacts: list[Artifact] | None = None,
+    history: list[Message] | None = None,
+) -> Task:
+    """Creates a Task object with a specified status."""
+    if not artifacts or not all(isinstance(a, Artifact) for a in artifacts):
+        raise ValueError(
+            'artifacts must be a non-empty list of Artifact objects'
+        )
+
+    if history is None:
+        history = []
+    return Task(
+        status=TaskStatus(state=state),
+        id=task_id,
+        context_id=context_id,
+        artifacts=artifacts,
+        history=history,
+    )
+
+
+# --- Part Helpers ---
+
+
+def get_text_parts(parts: Sequence[Part]) -> list[str]:
+    """Extracts text content from all text Parts."""
+    return [part.text for part in parts if part.HasField('text')]
+
+
+# --- Event & Stream Helpers ---
+
+
+def new_text_status_update_event(
+    task_id: str,
+    context_id: str,
+    state: TaskState,
+    text: str,
+) -> TaskStatusUpdateEvent:
+    """Creates a TaskStatusUpdateEvent with a single text message."""
+    return TaskStatusUpdateEvent(
+        task_id=task_id,
+        context_id=context_id,
+        status=TaskStatus(
+            state=state,
+            message=new_text_message(
+                text=text,
+                role=Role.ROLE_AGENT,
+                context_id=context_id,
+                task_id=task_id,
+            ),
+        ),
+    )
+
+
+def new_text_artifact_update_event(
+    task_id: str,
+    context_id: str,
+    name: str,
+    text: str,
+    append: bool = False,
+    last_chunk: bool = False,
+) -> TaskArtifactUpdateEvent:
+    """Creates a TaskArtifactUpdateEvent with a single text artifact."""
+    return TaskArtifactUpdateEvent(
+        task_id=task_id,
+        context_id=context_id,
+        artifact=Artifact(
+            artifact_id=str(uuid.uuid4()), name=name, parts=[Part(text=text)]
+        ),
+        append=append,
+        last_chunk=last_chunk,
+    )
+
+
+def get_stream_response_text(response: StreamResponse, delimiter: str = '\n') -> str:
+    """Extracts text content from a StreamResponse."""
+    if response.HasField('message'):
+        return get_message_text(response.message, delimiter)
+    elif response.HasField('task'):
+        if response.task.status.HasField('message'):
+            return get_message_text(response.task.status.message, delimiter)
+        return ''
+    elif response.HasField('status_update'):
+        if response.status_update.status.HasField('message'):
+            return get_message_text(response.status_update.status.message, delimiter)
+        return ''
+    elif response.HasField('artifact_update'):
+        return get_artifact_text(response.artifact_update.artifact, delimiter)
+    return ''

src/a2a/helpers/types.py

@@ -4,20 +4,83 @@
 from a2a.server.events.event_queue import Event
 from a2a.server.tasks.task_store import TaskStore
 from a2a.types.a2a_pb2 import (
+    Artifact,
     Message,
     Task,
     TaskArtifactUpdateEvent,
     TaskState,
     TaskStatus,
     TaskStatusUpdateEvent,
 )
-from a2a.utils import append_artifact_to_task
 from a2a.utils.errors import InvalidParamsError
 
 
 logger = logging.getLogger(__name__)
 
 
+def append_artifact_to_task(task: Task, event: TaskArtifactUpdateEvent) -> None:
+    """Helper method for updating a Task object with new artifact data from an event.
+
+    Handles creating the artifacts list if it doesn't exist, adding new artifacts,
+    and appending parts to existing artifacts based on the `append` flag in the event.
+
+    Args:
+        task: The `Task` object to modify.
+        event: The `TaskArtifactUpdateEvent` containing the artifact data.
+    """
+    new_artifact_data: Artifact = event.artifact
+    artifact_id: str = new_artifact_data.artifact_id
+    append_parts: bool = event.append
+
+    existing_artifact: Artifact | None = None
+    existing_artifact_list_index: int | None = None
+
+    # Find existing artifact by its id
+    for i, art in enumerate(task.artifacts):
+        if art.artifact_id == artifact_id:
+            existing_artifact = art
+            existing_artifact_list_index = i
+            break
+
+    if not append_parts:
+        # This represents the first chunk for this artifact index.
+        if existing_artifact_list_index is not None:
+            # Replace the existing artifact entirely with the new data
+            logger.debug(
+                'Replacing artifact at id %s for task %s', artifact_id, task.id
+            )
+            task.artifacts[existing_artifact_list_index].CopyFrom(
+                new_artifact_data
+            )
+        else:
+            # Append the new artifact since no artifact with this index exists yet
+            logger.debug(
+                'Adding new artifact with id %s for task %s',
+                artifact_id,
+                task.id,
+            )
+            task.artifacts.append(new_artifact_data)
+    elif existing_artifact:
+        # Append new parts to the existing artifact's part list
+        logger.debug(
+            'Appending parts to artifact id %s for task %s',
+            artifact_id,
+            task.id,
+        )
+        existing_artifact.parts.extend(new_artifact_data.parts)
+        existing_artifact.metadata.update(
+            dict(new_artifact_data.metadata.items())
+        )
+    else:
+        # We received a chunk to append, but we don't have an existing artifact.
+        # we will ignore this chunk
+        logger.warning(
+            'Received append=True for nonexistent artifact index %s in task %s. Ignoring chunk.',
+            artifact_id,
+            task.id,
+        )
+
+
 class TaskManager:
     """Helps manage a task's lifecycle during execution of a request.