Skip to content

Commit fa14dbf

Browse files
gaborfehergaborfeher
authored
feat: Implement a vertex based task store (#752)
Create a task store adapter that uses the Vertex Managed Task Store as a backend. Fixes #751
1 parent 8b647bd commit fa14dbf

14 files changed

+1815
-20
lines changed

‎pyproject.toml‎

Lines changed: 2 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>=1.140.0"]
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]

‎src/a2a/contrib/__init__.py‎

Whitespace-only changes.

‎src/a2a/contrib/tasks/__init__.py‎

Whitespace-only changes.

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

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

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

Lines changed: 209 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,209 @@
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 as vertexai_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.contrib.tasks import vertex_task_converter
17+
from a2a.server.context import ServerCallContext
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+
await self._create(task)
52+
else:
53+
await self._update(previous_task, task)
54+
55+
async def _create(self, sdk_task: Task) -> None:
56+
stored_task = vertex_task_converter.to_stored_task(sdk_task)
57+
await self._client.aio.agent_engines.a2a_tasks.create(
58+
name=self._agent_engine_resource_id,
59+
a2a_task_id=sdk_task.id,
60+
config=vertexai_types.CreateAgentEngineTaskConfig(
61+
context_id=stored_task.context_id,
62+
metadata=stored_task.metadata,
63+
output=stored_task.output,
64+
),
65+
)
66+
67+
def _get_status_change_event(
68+
self, previous_task: Task, task: Task, event_sequence_number: int
69+
) -> vertexai_types.TaskEvent | None:
70+
if task.status.state != previous_task.status.state:
71+
return vertexai_types.TaskEvent(
72+
event_data=vertexai_types.TaskEventData(
73+
state_change=vertexai_types.TaskStateChange(
74+
new_state=vertex_task_converter.to_stored_task_state(
75+
task.status.state
76+
),
77+
),
78+
),
79+
event_sequence_number=event_sequence_number,
80+
)
81+
return None
82+
83+
def _get_metadata_change_event(
84+
self, previous_task: Task, task: Task, event_sequence_number: int
85+
) -> vertexai_types.TaskEvent | None:
86+
if task.metadata != previous_task.metadata:
87+
return vertexai_types.TaskEvent(
88+
event_data=vertexai_types.TaskEventData(
89+
metadata_change=vertexai_types.TaskMetadataChange(
90+
new_metadata=task.metadata,
91+
)
92+
),
93+
event_sequence_number=event_sequence_number,
94+
)
95+
return None
96+
97+
def _get_artifacts_change_event(
98+
self, previous_task: Task, task: Task, event_sequence_number: int
99+
) -> vertexai_types.TaskEvent | None:
100+
if task.artifacts != previous_task.artifacts:
101+
task_artifact_change = vertexai_types.TaskArtifactChange()
102+
event = vertexai_types.TaskEvent(
103+
event_data=vertexai_types.TaskEventData(
104+
output_change=vertexai_types.TaskOutputChange(
105+
task_artifact_change=task_artifact_change
106+
)
107+
),
108+
event_sequence_number=event_sequence_number,
109+
)
110+
task_artifacts = (
111+
{artifact.artifact_id: artifact for artifact in task.artifacts}
112+
if task.artifacts
113+
else {}
114+
)
115+
previous_task_artifacts = (
116+
{
117+
artifact.artifact_id: artifact
118+
for artifact in previous_task.artifacts
119+
}
120+
if previous_task.artifacts
121+
else {}
122+
)
123+
for artifact in previous_task_artifacts.values():
124+
if artifact.artifact_id not in task_artifacts:
125+
if not task_artifact_change.deleted_artifact_ids:
126+
task_artifact_change.deleted_artifact_ids = []
127+
task_artifact_change.deleted_artifact_ids.append(
128+
artifact.artifact_id
129+
)
130+
for artifact in task_artifacts.values():
131+
if artifact.artifact_id not in previous_task_artifacts:
132+
if not task_artifact_change.added_artifacts:
133+
task_artifact_change.added_artifacts = []
134+
task_artifact_change.added_artifacts.append(
135+
vertex_task_converter.to_stored_artifact(artifact)
136+
)
137+
elif artifact != previous_task_artifacts[artifact.artifact_id]:
138+
if not task_artifact_change.updated_artifacts:
139+
task_artifact_change.updated_artifacts = []
140+
task_artifact_change.updated_artifacts.append(
141+
vertex_task_converter.to_stored_artifact(artifact)
142+
)
143+
if task_artifact_change != vertexai_types.TaskArtifactChange():
144+
return event
145+
return None
146+
147+
async def _update(
148+
self, previous_stored_task: vertexai_types.A2aTask, task: Task
149+
) -> None:
150+
previous_task = vertex_task_converter.to_sdk_task(previous_stored_task)
151+
events = []
152+
event_sequence_number = previous_stored_task.next_event_sequence_number
153+
154+
status_event = self._get_status_change_event(
155+
previous_task, task, event_sequence_number
156+
)
157+
if status_event:
158+
events.append(status_event)
159+
event_sequence_number += 1
160+
161+
metadata_event = self._get_metadata_change_event(
162+
previous_task, task, event_sequence_number
163+
)
164+
if metadata_event:
165+
events.append(metadata_event)
166+
event_sequence_number += 1
167+
168+
artifacts_event = self._get_artifacts_change_event(
169+
previous_task, task, event_sequence_number
170+
)
171+
if artifacts_event:
172+
events.append(artifacts_event)
173+
event_sequence_number += 1
174+
175+
if not events:
176+
return
177+
await self._client.aio.agent_engines.a2a_tasks.events.append(
178+
name=self._agent_engine_resource_id + '/a2aTasks/' + task.id,
179+
task_events=events,
180+
)
181+
182+
async def _get_stored_task(
183+
self, task_id: str
184+
) -> vertexai_types.A2aTask | None:
185+
try:
186+
a2a_task = await self._client.aio.agent_engines.a2a_tasks.get(
187+
name=self._agent_engine_resource_id + '/a2aTasks/' + task_id,
188+
)
189+
except genai_errors.APIError as e:
190+
if e.status == 'NOT_FOUND':
191+
logger.debug('Task %s not found in store.', task_id)
192+
return None
193+
raise
194+
return a2a_task
195+
196+
async def get(
197+
self, task_id: str, context: ServerCallContext | None = None
198+
) -> Task | None:
199+
"""Retrieves a task from the database by ID."""
200+
a2a_task = await self._get_stored_task(task_id)
201+
if a2a_task is None:
202+
return None
203+
return vertex_task_converter.to_sdk_task(a2a_task)
204+
205+
async def delete(
206+
self, task_id: str, context: ServerCallContext | None = None
207+
) -> None:
208+
"""The backend doesn't support deleting tasks, so this is not implemented."""
209+
raise NotImplementedError

0 commit comments

Comments
 (0)