feat: Implement artifact_version related methods in GcsArtifactService

PiperOrigin-RevId: 824646770

src/google/adk/artifacts/gcs_artifact_service.py

@@ -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 (
-      return f"{app_name}/{user_id}/user/{filename}/{version}"
+        f"{self._get_blob_prefix(app_name, user_id, filename, session_id)}/{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}"
 
   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)
+    prefix = self._get_blob_prefix(app_name, user_id, filename, session_id)
-    blobs = self.storage_client.list_blobs(self.bucket, prefix=prefix)
+    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.
+    return await asyncio.to_thread(
-    raise NotImplementedError("list_artifact_versions is not implemented yet.")
+        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.
+    return await asyncio.to_thread(
-    raise NotImplementedError("get_artifact_version is not implemented yet.")
+        self._get_artifact_version_sync,
+        app_name,
+        user_id,
+        session_id,
+        filename,
+        version,
+    )

tests/unittests/artifacts/test_artifact_service.py

@@ -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 = [
+    expected_artifact_versions = []
-        ArtifactVersion(
+    for i in range(4):
-            version=i,
+      metadata = {"key": "value" + str(i)}
-            canonical_uri=(
+      if service_type == ArtifactServiceType.GCS:
-                f"memory://apps/{app_name}/users/{user_id}/sessions/{session_id}/artifacts/{filename}/versions/{i}"
+        uri = (
-            ),
+            f"gs://test_bucket/{app_name}/{user_id}/{session_id}/{filename}/{i}"
-            custom_metadata={"key": "value" + str(i)},
-            mime_type="text/plain",
-            create_time=FIXED_DATETIME.timestamp(),
         )
-        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 = [
+    expected_artifact_versions = []
-        ArtifactVersion(
+    for i in range(4):
-            version=i,
+      metadata = {"key": "value" + str(i)}
-            canonical_uri=(
+      if service_type == ArtifactServiceType.GCS:
-                f"memory://apps/{app_name}/users/{user_id}/artifacts/{user_scoped_filename}/versions/{i}"
+        uri = f"gs://test_bucket/{app_name}/{user_id}/user/{user_scoped_filename}/{i}"
-            ),
+      else:
-            custom_metadata={"key": "value" + str(i)},
+        uri = f"memory://apps/{app_name}/users/{user_id}/artifacts/{user_scoped_filename}/versions/{i}"
-            mime_type="text/plain",
+      expected_artifact_versions.append(
-            create_time=FIXED_DATETIME.timestamp(),
+          ArtifactVersion(
-        )
+              version=i,
-        for i in range(4)
+              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"