encounter/adk-python
0
0
Fork 0
mirror of https://github.com/encounter/adk-python.git synced 2026-07-09 18:19:28 -07:00
Code Issues Packages Projects Releases Wiki Activity

feat: Implement artifact_version related methods in GcsArtifactService

Browse Source 
PiperOrigin-RevId: 824646770
This commit is contained in:
Shangjie Chen
2025-10-27 13:13:07 -07:00
committed by Copybara-Service
parent 1a4261ad4b
commit e194ebb33c
2 changed files with 196 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/google/adk/artifacts/gcs_artifact_service.py
+120 -20
View File
@@ -148,6 +148,23 @@ class GcsArtifactService(BaseArtifactService):
"""
return filename.startswith("user:")
def _get_blob_prefix(
self,
app_name: str,
user_id: str,
filename: str,
session_id: Optional[str] = None,
) -> str:
"""Constructs the blob name prefix in GCS for a given artifact."""
if self._file_has_user_namespace(filename):
return f"{app_name}/{user_id}/user/{filename}"
if session_id is None:
raise ValueError(
"Session ID must be provided for session-scoped artifacts."
)
return f"{app_name}/{user_id}/{session_id}/{filename}"
def _get_blob_name(
self,
app_name: str,
@@ -168,14 +185,9 @@ class GcsArtifactService(BaseArtifactService):
Returns:
The constructed blob name in GCS.
"""
if self._file_has_user_namespace(filename):
return f"{app_name}/{user_id}/user/{filename}/{version}"
if session_id is None:
raise ValueError(
"Session ID must be provided for session-scoped artifacts."
)
return f"{app_name}/{user_id}/{session_id}/{filename}/{version}"
return (
f"{self._get_blob_prefix(app_name, user_id, filename, session_id)}/{version}"
)
def _save_artifact(
self,
@@ -186,10 +198,6 @@ class GcsArtifactService(BaseArtifactService):
artifact: types.Part,
custom_metadata: Optional[dict[str, Any]] = None,
) -> int:
if custom_metadata:
# TODO: b/447451270 - support saving artifact with custom metadata.
raise NotImplementedError("custom_metadata is not supported yet.")
versions = self._list_versions(
app_name=app_name,
user_id=user_id,
@@ -202,6 +210,8 @@ class GcsArtifactService(BaseArtifactService):
app_name, user_id, filename, version, session_id
)
blob = self.bucket.blob(blob_name)
if custom_metadata:
blob.metadata = {k: str(v) for k, v in custom_metadata.items()}
if artifact.inline_data:
blob.upload_from_string(
@@ -211,6 +221,7 @@ class GcsArtifactService(BaseArtifactService):
elif artifact.text:
blob.upload_from_string(
data=artifact.text,
content_type="text/plain",
)
elif artifact.file_data:
raise NotImplementedError(
@@ -265,7 +276,12 @@ class GcsArtifactService(BaseArtifactService):
self.bucket, prefix=session_prefix
)
for blob in session_blobs:
*_, filename, _ = blob.name.split("/")
# blob.name is like session_prefix/filename/version
# or session_prefix/path/to/filename/version
# we need to extract filename including slashes, but remove prefix
# and /version
fn_and_version = blob.name[len(session_prefix) :]
filename = "/".join(fn_and_version.split("/")[:-1])
filenames.add(filename)
user_namespace_prefix = f"{app_name}/{user_id}/user/"
@@ -273,7 +289,9 @@ class GcsArtifactService(BaseArtifactService):
self.bucket, prefix=user_namespace_prefix
)
for blob in user_namespace_blobs:
*_, filename, _ = blob.name.split("/")
# blob.name is like user_namespace_prefix/filename/version
fn_and_version = blob.name[len(user_namespace_prefix) :]
filename = "/".join(fn_and_version.split("/")[:-1])
filenames.add(filename)
return sorted(list(filenames))
@@ -323,14 +341,85 @@ class GcsArtifactService(BaseArtifactService):
artifact.
Returns an empty list if no versions are found.
"""
prefix = self._get_blob_name(app_name, user_id, filename, "", session_id)
blobs = self.storage_client.list_blobs(self.bucket, prefix=prefix)
prefix = self._get_blob_prefix(app_name, user_id, filename, session_id)
blobs = self.storage_client.list_blobs(self.bucket, prefix=f"{prefix}/")
versions = []
for blob in blobs:
*_, version = blob.name.split("/")
versions.append(int(version))
return versions
def _get_artifact_version_sync(
self,
app_name: str,
user_id: str,
session_id: Optional[str],
filename: str,
version: Optional[int] = None,
) -> Optional[ArtifactVersion]:
if version is None:
versions = self._list_versions(
app_name=app_name,
user_id=user_id,
session_id=session_id,
filename=filename,
)
if not versions:
return None
version = max(versions)
blob_name = self._get_blob_name(
app_name, user_id, filename, version, session_id
)
blob = self.bucket.get_blob(blob_name)
if not blob:
return None
canonical_uri = f"gs://{self.bucket_name}/{blob.name}"
return ArtifactVersion(
version=version,
canonical_uri=canonical_uri,
create_time=blob.time_created.timestamp(),
mime_type=blob.content_type,
custom_metadata=blob.metadata if blob.metadata else {},
)
def _list_artifact_versions_sync(
self,
app_name: str,
user_id: str,
session_id: Optional[str],
filename: str,
) -> list[ArtifactVersion]:
"""Lists all versions and their metadata of an artifact."""
prefix = self._get_blob_prefix(app_name, user_id, filename, session_id)
blobs = self.storage_client.list_blobs(self.bucket, prefix=f"{prefix}/")
artifact_versions = []
for blob in blobs:
try:
version = int(blob.name.split("/")[-1])
except ValueError:
logger.warning(
"Skipping blob %s because it does not end with a version number.",
blob.name,
)
continue
canonical_uri = f"gs://{self.bucket_name}/{blob.name}"
av = ArtifactVersion(
version=version,
canonical_uri=canonical_uri,
create_time=blob.time_created.timestamp(),
mime_type=blob.content_type,
custom_metadata=blob.metadata if blob.metadata else {},
)
artifact_versions.append(av)
artifact_versions.sort(key=lambda x: x.version)
return artifact_versions
@override
async def list_artifact_versions(
self,
@@ -340,8 +429,13 @@ class GcsArtifactService(BaseArtifactService):
filename: str,
session_id: Optional[str] = None,
) -> list[ArtifactVersion]:
# TODO: b/447451270 - Support list_artifact_versions.
raise NotImplementedError("list_artifact_versions is not implemented yet.")
return await asyncio.to_thread(
self._list_artifact_versions_sync,
app_name,
user_id,
session_id,
filename,
)
@override
async def get_artifact_version(
@@ -353,5 +447,11 @@ class GcsArtifactService(BaseArtifactService):
session_id: Optional[str] = None,
version: Optional[int] = None,
) -> Optional[ArtifactVersion]:
# TODO: b/447451270 - Support get_artifact_version.
raise NotImplementedError("get_artifact_version is not implemented yet.")
return await asyncio.to_thread(
self._get_artifact_version_sync,
app_name,
user_id,
session_id,
filename,
version,
)
tests/unittests/artifacts/test_artifact_service.py
+76 -37
View File
@@ -16,6 +16,7 @@
from datetime import datetime
import enum
from typing import Any
from typing import Optional
from typing import Union
from unittest import mock
@@ -55,6 +56,8 @@ class MockBlob:
self.name = name
self.content: Optional[bytes] = None
self.content_type: Optional[str] = None
self.time_created = FIXED_DATETIME
self.metadata: dict[str, Any] = {}
def upload_from_string(
self, data: Union[str, bytes], content_type: Optional[str] = None
@@ -119,6 +122,13 @@ class MockBucket:
self.blobs[blob_name] = MockBlob(blob_name)
return self.blobs[blob_name]
def get_blob(self, blob_name: str) -> Optional[MockBlob]:
"""Mocks getting a blob from storage if it exists and has content."""
blob = self.blobs.get(blob_name)
if blob and blob.content is not None:
return blob
return None
class MockClient:
"""Mocks the GCS Client."""
@@ -137,9 +147,11 @@ class MockClient:
"""Mocks listing blobs in a bucket, optionally with a prefix."""
if prefix:
return [
blob for name, blob in bucket.blobs.items() if name.startswith(prefix)
blob
for name, blob in bucket.blobs.items()
if name.startswith(prefix) and blob.content is not None
]
return list(bucket.blobs.values())
return [blob for blob in bucket.blobs.values() if blob.content is not None]
def mock_gcs_artifact_service():
@@ -295,9 +307,12 @@ async def test_list_versions(service_type):
@pytest.mark.asyncio
async def test_list_keys_preserves_user_prefix():
@pytest.mark.parametrize(
"service_type", [ArtifactServiceType.IN_MEMORY, ArtifactServiceType.GCS]
)
async def test_list_keys_preserves_user_prefix(service_type):
"""Tests that list_artifact_keys preserves 'user:' prefix in returned names."""
artifact_service = InMemoryArtifactService()
artifact_service = get_artifact_service(service_type)
artifact = types.Part.from_bytes(data=b"test_data", mime_type="text/plain")
app_name = "app0"
user_id = "user0"
@@ -340,9 +355,12 @@ async def test_list_keys_preserves_user_prefix():
@pytest.mark.asyncio
async def test_list_artifact_versions_and_get_artifact_version():
@pytest.mark.parametrize(
"service_type", [ArtifactServiceType.IN_MEMORY, ArtifactServiceType.GCS]
)
async def test_list_artifact_versions_and_get_artifact_version(service_type):
"""Tests listing artifact versions and getting a specific version."""
artifact_service = InMemoryArtifactService()
artifact_service = get_artifact_service(service_type)
app_name = "app0"
user_id = "user0"
session_id = "123"
@@ -360,13 +378,14 @@ async def test_list_artifact_versions_and_get_artifact_version():
mock_datetime.now.return_value = FIXED_DATETIME
for i in range(4):
custom_metadata = {"key": "value" + str(i)}
await artifact_service.save_artifact(
app_name=app_name,
user_id=user_id,
session_id=session_id,
filename=filename,
artifact=versions[i],
custom_metadata={"key": "value" + str(i)},
custom_metadata=custom_metadata,
)
artifact_versions = await artifact_service.list_artifact_versions(
@@ -376,18 +395,24 @@ async def test_list_artifact_versions_and_get_artifact_version():
filename=filename,
)
expected_artifact_versions = [
ArtifactVersion(
version=i,
canonical_uri=(
f"memory://apps/{app_name}/users/{user_id}/sessions/{session_id}/artifacts/{filename}/versions/{i}"
),
custom_metadata={"key": "value" + str(i)},
mime_type="text/plain",
create_time=FIXED_DATETIME.timestamp(),
expected_artifact_versions = []
for i in range(4):
metadata = {"key": "value" + str(i)}
if service_type == ArtifactServiceType.GCS:
uri = (
f"gs://test_bucket/{app_name}/{user_id}/{session_id}/{filename}/{i}"
)
for i in range(4)
]
else:
uri = f"memory://apps/{app_name}/users/{user_id}/sessions/{session_id}/artifacts/{filename}/versions/{i}"
expected_artifact_versions.append(
ArtifactVersion(
version=i,
canonical_uri=uri,
custom_metadata=metadata,
mime_type="text/plain",
create_time=FIXED_DATETIME.timestamp(),
)
)
assert artifact_versions == expected_artifact_versions
# Get latest artifact version when version is not specified
@@ -415,9 +440,12 @@ async def test_list_artifact_versions_and_get_artifact_version():
@pytest.mark.asyncio
async def test_list_artifact_versions_with_user_prefix():
@pytest.mark.parametrize(
"service_type", [ArtifactServiceType.IN_MEMORY, ArtifactServiceType.GCS]
)
async def test_list_artifact_versions_with_user_prefix(service_type):
"""Tests listing artifact versions with user prefix."""
artifact_service = InMemoryArtifactService()
artifact_service = get_artifact_service(service_type)
app_name = "app0"
user_id = "user0"
session_id = "123"
@@ -435,6 +463,7 @@ async def test_list_artifact_versions_with_user_prefix():
mock_datetime.now.return_value = FIXED_DATETIME
for i in range(4):
custom_metadata = {"key": "value" + str(i)}
# Save artifacts with "user:" prefix (cross-session artifacts)
await artifact_service.save_artifact(
app_name=app_name,
@@ -442,7 +471,7 @@ async def test_list_artifact_versions_with_user_prefix():
session_id=session_id,
filename=user_scoped_filename,
artifact=versions[i],
custom_metadata={"key": "value" + str(i)},
custom_metadata=custom_metadata,
)
artifact_versions = await artifact_service.list_artifact_versions(
@@ -452,25 +481,32 @@ async def test_list_artifact_versions_with_user_prefix():
filename=user_scoped_filename,
)
expected_artifact_versions = [
ArtifactVersion(
version=i,
canonical_uri=(
f"memory://apps/{app_name}/users/{user_id}/artifacts/{user_scoped_filename}/versions/{i}"
),
custom_metadata={"key": "value" + str(i)},
mime_type="text/plain",
create_time=FIXED_DATETIME.timestamp(),
)
for i in range(4)
]
expected_artifact_versions = []
for i in range(4):
metadata = {"key": "value" + str(i)}
if service_type == ArtifactServiceType.GCS:
uri = f"gs://test_bucket/{app_name}/{user_id}/user/{user_scoped_filename}/{i}"
else:
uri = f"memory://apps/{app_name}/users/{user_id}/artifacts/{user_scoped_filename}/versions/{i}"
expected_artifact_versions.append(
ArtifactVersion(
version=i,
canonical_uri=uri,
custom_metadata=metadata,
mime_type="text/plain",
create_time=FIXED_DATETIME.timestamp(),
)
)
assert artifact_versions == expected_artifact_versions
@pytest.mark.asyncio
async def test_get_artifact_version_artifact_does_not_exist():
@pytest.mark.parametrize(
"service_type", [ArtifactServiceType.IN_MEMORY, ArtifactServiceType.GCS]
)
async def test_get_artifact_version_artifact_does_not_exist(service_type):
"""Tests getting an artifact version when artifact does not exist."""
artifact_service = InMemoryArtifactService()
artifact_service = get_artifact_service(service_type)
assert not await artifact_service.get_artifact_version(
app_name="test_app",
user_id="test_user",
@@ -480,9 +516,12 @@ async def test_get_artifact_version_artifact_does_not_exist():
@pytest.mark.asyncio
async def test_get_artifact_version_out_of_index():
@pytest.mark.parametrize(
"service_type", [ArtifactServiceType.IN_MEMORY, ArtifactServiceType.GCS]
)
async def test_get_artifact_version_out_of_index(service_type):
"""Tests loading an artifact with an out-of-index version."""
artifact_service = InMemoryArtifactService()
artifact_service = get_artifact_service(service_type)
app_name = "app0"
user_id = "user0"
session_id = "123"