fix: Allow artifact services to accept dictionary representations of types.Part

This change introduces an `ensure_part` helper function that normalizes input to `types.Part`. This allows `save_artifact` methods in `FileArtifactService`, `GcsArtifactService`, and `InMemoryArtifactService` to accept dictionaries, including those with camelCase keys as used by Agentspace, and convert them into proper `types.Part` instances before saving

Close #2886

Co-authored-by: George Weale <gweale@google.com>
PiperOrigin-RevId: 878131948
This commit is contained in:
George Weale
2026-03-03 14:41:54 -08:00
committed by Copybara-Service
parent 2e434ca7be
commit b004da5027
5 changed files with 175 additions and 10 deletions
@@ -16,8 +16,10 @@ from __future__ import annotations
from abc import ABC
from abc import abstractmethod
from datetime import datetime
import logging
from typing import Any
from typing import Optional
from typing import Union
from google.genai import types
from pydantic import alias_generators
@@ -25,6 +27,8 @@ from pydantic import BaseModel
from pydantic import ConfigDict
from pydantic import Field
logger = logging.getLogger("google_adk." + __name__)
class ArtifactVersion(BaseModel):
"""Metadata describing a specific version of an artifact."""
@@ -60,6 +64,26 @@ class ArtifactVersion(BaseModel):
)
def ensure_part(artifact: Union[types.Part, dict[str, Any]]) -> types.Part:
"""Normalizes an artifact to a ``types.Part`` instance.
External callers may provide artifacts as
plain dictionaries with camelCase keys (``inlineData``) instead of properly
deserialized ``types.Part`` objects. ``model_validate`` handles both
camelCase and snake_case dictionaries transparently via Pydantic aliases.
Args:
artifact: A ``types.Part`` instance or a dictionary representation.
Returns:
A validated ``types.Part`` instance.
"""
if isinstance(artifact, dict):
logger.debug("Normalizing artifact dict to types.Part: %s", list(artifact))
return types.Part.model_validate(artifact)
return artifact
class BaseArtifactService(ABC):
"""Abstract base class for artifact services."""
@@ -70,7 +94,7 @@ class BaseArtifactService(ABC):
app_name: str,
user_id: str,
filename: str,
artifact: types.Part,
artifact: Union[types.Part, dict[str, Any]],
session_id: Optional[str] = None,
custom_metadata: Optional[dict[str, Any]] = None,
) -> int:
@@ -84,10 +108,12 @@ class BaseArtifactService(ABC):
app_name: The app name.
user_id: The user ID.
filename: The filename of the artifact.
artifact: The artifact to save. If the artifact consists of `file_data`,
the artifact service assumes its content has been uploaded separately,
and this method will associate the `file_data` with the artifact if
necessary.
artifact: The artifact to save. Accepts a ``types.Part`` instance or a
plain dictionary (camelCase or snake_case keys) which will be
normalized via ``ensure_part``. If the artifact consists of
``file_data``, the artifact service assumes its content has been
uploaded separately, and this method will associate the ``file_data``
with the artifact if necessary.
session_id: The session ID. If `None`, the artifact is user-scoped.
custom_metadata: custom metadata to associate with the artifact.
@@ -22,6 +22,7 @@ from pathlib import PureWindowsPath
import shutil
from typing import Any
from typing import Optional
from typing import Union
from urllib.parse import unquote
from urllib.parse import urlparse
@@ -35,6 +36,7 @@ from typing_extensions import override
from ..errors.input_validation_error import InputValidationError
from .base_artifact_service import ArtifactVersion
from .base_artifact_service import BaseArtifactService
from .base_artifact_service import ensure_part
logger = logging.getLogger("google_adk." + __name__)
@@ -314,7 +316,7 @@ class FileArtifactService(BaseArtifactService):
app_name: str,
user_id: str,
filename: str,
artifact: types.Part,
artifact: Union[types.Part, dict[str, Any]],
session_id: Optional[str] = None,
custom_metadata: Optional[dict[str, Any]] = None,
) -> int:
@@ -339,11 +341,12 @@ class FileArtifactService(BaseArtifactService):
self,
user_id: str,
filename: str,
artifact: types.Part,
artifact: Union[types.Part, dict[str, Any]],
session_id: Optional[str],
custom_metadata: Optional[dict[str, Any]],
) -> int:
"""Saves an artifact to disk and returns its version."""
artifact = ensure_part(artifact)
artifact_dir = self._artifact_dir(
user_id=user_id,
session_id=session_id,
@@ -27,6 +27,7 @@ import asyncio
import logging
from typing import Any
from typing import Optional
from typing import Union
from google.genai import types
from typing_extensions import override
@@ -34,6 +35,7 @@ from typing_extensions import override
from ..errors.input_validation_error import InputValidationError
from .base_artifact_service import ArtifactVersion
from .base_artifact_service import BaseArtifactService
from .base_artifact_service import ensure_part
logger = logging.getLogger("google_adk." + __name__)
@@ -61,7 +63,7 @@ class GcsArtifactService(BaseArtifactService):
app_name: str,
user_id: str,
filename: str,
artifact: types.Part,
artifact: Union[types.Part, dict[str, Any]],
session_id: Optional[str] = None,
custom_metadata: Optional[dict[str, Any]] = None,
) -> int:
@@ -198,9 +200,10 @@ class GcsArtifactService(BaseArtifactService):
user_id: str,
session_id: Optional[str],
filename: str,
artifact: types.Part,
artifact: Union[types.Part, dict[str, Any]],
custom_metadata: Optional[dict[str, Any]] = None,
) -> int:
artifact = ensure_part(artifact)
versions = self._list_versions(
app_name=app_name,
user_id=user_id,
@@ -17,6 +17,7 @@ import dataclasses
import logging
from typing import Any
from typing import Optional
from typing import Union
from google.genai import types
from pydantic import BaseModel
@@ -27,6 +28,7 @@ from . import artifact_util
from ..errors.input_validation_error import InputValidationError
from .base_artifact_service import ArtifactVersion
from .base_artifact_service import BaseArtifactService
from .base_artifact_service import ensure_part
logger = logging.getLogger("google_adk." + __name__)
@@ -99,10 +101,11 @@ class InMemoryArtifactService(BaseArtifactService, BaseModel):
app_name: str,
user_id: str,
filename: str,
artifact: types.Part,
artifact: Union[types.Part, dict[str, Any]],
session_id: Optional[str] = None,
custom_metadata: Optional[dict[str, Any]] = None,
) -> int:
artifact = ensure_part(artifact)
path = self._artifact_path(app_name, user_id, filename, session_id)
if path not in self.artifacts:
self.artifacts[path] = []
@@ -29,6 +29,7 @@ from urllib.parse import unquote
from urllib.parse import urlparse
from google.adk.artifacts.base_artifact_service import ArtifactVersion
from google.adk.artifacts.base_artifact_service import ensure_part
from google.adk.artifacts.file_artifact_service import FileArtifactService
from google.adk.artifacts.gcs_artifact_service import GcsArtifactService
from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
@@ -766,3 +767,132 @@ async def test_file_save_artifact_rejects_absolute_path_within_scope(tmp_path):
filename=str(absolute_in_scope),
artifact=part,
)
class TestEnsurePart:
"""Tests for the ensure_part normalization helper."""
def test_returns_part_unchanged(self):
"""A types.Part instance passes through without modification."""
part = types.Part.from_bytes(data=b"hello", mime_type="text/plain")
result = ensure_part(part)
assert result is part
def test_converts_camel_case_dict(self):
"""A camelCase dict (Agentspace format) is converted to types.Part."""
raw = {"inlineData": {"mimeType": "image/png", "data": "dGVzdA=="}}
result = ensure_part(raw)
assert isinstance(result, types.Part)
assert result.inline_data is not None
assert result.inline_data.mime_type == "image/png"
def test_converts_snake_case_dict(self):
"""A snake_case dict is converted to types.Part."""
raw = {"inline_data": {"mime_type": "text/plain", "data": "aGVsbG8="}}
result = ensure_part(raw)
assert isinstance(result, types.Part)
assert result.inline_data is not None
assert result.inline_data.mime_type == "text/plain"
def test_converts_text_dict(self):
"""A dict with 'text' key is converted to types.Part."""
raw = {"text": "hello world"}
result = ensure_part(raw)
assert isinstance(result, types.Part)
assert result.text == "hello world"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"service_type",
[
ArtifactServiceType.IN_MEMORY,
ArtifactServiceType.GCS,
ArtifactServiceType.FILE,
],
)
async def test_save_artifact_with_camel_case_dict(
service_type, artifact_service_factory
):
"""Artifact services accept camelCase dicts (Agentspace format).
Regression test for https://github.com/google/adk-python/issues/2886
"""
artifact_service = artifact_service_factory(service_type)
app_name = "app0"
user_id = "user0"
session_id = "sess0"
filename = "uploaded.png"
# Simulate what Agentspace sends: a plain dict with camelCase keys.
raw_artifact = {
"inlineData": {
"mimeType": "image/png",
"data": "dGVzdF9pbWFnZV9kYXRh",
}
}
version = await artifact_service.save_artifact(
app_name=app_name,
user_id=user_id,
session_id=session_id,
filename=filename,
artifact=raw_artifact,
)
assert version == 0
loaded = await artifact_service.load_artifact(
app_name=app_name,
user_id=user_id,
session_id=session_id,
filename=filename,
)
assert loaded is not None
assert loaded.inline_data is not None
assert loaded.inline_data.mime_type == "image/png"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"service_type",
[
ArtifactServiceType.IN_MEMORY,
ArtifactServiceType.GCS,
ArtifactServiceType.FILE,
],
)
async def test_save_artifact_with_snake_case_dict(
service_type, artifact_service_factory
):
"""Artifact services accept snake_case dicts."""
artifact_service = artifact_service_factory(service_type)
app_name = "app0"
user_id = "user0"
session_id = "sess0"
filename = "uploaded.txt"
raw_artifact = {
"inline_data": {
"mime_type": "text/plain",
"data": "aGVsbG8=",
}
}
version = await artifact_service.save_artifact(
app_name=app_name,
user_id=user_id,
session_id=session_id,
filename=filename,
artifact=raw_artifact,
)
assert version == 0
loaded = await artifact_service.load_artifact(
app_name=app_name,
user_id=user_id,
session_id=session_id,
filename=filename,
)
assert loaded is not None
assert loaded.inline_data is not None
assert loaded.inline_data.mime_type == "text/plain"