Skip to content

Commit 11598fa

Browse files
committed
refactor: simplify task deserialization in _from_orm by using ParseDict and model_validate for mapping ORM models to task types.
1 parent 4e03b9d commit 11598fa

1 file changed

Lines changed: 10 additions & 35 deletions

File tree

src/a2a/server/tasks/database_task_store.py

Lines changed: 10 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -139,42 +139,17 @@ def _to_orm(self, task: Task, owner: str) -> TaskModel:
139139

140140
def _from_orm(self, task_model: TaskModel) -> Task:
141141
"""Maps a SQLAlchemy TaskModel to a Proto Task instance."""
142-
# Data is stored as raw JSON (dicts/lists), so we parse it manually
143-
task = Task(
144-
id=task_model.id,
145-
context_id=task_model.context_id,
146-
)
142+
task_dict = {
143+
'id': task_model.id,
144+
'context_id': task_model.context_id,
145+
'status': task_model.status,
146+
'artifacts': task_model.artifacts,
147+
'history': task_model.history,
148+
'metadata': task_model.task_metadata,
149+
}
147150
if task_model.protocol_version == '1.0':
148-
if task_model.status:
149-
ParseDict(
150-
cast('dict[str, Any]', task_model.status), task.status
151-
)
152-
if task_model.artifacts:
153-
for art_dict in cast(
154-
'list[dict[str, Any]]', task_model.artifacts
155-
):
156-
art = task.artifacts.add()
157-
ParseDict(art_dict, art)
158-
if task_model.history:
159-
for msg_dict in cast(
160-
'list[dict[str, Any]]', task_model.history
161-
):
162-
msg = task.history.add()
163-
ParseDict(msg_dict, msg)
164-
if task_model.task_metadata:
165-
task.metadata.update(
166-
cast('dict[str, Any]', task_model.task_metadata)
167-
)
168-
return task
169-
# Reconstruct legacy task from raw columns (which are dicts/lists here)
170-
legacy_task = types_v03.Task(
171-
id=task_model.id,
172-
context_id=task_model.context_id,
173-
status=cast('dict[str, Any]', task_model.status),
174-
artifacts=cast('list[dict[str, Any]]', task_model.artifacts),
175-
history=cast('list[dict[str, Any]]', task_model.history),
176-
metadata=cast('dict[str, Any]', task_model.task_metadata),
177-
)
151+
return ParseDict(task_dict, Task())
152+
legacy_task = types_v03.Task.model_validate(task_dict)
178153
return conversions.to_core_task(legacy_task)
179154

180155
async def save(

0 commit comments

Comments
 (0)