feat: Implement a vertex based task store for the 1.0 branch (#791)

For #751

Author: gaborfeher

pyproject.toml

@@ -40,6 +40,7 @@ mysql = ["sqlalchemy[asyncio,aiomysql]>=2.0.0"]
 signing = ["PyJWT>=2.0.0"]
 sqlite = ["sqlalchemy[asyncio,aiosqlite]>=2.0.0"]
 db-cli = ["alembic>=1.14.0"]
+vertex = ["google-cloud-aiplatform>=1.140.0"]
 
 sql = ["a2a-sdk[postgresql,mysql,sqlite]"]
 
@@ -51,6 +52,7 @@ all = [
   "a2a-sdk[telemetry]",
   "a2a-sdk[signing]",
   "a2a-sdk[db-cli]",
+  "a2a-sdk[vertex]",
 ]
 
 [project.urls]

src/a2a/contrib/tasks/__init__.py

@@ -0,0 +1,158 @@
+try:
+    from vertexai import types as vertexai_types
+except ImportError as e:
+    raise ImportError(
+        'vertex_task_converter requires vertexai. '
+        'Install with: '
+        "'pip install a2a-sdk[vertex]'"
+    ) from e
+
+import base64
+import json
+
+from a2a.compat.v0_3.types import (
+    Artifact,
+    DataPart,
+    FilePart,
+    FileWithBytes,
+    FileWithUri,
+    Part,
+    Task,
+    TaskState,
+    TaskStatus,
+    TextPart,
+)
+
+
+_TO_SDK_TASK_STATE = {
+    vertexai_types.State.STATE_UNSPECIFIED: TaskState.unknown,
+    vertexai_types.State.SUBMITTED: TaskState.submitted,
+    vertexai_types.State.WORKING: TaskState.working,
+    vertexai_types.State.COMPLETED: TaskState.completed,
+    vertexai_types.State.CANCELLED: TaskState.canceled,
+    vertexai_types.State.FAILED: TaskState.failed,
+    vertexai_types.State.REJECTED: TaskState.rejected,
+    vertexai_types.State.INPUT_REQUIRED: TaskState.input_required,
+    vertexai_types.State.AUTH_REQUIRED: TaskState.auth_required,
+}
+
+_SDK_TO_STORED_TASK_STATE = {v: k for k, v in _TO_SDK_TASK_STATE.items()}
+
+
+def to_sdk_task_state(stored_state: vertexai_types.State) -> TaskState:
+    """Converts a proto A2aTask.State to a TaskState enum."""
+    return _TO_SDK_TASK_STATE.get(stored_state, TaskState.unknown)
+
+
+def to_stored_task_state(task_state: TaskState) -> vertexai_types.State:
+    """Converts a TaskState enum to a proto A2aTask.State enum value."""
+    return _SDK_TO_STORED_TASK_STATE.get(
+        task_state, vertexai_types.State.STATE_UNSPECIFIED
+    )
+
+
+def to_stored_part(part: Part) -> vertexai_types.Part:
+    """Converts a SDK Part to a proto Part."""
+    if isinstance(part.root, TextPart):
+        return vertexai_types.Part(text=part.root.text)
+    if isinstance(part.root, DataPart):
+        data_bytes = json.dumps(part.root.data).encode('utf-8')
+        return vertexai_types.Part(
+            inline_data=vertexai_types.Blob(
+                mime_type='application/json', data=data_bytes
+            )
+        )
+    if isinstance(part.root, FilePart):
+        file_content = part.root.file
+        if isinstance(file_content, FileWithBytes):
+            decoded_bytes = base64.b64decode(file_content.bytes)
+            return vertexai_types.Part(
+                inline_data=vertexai_types.Blob(
+                    mime_type=file_content.mime_type or '', data=decoded_bytes
+                )
+            )
+        if isinstance(file_content, FileWithUri):
+            return vertexai_types.Part(
+                file_data=vertexai_types.FileData(
+                    mime_type=file_content.mime_type or '',
+                    file_uri=file_content.uri,
+                )
+            )
+    raise ValueError(f'Unsupported part type: {type(part.root)}')
+
+
+def to_sdk_part(stored_part: vertexai_types.Part) -> Part:
+    """Converts a proto Part to a SDK Part."""
+    if stored_part.text:
+        return Part(root=TextPart(text=stored_part.text))
+    if stored_part.inline_data:
+        encoded_bytes = base64.b64encode(stored_part.inline_data.data).decode(
+            'utf-8'
+        )
+        return Part(
+            root=FilePart(
+                file=FileWithBytes(
+                    mime_type=stored_part.inline_data.mime_type,
+                    bytes=encoded_bytes,
+                )
+            )
+        )
+    if stored_part.file_data:
+        return Part(
+            root=FilePart(
+                file=FileWithUri(
+                    mime_type=stored_part.file_data.mime_type,
+                    uri=stored_part.file_data.file_uri,
+                )
+            )
+        )
+
+    raise ValueError(f'Unsupported part: {stored_part}')
+
+
+def to_stored_artifact(artifact: Artifact) -> vertexai_types.TaskArtifact:
+    """Converts a SDK Artifact to a proto TaskArtifact."""
+    return vertexai_types.TaskArtifact(
+        artifact_id=artifact.artifact_id,
+        parts=[to_stored_part(part) for part in artifact.parts],
+    )
+
+
+def to_sdk_artifact(stored_artifact: vertexai_types.TaskArtifact) -> Artifact:
+    """Converts a proto TaskArtifact to a SDK Artifact."""
+    return Artifact(
+        artifact_id=stored_artifact.artifact_id,
+        parts=[to_sdk_part(part) for part in stored_artifact.parts],
+    )
+
+
+def to_stored_task(task: Task) -> vertexai_types.A2aTask:
+    """Converts a SDK Task to a proto A2aTask."""
+    return vertexai_types.A2aTask(
+        context_id=task.context_id,
+        metadata=task.metadata,
+        state=to_stored_task_state(task.status.state),
+        output=vertexai_types.TaskOutput(
+            artifacts=[
+                to_stored_artifact(artifact)
+                for artifact in task.artifacts or []
+            ]
+        ),
+    )
+
+
+def to_sdk_task(a2a_task: vertexai_types.A2aTask) -> Task:
+    """Converts a proto A2aTask to a SDK Task."""
+    return Task(
+        id=a2a_task.name.split('/')[-1],
+        context_id=a2a_task.context_id,
+        status=TaskStatus(state=to_sdk_task_state(a2a_task.state)),
+        metadata=a2a_task.metadata or {},
+        artifacts=[
+            to_sdk_artifact(artifact)
+            for artifact in a2a_task.output.artifacts or []
+        ]
+        if a2a_task.output
+        else [],
+        history=[],
+    )