Skip to content

Commit baafd8f

Browse files
gaborfehergaborfeher
committed
Implement a vertex based task store
1 parent cced34d commit baafd8f

9 files changed

Lines changed: 2633 additions & 1242 deletions

File tree

‎pyproject.toml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ postgresql = ["sqlalchemy[asyncio,postgresql-asyncpg]>=2.0.0"]
3737
mysql = ["sqlalchemy[asyncio,aiomysql]>=2.0.0"]
3838
signing = ["PyJWT>=2.0.0"]
3939
sqlite = ["sqlalchemy[asyncio,aiosqlite]>=2.0.0"]
40+
vertex = ["google-cloud-aiplatform @ git+https://github.com/googleapis/python-aiplatform.git@7d58adb9b32bce07152e0df2523eb2de32c66740"]
4041

4142
sql = ["a2a-sdk[postgresql,mysql,sqlite]"]
4243

@@ -47,6 +48,7 @@ all = [
4748
"a2a-sdk[grpc]",
4849
"a2a-sdk[telemetry]",
4950
"a2a-sdk[signing]",
51+
"a2a-sdk[vertex]",
5052
]
5153

5254
[project.urls]
@@ -62,6 +64,9 @@ build-backend = "hatchling.build"
6264
[tool.hatch.version]
6365
source = "uv-dynamic-versioning"
6466

67+
[tool.hatch.metadata]
68+
allow-direct-references = true
69+
6570
[tool.hatch.build.targets.wheel]
6671
packages = ["src/a2a"]
6772

‎scripts/run_vertex_tests.sh‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
#!/bin/bash
2+
set -e
3+
4+
for var in VERTEX_PROJECT VERTEX_LOCATION VERTEX_BASE_URL VERTEX_API_VERSION; do
5+
if [ -z "${!var}" ]; then
6+
echo "Error: Environment variable $var is undefined or empty." >&2
7+
exit 1
8+
fi
9+
done
10+
11+
# Get the directory of this script
12+
SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
13+
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"
14+
PYTEST_ARGS=("$@")
15+
16+
uv run pytest -v "${PYTEST_ARGS[@]}" tests/server/tasks/test_vertex_task_store.py tests/server/tasks/test_vertex_task_converter.py

‎src/a2a/server/tasks/vertex_task_converter.py‎

