Skip to content

Commit 6fe9958

Browse files
committed
fix(artifacts): Preserve inline_data display_name on artifact load
File and GCS artifact services dropped inline_data.display_name after save/load, so clients saw generic attachment names instead of user-facing filenames. Persist display_name in file metadata and GCS blob metadata, then restore it when rebuilding binary Parts. Fixes #5833
1 parent 7ad7994 commit 6fe9958

3 files changed

Lines changed: 85 additions & 5 deletions

File tree

src/google/adk/artifacts/file_artifact_service.py

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -210,6 +210,12 @@ class FileArtifactVersion(ArtifactVersion):
210210
file_name: str = Field(
211211
description="Original filename supplied by the caller."
212212
)
213+
display_name: Optional[str] = Field(
214+
default=None,
215+
description=(
216+
"User-facing filename from inline_data.display_name when persisted."
217+
),
218+
)
213219

214220

215221
class FileArtifactService(BaseArtifactService):
@@ -391,13 +397,15 @@ def _save_artifact_sync(
391397
stored_filename = artifact_dir.name
392398
content_path = version_dir / stored_filename
393399

400+
display_name: Optional[str] = None
394401
if artifact.inline_data:
395402
content_path.write_bytes(artifact.inline_data.data)
396403
mime_type = (
397404
artifact.inline_data.mime_type
398405
if artifact.inline_data.mime_type
399406
else "application/octet-stream"
400407
)
408+
display_name = artifact.inline_data.display_name
401409
elif artifact.text is not None:
402410
content_path.write_text(artifact.text, encoding="utf-8")
403411
mime_type = None
@@ -419,6 +427,7 @@ def _save_artifact_sync(
419427
version=next_version,
420428
canonical_uri=canonical_uri,
421429
custom_metadata=custom_metadata,
430+
display_name=display_name,
422431
)
423432

424433
logger.debug(
@@ -491,7 +500,13 @@ def _load_artifact_sync(
491500
)
492501
return None
493502
data = content_path.read_bytes()
494-
return types.Part(inline_data=types.Blob(mime_type=mime_type, data=data))
503+
return types.Part(
504+
inline_data=types.Blob(
505+
mime_type=mime_type,
506+
data=data,
507+
display_name=metadata.display_name if metadata else None,
508+
)
509+
)
495510

496511
if not content_path.exists():
497512
logger.warning("Text artifact %s missing at %s", filename, content_path)
@@ -719,13 +734,15 @@ def _write_metadata(
719734
version: int,
720735
canonical_uri: str,
721736
custom_metadata: Optional[dict[str, Any]],
737+
display_name: Optional[str] = None,
722738
) -> None:
723739
"""Persists metadata describing an artifact version."""
724740
metadata = FileArtifactVersion(
725741
file_name=filename,
726742
mime_type=mime_type,
727743
canonical_uri=canonical_uri,
728744
version=version,
745+
display_name=display_name,
729746
# Persist caller supplied metadata for feature parity with other
730747
# artifact services (e.g. GCS).
731748
custom_metadata=dict(custom_metadata or {}),

src/google/adk/artifacts/gcs_artifact_service.py

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,8 @@
3939

4040
logger = logging.getLogger("google_adk." + __name__)
4141

42+
_GCS_DISPLAY_NAME_METADATA_KEY = "adkDisplayName"
43+
4244

4345
class GcsArtifactService(BaseArtifactService):
4446
"""An artifact service implementation using Google Cloud Storage (GCS)."""
@@ -216,8 +218,13 @@ def _save_artifact(
216218
app_name, user_id, filename, version, session_id
217219
)
218220
blob = self.bucket.blob(blob_name)
219-
if custom_metadata:
220-
blob.metadata = {k: str(v) for k, v in custom_metadata.items()}
221+
blob_metadata = {k: str(v) for k, v in (custom_metadata or {}).items()}
222+
if artifact.inline_data and artifact.inline_data.display_name:
223+
blob_metadata[_GCS_DISPLAY_NAME_METADATA_KEY] = (
224+
artifact.inline_data.display_name
225+
)
226+
if blob_metadata:
227+
blob.metadata = blob_metadata
221228

222229
if artifact.inline_data:
223230
blob.upload_from_string(
@@ -268,10 +275,20 @@ def _load_artifact(
268275
artifact_bytes = blob.download_as_bytes()
269276
if not artifact_bytes:
270277
return None
271-
artifact = types.Part.from_bytes(
278+
display_name = None
279+
if blob.metadata:
280+
display_name = blob.metadata.get(_GCS_DISPLAY_NAME_METADATA_KEY)
281+
if display_name:
282+
return types.Part(
283+
inline_data=types.Blob(
284+
mime_type=blob.content_type,
285+
data=artifact_bytes,
286+
display_name=display_name,
287+
)
288+
)
289+
return types.Part.from_bytes(
272290
data=artifact_bytes, mime_type=blob.content_type
273291
)
274-
return artifact
275292

276293
def _list_artifact_keys(
277294
self, app_name: str, user_id: str, session_id: Optional[str]

tests/unittests/artifacts/test_artifact_service.py

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -959,3 +959,49 @@ async def test_save_artifact_with_snake_case_dict(
959959
assert loaded is not None
960960
assert loaded.inline_data is not None
961961
assert loaded.inline_data.mime_type == "text/plain"
962+
963+
964+
@pytest.mark.asyncio
965+
@pytest.mark.parametrize(
966+
"service_type",
967+
[
968+
ArtifactServiceType.IN_MEMORY,
969+
ArtifactServiceType.GCS,
970+
ArtifactServiceType.FILE,
971+
],
972+
)
973+
async def test_load_artifact_preserves_inline_data_display_name(
974+
service_type, artifact_service_factory
975+
):
976+
"""Binary artifact load restores inline_data.display_name after save."""
977+
artifact_service = artifact_service_factory(service_type)
978+
app_name = "app0"
979+
user_id = "user0"
980+
session_id = "sess0"
981+
filename = "artifact.bin"
982+
display_name = "My Report (final).png"
983+
artifact = types.Part(
984+
inline_data=types.Blob(
985+
mime_type="image/png",
986+
data=b"\x89PNG\r\n\x1a\n",
987+
display_name=display_name,
988+
)
989+
)
990+
991+
await artifact_service.save_artifact(
992+
app_name=app_name,
993+
user_id=user_id,
994+
session_id=session_id,
995+
filename=filename,
996+
artifact=artifact,
997+
)
998+
loaded = await artifact_service.load_artifact(
999+
app_name=app_name,
1000+
user_id=user_id,
1001+
session_id=session_id,
1002+
filename=filename,
1003+
)
1004+
1005+
assert loaded is not None
1006+
assert loaded.inline_data is not None
1007+
assert loaded.inline_data.display_name == display_name

0 commit comments

Comments
 (0)