mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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:
committed by
Copybara-Service
parent
2e434ca7be
commit
b004da5027
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user