feat: implement new methods in in-memory artifact service

* save_artifact with custom_metadata
* list_artifact_versions
* get_artifact_version

PiperOrigin-RevId: 820321444

src/google/adk/artifacts/in_memory_artifact_service.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 from __future__ import annotations
 
+import dataclasses
 import logging
 from typing import Any
 from typing import Optional
@@ -28,6 +29,19 @@ from .base_artifact_service import BaseArtifactService
 logger = logging.getLogger("google_adk." + __name__)
 
 
+@dataclasses.dataclass
+class _ArtifactEntry:
+  """Represents a single version of an artifact stored in memory.
+
+  Attributes:
+    data: The actual data of the artifact.
+    artifact_version: Metadata about this specific version of the artifact.
+  """
+
+  data: types.Part
+  artifact_version: ArtifactVersion
+
+
 class InMemoryArtifactService(BaseArtifactService, BaseModel):
   """An in-memory implementation of the artifact service.
 
@@ -35,7 +49,7 @@ class InMemoryArtifactService(BaseArtifactService, BaseModel):
   testing and development only.
   """
 
-  artifacts: dict[str, list[types.Part]] = Field(default_factory=dict)
+  artifacts: dict[str, list[_ArtifactEntry]] = Field(default_factory=dict)
 
   def _file_has_user_namespace(self, filename: str) -> bool:
     """Checks if the filename has a user namespace.
@@ -87,15 +101,34 @@ class InMemoryArtifactService(BaseArtifactService, BaseModel):
       session_id: Optional[str] = None,
       custom_metadata: Optional[dict[str, Any]] = None,
   ) -> int:
-    # TODO: b/447451270 - Support saving artifact with custom metadata.
-    if custom_metadata:
-      raise NotImplementedError("custom_metadata is not supported yet.")
-
     path = self._artifact_path(app_name, user_id, filename, session_id)
     if path not in self.artifacts:
       self.artifacts[path] = []
     version = len(self.artifacts[path])
-    self.artifacts[path].append(artifact)
+    if self._file_has_user_namespace(filename):
+      canonical_uri = f"memory://apps/{app_name}/users/{user_id}/artifacts/{filename}/versions/{version}"
+    else:
+      canonical_uri = f"memory://apps/{app_name}/users/{user_id}/sessions/{session_id}/artifacts/{filename}/versions/{version}"
+
+    artifact_version = ArtifactVersion(
+        version=version,
+        canonical_uri=canonical_uri,
+    )
+    if custom_metadata:
+      artifact_version.custom_metadata = custom_metadata
+
+    if artifact.inline_data is not None:
+      artifact_version.mime_type = artifact.inline_data.mime_type
+    elif artifact.text is not None:
+      artifact_version.mime_type = "text/plain"
+    elif artifact.file_data is not None:
+      artifact_version.mime_type = artifact.file_data.mime_type
+    else:
+      raise ValueError("Not supported artifact type.")
+
+    self.artifacts[path].append(
+        _ArtifactEntry(data=artifact, artifact_version=artifact_version)
+    )
     return version
 
   @override
@@ -114,7 +147,10 @@ class InMemoryArtifactService(BaseArtifactService, BaseModel):
       return None
     if version is None:
       version = -1
-    return versions[version]
+    try:
+      return versions[version].data
+    except IndexError:
+      return None
 
   @override
   async def list_artifact_keys(
@@ -172,8 +208,11 @@ class InMemoryArtifactService(BaseArtifactService, BaseModel):
       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.")
+    path = self._artifact_path(app_name, user_id, filename, session_id)
+    entries = self.artifacts.get(path)
+    if not entries:
+      return []
+    return [entry.artifact_version for entry in entries]
 
   @override
   async def get_artifact_version(
@@ -185,5 +224,14 @@ class InMemoryArtifactService(BaseArtifactService, BaseModel):
       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.")
+    path = self._artifact_path(app_name, user_id, filename, session_id)
+    entries = self.artifacts.get(path)
+    if not entries:
+      return None
+
+    if version is None:
+      version = -1
+    try:
+      return entries[version].artifact_version
+    except IndexError:
+      return None

tests/unittests/artifacts/test_artifact_service.py

@@ -14,11 +14,14 @@
 
 """Tests for the artifact service."""
 
+from datetime import datetime
 import enum
 from typing import Optional
 from typing import Union
 from unittest import mock
+from unittest.mock import patch
 
+from google.adk.artifacts.base_artifact_service import ArtifactVersion
 from google.adk.artifacts.gcs_artifact_service import GcsArtifactService
 from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
 from google.genai import types
@@ -26,6 +29,9 @@ import pytest
 
 Enum = enum.Enum
 
+# Define a fixed datetime object to be returned by datetime.now()
+FIXED_DATETIME = datetime(2025, 1, 1, 12, 0, 0)
+
 
 class ArtifactServiceType(Enum):
   IN_MEMORY = "IN_MEMORY"
@@ -195,6 +201,15 @@ async def test_save_load_delete(service_type):
       == artifact
   )
 
+  # Attempt to load a version that doesn't exist
+  assert not await artifact_service.load_artifact(
+      app_name=app_name,
+      user_id=user_id,
+      session_id=session_id,
+      filename=filename,
+      version=3,
+  )
+
   await artifact_service.delete_artifact(
       app_name=app_name,
       user_id=user_id,
@@ -322,3 +337,171 @@ async def test_list_keys_preserves_user_prefix():
   # Should contain prefixed names and session file
   expected_keys = ["user:document.pdf", "user:image.png", "session_file.txt"]
   assert sorted(artifact_keys) == sorted(expected_keys)
+
+
+@pytest.mark.asyncio
+async def test_list_artifact_versions_and_get_artifact_version():
+  """Tests listing artifact versions and getting a specific version."""
+  artifact_service = InMemoryArtifactService()
+  app_name = "app0"
+  user_id = "user0"
+  session_id = "123"
+  filename = "filename"
+  versions = [
+      types.Part.from_bytes(
+          data=i.to_bytes(2, byteorder="big"), mime_type="text/plain"
+      )
+      for i in range(4)
+  ]
+
+  with patch(
+      "google.adk.artifacts.base_artifact_service.datetime"
+  ) as mock_datetime:
+    mock_datetime.now.return_value = FIXED_DATETIME
+
+    for i in range(4):
+      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)},
+      )
+
+    artifact_versions = await artifact_service.list_artifact_versions(
+        app_name=app_name,
+        user_id=user_id,
+        session_id=session_id,
+        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(),
+        )
+        for i in range(4)
+    ]
+    assert artifact_versions == expected_artifact_versions
+
+    # Get latest artifact version when version is not specified
+    assert (
+        await artifact_service.get_artifact_version(
+            app_name=app_name,
+            user_id=user_id,
+            session_id=session_id,
+            filename=filename,
+        )
+        == expected_artifact_versions[-1]
+    )
+
+    # Get artifact version by version number
+    assert (
+        await artifact_service.get_artifact_version(
+            app_name=app_name,
+            user_id=user_id,
+            session_id=session_id,
+            filename=filename,
+            version=2,
+        )
+        == expected_artifact_versions[2]
+    )
+
+
+@pytest.mark.asyncio
+async def test_list_artifact_versions_with_user_prefix():
+  """Tests listing artifact versions with user prefix."""
+  artifact_service = InMemoryArtifactService()
+  app_name = "app0"
+  user_id = "user0"
+  session_id = "123"
+  user_scoped_filename = "user:document.pdf"
+  versions = [
+      types.Part.from_bytes(
+          data=i.to_bytes(2, byteorder="big"), mime_type="text/plain"
+      )
+      for i in range(4)
+  ]
+
+  with patch(
+      "google.adk.artifacts.base_artifact_service.datetime"
+  ) as mock_datetime:
+    mock_datetime.now.return_value = FIXED_DATETIME
+
+    for i in range(4):
+      # Save artifacts with "user:" prefix (cross-session artifacts)
+      await artifact_service.save_artifact(
+          app_name=app_name,
+          user_id=user_id,
+          session_id=session_id,
+          filename=user_scoped_filename,
+          artifact=versions[i],
+          custom_metadata={"key": "value" + str(i)},
+      )
+
+    artifact_versions = await artifact_service.list_artifact_versions(
+        app_name=app_name,
+        user_id=user_id,
+        session_id=session_id,
+        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)
+    ]
+    assert artifact_versions == expected_artifact_versions
+
+
+@pytest.mark.asyncio
+async def test_get_artifact_version_artifact_does_not_exist():
+  """Tests getting an artifact version when artifact does not exist."""
+  artifact_service = InMemoryArtifactService()
+  assert not await artifact_service.get_artifact_version(
+      app_name="test_app",
+      user_id="test_user",
+      session_id="session_id",
+      filename="filename",
+  )
+
+
+@pytest.mark.asyncio
+async def test_get_artifact_version_out_of_index():
+  """Tests loading an artifact with an out-of-index version."""
+  artifact_service = InMemoryArtifactService()
+  app_name = "app0"
+  user_id = "user0"
+  session_id = "123"
+  filename = "filename"
+  artifact = types.Part.from_bytes(data=b"test_data", mime_type="text/plain")
+
+  await artifact_service.save_artifact(
+      app_name=app_name,
+      user_id=user_id,
+      session_id=session_id,
+      filename=filename,
+      artifact=artifact,
+  )
+
+  # Attempt to get a version that doesn't exist
+  assert not await artifact_service.get_artifact_version(
+      app_name=app_name,
+      user_id=user_id,
+      session_id=session_id,
+      filename=filename,
+      version=3,
+  )

tests/unittests/tools/test_agent_tool.py

@@ -178,7 +178,8 @@ def test_update_state():
   assert runner.session.state['state_1'] == 'changed_value'
 
 
-def test_update_artifacts():
+@mark.asyncio
+async def test_update_artifacts():
   """The agent tool can read and write artifacts."""
 
   async def before_tool_agent(callback_context: CallbackContext):
@@ -219,12 +220,21 @@ def test_update_artifacts():
   runner = testing_utils.InMemoryRunner(root_agent)
   runner.run('test1')
 
-  artifacts_path = f'test_app/test_user/{runner.session_id}'
-  assert runner.runner.artifact_service.artifacts == {
-      f'{artifacts_path}/artifact_1': [Part.from_text(text='test')],
-      f'{artifacts_path}/artifact_2': [Part.from_text(text='test 2')],
-      f'{artifacts_path}/artifact_3': [Part.from_text(text='test 2 3')],
-  }
+  async def load_artifact(filename: str):
+    return await runner.runner.artifact_service.load_artifact(
+        app_name='test_app',
+        user_id='test_user',
+        session_id=runner.session_id,
+        filename=filename,
+    )
+
+  assert await runner.runner.artifact_service.list_artifact_keys(
+      app_name='test_app', user_id='test_user', session_id=runner.session_id
+  ) == ['artifact_1', 'artifact_2', 'artifact_3']
+
+  assert await load_artifact('artifact_1') == Part.from_text(text='test')
+  assert await load_artifact('artifact_2') == Part.from_text(text='test 2')
+  assert await load_artifact('artifact_3') == Part.from_text(text='test 2 3')
 
 
 @mark.parametrize(