diff --git a/src/google/adk/artifacts/in_memory_artifact_service.py b/src/google/adk/artifacts/in_memory_artifact_service.py index ab556f2b..f7b07d35 100644 --- a/src/google/adk/artifacts/in_memory_artifact_service.py +++ b/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 diff --git a/tests/unittests/artifacts/test_artifact_service.py b/tests/unittests/artifacts/test_artifact_service.py index f9065959..920aad68 100644 --- a/tests/unittests/artifacts/test_artifact_service.py +++ b/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, + ) diff --git a/tests/unittests/tools/test_agent_tool.py b/tests/unittests/tools/test_agent_tool.py index 0b476966..1743a30f 100644 --- a/tests/unittests/tools/test_agent_tool.py +++ b/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(