Lines changed: 161 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,161 @@
1+
try:
2+
from vertexai import types
3+
except ImportError as e:
4+
raise ImportError(
5+
'vertex_task_converter requires vertexai. '
6+
'Install with: '
7+
"'pip install a2a-sdk[vertex]'"
8+
) from e
9+
10+
import base64
11+
import json
12+
13+
from a2a.types import (
14+
Artifact,
15+
DataPart,
16+
FilePart,
17+
FileWithBytes,
18+
FileWithUri,
19+
Part,
20+
Task,
21+
TaskState,
22+
TaskStatus,
23+
TextPart,
24+
)
25+
26+
27+
def to_sdk_task_state(stored_state: types.State) -> TaskState:
28+
"""Converts a proto A2aTask.State to a TaskState enum."""
29+
return {
30+
types.State.STATE_UNSPECIFIED: TaskState.unknown,
31+
types.State.SUBMITTED: TaskState.submitted,
32+
types.State.WORKING: TaskState.working,
33+
types.State.COMPLETED: TaskState.completed,
34+
types.State.CANCELLED: TaskState.canceled,
35+
types.State.FAILED: TaskState.failed,
36+
types.State.REJECTED: TaskState.rejected,
37+
types.State.INPUT_REQUIRED: TaskState.input_required,
38+
types.State.AUTH_REQUIRED: TaskState.auth_required,
39+
}.get(stored_state, TaskState.unknown)
40+
41+
42+
def to_stored_task_state(task_state: TaskState) -> types.State:
43+
"""Converts a TaskState enum to a proto A2aTask.State enum value."""
44+
return {
45+
TaskState.unknown: types.State.STATE_UNSPECIFIED,
46+
TaskState.submitted: types.State.SUBMITTED,
47+
TaskState.working: types.State.WORKING,
48+
TaskState.completed: types.State.COMPLETED,
49+
TaskState.canceled: types.State.CANCELLED,
50+
TaskState.failed: types.State.FAILED,
51+
TaskState.rejected: types.State.REJECTED,
52+
TaskState.input_required: types.State.INPUT_REQUIRED,
53+
TaskState.auth_required: types.State.AUTH_REQUIRED,
54+
}.get(task_state, types.State.STATE_UNSPECIFIED)
55+
56+
57+
def to_stored_part(part: Part) -> types.Part:
58+
"""Converts a SDK Part to a proto Part."""
59+
if isinstance(part.root, TextPart):
60+
return types.Part(text=part.root.text)
61+
if isinstance(part.root, DataPart):
62+
data_bytes = json.dumps(part.root.data).encode('utf-8')
63+
return types.Part(
64+
inline_data=types.Blob(
65+
mime_type='application/json', data=data_bytes
66+
)
67+
)
68+
if isinstance(part.root, FilePart):
69+
file_content = part.root.file
70+
if isinstance(file_content, FileWithBytes):
71+
decoded_bytes = base64.b64decode(file_content.bytes)
72+
return types.Part(
73+
inline_data=types.Blob(
74+
mime_type=file_content.mime_type or '', data=decoded_bytes
75+
)
76+
)
77+
if isinstance(file_content, FileWithUri):
78+
return types.Part(
79+
file_data=types.FileData(
80+
mime_type=file_content.mime_type or '',
81+
file_uri=file_content.uri,
82+
)
83+
)
84+
raise ValueError(f'Unsupported part type: {type(part.root)}')
85+
86+
87+
def to_sdk_part(stored_part: types.Part) -> Part:
88+
"""Converts a proto Part to a SDK Part."""
89+
if stored_part.text:
90+
return Part(root=TextPart(text=stored_part.text))
91+
if stored_part.inline_data:
92+
encoded_bytes = base64.b64encode(stored_part.inline_data.data).decode(
93+
'utf-8'
94+
)
95+
return Part(
96+
root=FilePart(
97+
file=FileWithBytes(
98+
mime_type=stored_part.inline_data.mime_type,
99+
bytes=encoded_bytes,
100+
)
101+
)
102+
)
103+
if stored_part.file_data:
104+
return Part(
105+
root=FilePart(
106+
file=FileWithUri(
107+
mime_type=stored_part.file_data.mime_type,
108+
uri=stored_part.file_data.file_uri,
109+
)
110+
)
111+
)
112+
113+
return Part(root=TextPart(text=''))
114+
115+
116+
def to_stored_artifact(artifact: Artifact) -> types.TaskArtifact:
117+
"""Converts a SDK Artifact to a proto TaskArtifact."""
118+
return types.TaskArtifact(
119+
artifact_id=artifact.artifact_id,
120+
parts=[to_stored_part(part) for part in artifact.parts],
121+
)
122+
123+
124+
def to_sdk_artifact(stored_artifact: types.TaskArtifact) -> Artifact:
125+
"""Converts a proto TaskArtifact to a SDK Artifact."""
126+
return Artifact(
127+
artifact_id=stored_artifact.artifact_id,
128+
parts=[to_sdk_part(part) for part in stored_artifact.parts],
129+
)
130+
131+
132+
def to_stored_task(task: Task) -> types.A2aTask:
133+
"""Converts a SDK Task to a proto A2aTask."""
134+
return types.A2aTask(
135+
context_id=task.context_id,
136+
metadata=task.metadata,
137+
state=to_stored_task_state(task.status.state),
138+
output=types.TaskOutput(
139+
artifacts=[
140+
to_stored_artifact(artifact)
141+
for artifact in task.artifacts or []
142+
]
143+
),
144+
)
145+
146+
147+
def to_sdk_task(a2a_task: types.A2aTask) -> Task:
148+
"""Converts a proto A2aTask to a SDK Task."""
149+
return Task(
150+
id=a2a_task.name.split('/')[-1],
151+
context_id=a2a_task.context_id,
152+
status=TaskStatus(state=to_sdk_task_state(a2a_task.state)),
153+
metadata=a2a_task.metadata or {},
154+
artifacts=[
155+
to_sdk_artifact(artifact)
156+
for artifact in a2a_task.output.artifacts or []
157+
]
158+
if a2a_task.output
159+
else [],
160+
history=[],
161+
)

‎src/a2a/server/tasks/vertex_task_store.py‎

Lines changed: 175 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,175 @@
1+
import logging
2+
3+
4+
try:
5+
import vertexai
6+
7+
from google.genai import errors as genai_errors
8+
from vertexai import types
9+
except ImportError as e:
10+
raise ImportError(
11+
'VertexTaskStore requires vertexai. '
12+
'Install with: '
13+
"'pip install a2a-sdk[vertex]'"
14+
) from e
15+
16+
from a2a.server.context import ServerCallContext
17+
from a2a.server.tasks import vertex_task_converter
18+
from a2a.server.tasks.task_store import TaskStore
19+
from a2a.types import Task # Task is the Pydantic model
20+
21+
22+
logger = logging.getLogger(__name__)
23+
24+
25+
class VertexTaskStore(TaskStore):
26+
"""Implementation of TaskStore using Vertex AI Agent Engine Task Store.
27+
28+
Stores task objects in Vertex AI Agent Engine Task Store.
29+
"""
30+
31+
def __init__(
32+
self,
33+
client: vertexai.Client,
34+
agent_engine_resource_id: str,
35+
) -> None:
36+
"""Initializes the VertexTaskStore.
37+
38+
Args:
39+
client: The Vertex AI client.
40+
agent_engine_resource_id: The resource ID of the agent engine.
41+
"""
42+
self._client = client
43+
self._agent_engine_resource_id = agent_engine_resource_id
44+
45+
async def save(
46+
self, task: Task, context: ServerCallContext | None = None
47+
) -> None:
48+
"""Saves or updates a task in the store."""
49+
previous_task = await self._get_stored_task(task.id)
50+
if previous_task is None:
51+
self._create(task)
52+
else:
53+
self._update(previous_task, task)
54+
55+
def _create(self, sdk_task: Task) -> None:
56+
stored_task = vertex_task_converter.to_stored_task(sdk_task)
57+
self._client.agent_engines.a2a_tasks.create(
58+
name=self._agent_engine_resource_id,
59+
a2a_task_id=sdk_task.id,
60+
config=types.CreateAgentEngineTaskConfig(
61+
context_id=stored_task.context_id,
62+
metadata=stored_task.metadata,
63+
output=stored_task.output,
64+
),
65+
)
66+
67+
def _update(self, previous_stored_task: types.A2aTask, task: Task) -> None:
68+
previous_task = vertex_task_converter.to_sdk_task(previous_stored_task)
69+
events = []
70+
event_sequence_number = previous_stored_task.next_event_sequence_number
71+
if task.status.state != previous_task.status.state:
72+
events.append(
73+
types.TaskEvent(
74+
event_data=types.TaskEventData(
75+
state_change=types.TaskStateChange(
76+
new_state=vertex_task_converter.to_stored_task_state(
77+
task.status.state
78+
),
79+
),
80+
),
81+
event_sequence_number=event_sequence_number,
82+
),
83+
)
84+
event_sequence_number += 1
85+
if task.metadata != previous_task.metadata:
86+
events.append(
87+
types.TaskEvent(
88+
event_data=types.TaskEventData(
89+
metadata_change=types.TaskMetadataChange(
90+
new_metadata=task.metadata,
91+
)
92+
),
93+
event_sequence_number=event_sequence_number,
94+
),
95+
)
96+
event_sequence_number += 1
97+
if task.artifacts != previous_task.artifacts:
98+
task_artifact_change = types.TaskArtifactChange()
99+
event = types.TaskEvent(
100+
event_data=types.TaskEventData(
101+
output_change=types.TaskOutputChange(
102+
task_artifact_change=task_artifact_change
103+
)
104+
),
105+
event_sequence_number=event_sequence_number,
106+
)
107+
task_artifacts = (
108+
{artifact.artifact_id: artifact for artifact in task.artifacts}
109+
if task.artifacts
110+
else {}
111+
)
112+
previous_task_artifacts = (
113+
{
114+
artifact.artifact_id: artifact
115+
for artifact in previous_task.artifacts
116+
}
117+
if previous_task.artifacts
118+
else {}
119+
)
120+
for artifact in previous_task_artifacts.values():
121+
if artifact.artifact_id not in task_artifacts:
122+
if not task_artifact_change.deleted_artifact_ids:
123+
task_artifact_change.deleted_artifact_ids = []
124+
task_artifact_change.deleted_artifact_ids.append(
125+
artifact.artifact_id
126+
)
127+
for artifact in task_artifacts.values():
128+
if artifact.artifact_id not in previous_task_artifacts:
129+
if not task_artifact_change.added_artifacts:
130+
task_artifact_change.added_artifacts = []
131+
task_artifact_change.added_artifacts.append(
132+
vertex_task_converter.to_stored_artifact(artifact)
133+
)
134+
elif artifact != previous_task_artifacts[artifact.artifact_id]:
135+
if not task_artifact_change.updated_artifacts:
136+
task_artifact_change.updated_artifacts = []
137+
task_artifact_change.updated_artifacts.append(
138+
vertex_task_converter.to_stored_artifact(artifact)
139+
)
140+
if task_artifact_change != types.TaskArtifactChange():
141+
events.append(event)
142+
event_sequence_number += 1
143+
if not events:
144+
return
145+
self._client.agent_engines.a2a_tasks.events.append(
146+
name=self._agent_engine_resource_id + '/a2aTasks/' + task.id,
147+
task_events=events,
148+
)
149+
150+
async def _get_stored_task(self, task_id: str) -> types.A2aTask | None:
151+
try:
152+
a2a_task = self._client.agent_engines.a2a_tasks.get(
153+
name=self._agent_engine_resource_id + '/a2aTasks/' + task_id,
154+
)
155+
except genai_errors.APIError as e:
156+
if e.status == 'NOT_FOUND':
157+
logger.debug('Task %s not found in store.', task_id)
158+
return None
159+
raise
160+
return a2a_task
161+
162+
async def get(
163+
self, task_id: str, context: ServerCallContext | None = None
164+
) -> Task | None:
165+
"""Retrieves a task from the database by ID."""
166+
a2a_task = await self._get_stored_task(task_id)
167+
if a2a_task is None:
168+
return None
169+
return vertex_task_converter.to_sdk_task(a2a_task)
170+
171+
async def delete(
172+
self, task_id: str, context: ServerCallContext | None = None
173+
) -> None:
174+
"""The backend doesn't support deleting tasks, so this is not implemented."""
175+
raise NotImplementedError

0 commit comments

Comments
 (0)