From 82cbd26c9ed244032e2ee7040bdd2a657315e6dd Mon Sep 17 00:00:00 2001 From: sarojrout Date: Wed, 26 Nov 2025 23:55:08 -0800 Subject: [PATCH] feat(web): expose artifact metadata endpoints --- src/google/adk/cli/adk_web_server.py | 44 +++++ .../cli/conformance/adk_web_server_client.py | 42 +++++ .../conformance/test_adk_web_server_client.py | 79 +++++++++ tests/unittests/cli/test_fast_api.py | 165 ++++++++++++++++-- 4 files changed, 311 insertions(+), 19 deletions(-) diff --git a/src/google/adk/cli/adk_web_server.py b/src/google/adk/cli/adk_web_server.py index 45747a52a1..39ec3592c6 100644 --- a/src/google/adk/cli/adk_web_server.py +++ b/src/google/adk/cli/adk_web_server.py @@ -61,6 +61,7 @@ from ..agents.run_config import RunConfig from ..agents.run_config import StreamingMode from ..apps.app import App +from ..artifacts.base_artifact_service import ArtifactVersion from ..artifacts.base_artifact_service import BaseArtifactService from ..auth.credential_service.base_credential_service import BaseCredentialService from ..errors.already_exists_error import AlreadyExistsError @@ -1294,6 +1295,24 @@ async def load_artifact( raise HTTPException(status_code=404, detail="Artifact not found") return artifact + @app.get( + "/apps/{app_name}/users/{user_id}/sessions/{session_id}/artifacts/{artifact_name}/versions/metadata", + response_model=list[ArtifactVersion], + response_model_exclude_none=True, + ) + async def list_artifact_versions_metadata( + app_name: str, + user_id: str, + session_id: str, + artifact_name: str, + ) -> list[ArtifactVersion]: + return await self.artifact_service.list_artifact_versions( + app_name=app_name, + user_id=user_id, + session_id=session_id, + filename=artifact_name, + ) + @app.get( "/apps/{app_name}/users/{user_id}/sessions/{session_id}/artifacts/{artifact_name}/versions/{version_id}", response_model_exclude_none=True, @@ -1316,6 +1335,31 @@ async def load_artifact_version( raise HTTPException(status_code=404, detail="Artifact not found") return artifact + @app.get( + "/apps/{app_name}/users/{user_id}/sessions/{session_id}/artifacts/{artifact_name}/versions/{version_id}/metadata", + response_model=ArtifactVersion, + response_model_exclude_none=True, + ) + async def get_artifact_version_metadata( + app_name: str, + user_id: str, + session_id: str, + artifact_name: str, + version_id: int, + ) -> ArtifactVersion: + artifact_version = await self.artifact_service.get_artifact_version( + app_name=app_name, + user_id=user_id, + session_id=session_id, + filename=artifact_name, + version=version_id, + ) + if not artifact_version: + raise HTTPException( + status_code=404, detail="Artifact version not found" + ) + return artifact_version + @app.get( "/apps/{app_name}/users/{user_id}/sessions/{session_id}/artifacts", response_model_exclude_none=True, diff --git a/src/google/adk/cli/conformance/adk_web_server_client.py b/src/google/adk/cli/conformance/adk_web_server_client.py index 88fe2ead0c..33f2442bf6 100644 --- a/src/google/adk/cli/conformance/adk_web_server_client.py +++ b/src/google/adk/cli/conformance/adk_web_server_client.py @@ -27,6 +27,7 @@ import httpx +from ...artifacts.base_artifact_service import ArtifactVersion from ...events.event import Event from ...sessions.session import Session from ..adk_web_server import RunAgentRequest @@ -265,3 +266,44 @@ async def run_agent( yield Event.model_validate(event_data) else: logger.debug("Non data line received: %s", line) + + async def get_artifact_version_metadata( + self, + *, + app_name: str, + user_id: str, + session_id: str, + artifact_name: str, + version: int, + ) -> ArtifactVersion: + """Retrieve metadata for a specific artifact version.""" + async with self._get_client() as client: + response = await client.get( + ( + f"/apps/{app_name}/users/{user_id}/sessions/{session_id}" + f"/artifacts/{artifact_name}/versions/{version}/metadata" + ) + ) + response.raise_for_status() + return ArtifactVersion.model_validate(response.json()) + + async def list_artifact_versions_metadata( + self, + *, + app_name: str, + user_id: str, + session_id: str, + artifact_name: str, + ) -> list[ArtifactVersion]: + """List metadata for all versions of an artifact.""" + async with self._get_client() as client: + response = await client.get( + ( + f"/apps/{app_name}/users/{user_id}/sessions/{session_id}" + f"/artifacts/{artifact_name}/versions/metadata" + ) + ) + response.raise_for_status() + return [ + ArtifactVersion.model_validate(item) for item in response.json() + ] diff --git a/tests/unittests/cli/conformance/test_adk_web_server_client.py b/tests/unittests/cli/conformance/test_adk_web_server_client.py index b2bfc43c6d..745ce3a07c 100644 --- a/tests/unittests/cli/conformance/test_adk_web_server_client.py +++ b/tests/unittests/cli/conformance/test_adk_web_server_client.py @@ -17,6 +17,7 @@ from unittest.mock import MagicMock from unittest.mock import patch +from google.adk.artifacts.base_artifact_service import ArtifactVersion from google.adk.cli.adk_web_server import RunAgentRequest from google.adk.cli.conformance.adk_web_server_client import AdkWebServerClient from google.adk.events.event import Event @@ -224,6 +225,84 @@ def mock_stream(*_args, **_kwargs): assert events[1].invocation_id == "test_invocation_2" +@pytest.mark.asyncio +async def test_get_artifact_version_metadata(): + client = AdkWebServerClient() + mock_response = MagicMock() + mock_response.json.return_value = { + "version": 2, + "canonicalUri": ( + "artifact://apps/app/users/user/sessions/session/" + "artifacts/report/versions/2" + ), + "customMetadata": {"foo": "bar"}, + "createTime": 123.4, + "mimeType": "text/plain", + } + + with patch("httpx.AsyncClient") as mock_client_class: + mock_client = AsyncMock() + mock_client.get.return_value = mock_response + mock_client_class.return_value = mock_client + + metadata = await client.get_artifact_version_metadata( + app_name="app", + user_id="user", + session_id="session", + artifact_name="report", + version=2, + ) + + assert isinstance(metadata, ArtifactVersion) + assert metadata.version == 2 + assert metadata.custom_metadata == {"foo": "bar"} + mock_client.get.assert_called_once_with( + "/apps/app/users/user/sessions/session/artifacts/report/versions/2/metadata" + ) + mock_response.raise_for_status.assert_called_once() + + +@pytest.mark.asyncio +async def test_list_artifact_versions_metadata(): + client = AdkWebServerClient() + mock_response = MagicMock() + mock_response.json.return_value = [ + { + "version": 0, + "canonicalUri": "artifact://.../versions/0", + "customMetadata": {}, + "createTime": 100.0, + }, + { + "version": 1, + "canonicalUri": "artifact://.../versions/1", + "customMetadata": {"foo": "bar"}, + "createTime": 200.0, + "mimeType": "application/json", + }, + ] + + with patch("httpx.AsyncClient") as mock_client_class: + mock_client = AsyncMock() + mock_client.get.return_value = mock_response + mock_client_class.return_value = mock_client + + metadata_list = await client.list_artifact_versions_metadata( + app_name="app", + user_id="user", + session_id="session", + artifact_name="report", + ) + + assert len(metadata_list) == 2 + assert all(isinstance(item, ArtifactVersion) for item in metadata_list) + assert metadata_list[1].custom_metadata == {"foo": "bar"} + mock_client.get.assert_called_once_with( + "/apps/app/users/user/sessions/session/artifacts/report/versions/metadata" + ) + mock_response.raise_for_status.assert_called_once() + + @pytest.mark.asyncio async def test_close(): client = AdkWebServerClient() diff --git a/tests/unittests/cli/test_fast_api.py b/tests/unittests/cli/test_fast_api.py index a8b1ef2f2f..041b1075ed 100755 --- a/tests/unittests/cli/test_fast_api.py +++ b/tests/unittests/cli/test_fast_api.py @@ -30,6 +30,7 @@ from google.adk.agents.base_agent import BaseAgent from google.adk.agents.run_config import RunConfig from google.adk.apps.app import App +from google.adk.artifacts.base_artifact_service import ArtifactVersion from google.adk.cli.fast_api import get_fast_api_app from google.adk.evaluation.eval_case import EvalCase from google.adk.evaluation.eval_case import Invocation @@ -211,48 +212,111 @@ def mock_session_service(): def mock_artifact_service(): """Create a mock artifact service.""" - # Storage for artifacts - artifacts = {} - class MockArtifactService: + def __init__(self): + self.artifacts: dict[str, list[dict[str, Any]]] = {} + + def _make_key( + self, app_name: str, user_id: str, session_id: Optional[str], filename: str + ) -> str: + return f"{app_name}:{user_id}:{session_id}:{filename}" + + def add_artifact( + self, + *, + app_name: str, + user_id: str, + session_id: str, + filename: str, + artifact: types.Part, + custom_metadata: Optional[dict[str, Any]] = None, + canonical_uri: Optional[str] = None, + mime_type: Optional[str] = None, + ) -> int: + key = self._make_key(app_name, user_id, session_id, filename) + entries = self.artifacts.setdefault(key, []) + version = len(entries) + metadata = ArtifactVersion( + version=version, + canonical_uri=( + canonical_uri + or "artifact://apps/" + f"{app_name}/users/{user_id}/sessions/{session_id}/artifacts/" + f"{filename}/versions/{version}" + ), + custom_metadata=custom_metadata or {}, + ) + if mime_type: + metadata.mime_type = mime_type + entries.append({"artifact": artifact, "metadata": metadata}) + return version + async def load_artifact( - self, app_name, user_id, session_id, filename, version=None + self, + app_name, + user_id, + session_id, + filename, + version=None, ): """Load an artifact by filename.""" - key = f"{app_name}:{user_id}:{session_id}:{filename}" - if key not in artifacts: + key = self._make_key(app_name, user_id, session_id, filename) + entries = self.artifacts.get(key) + if not entries: return None if version is not None: - # Get a specific version - for v in artifacts[key]: - if v["version"] == version: - return v["artifact"] + for entry in entries: + if entry["metadata"].version == version: + return entry["artifact"] return None - # Get the latest version - return sorted(artifacts[key], key=lambda x: x["version"])[-1]["artifact"] + return entries[-1]["artifact"] async def list_artifact_keys(self, app_name, user_id, session_id): """List artifact names for a session.""" prefix = f"{app_name}:{user_id}:{session_id}:" return [ - k.split(":")[-1] for k in artifacts.keys() if k.startswith(prefix) + k.split(":")[-1] for k in self.artifacts.keys() if k.startswith(prefix) ] async def list_versions(self, app_name, user_id, session_id, filename): """List versions of an artifact.""" - key = f"{app_name}:{user_id}:{session_id}:{filename}" - if key not in artifacts: + key = self._make_key(app_name, user_id, session_id, filename) + entries = self.artifacts.get(key) + if not entries: return [] - return [a["version"] for a in artifacts[key]] + return [entry["metadata"].version for entry in entries] + + async def list_artifact_versions( + self, app_name, user_id, session_id, filename + ): + key = self._make_key(app_name, user_id, session_id, filename) + entries = self.artifacts.get(key) + if not entries: + return [] + return [entry["metadata"] for entry in entries] + + async def get_artifact_version( + self, app_name, user_id, session_id, filename, version=None + ): + key = self._make_key(app_name, user_id, session_id, filename) + entries = self.artifacts.get(key) + if not entries: + return None + if version is None: + return entries[-1]["metadata"] + for entry in entries: + if entry["metadata"].version == version: + return entry["metadata"] + return None async def delete_artifact(self, app_name, user_id, session_id, filename): """Delete an artifact.""" - key = f"{app_name}:{user_id}:{session_id}:{filename}" - if key in artifacts: - del artifacts[key] + key = self._make_key(app_name, user_id, session_id, filename) + if key in self.artifacts: + del self.artifacts[key] return MockArtifactService() @@ -810,6 +874,69 @@ def test_list_artifact_names(test_app, create_test_session): logger.info(f"Listed {len(data)} artifacts") +def test_get_artifact_version_metadata( + test_app, create_test_session, mock_artifact_service +): + """Test retrieving metadata for a specific artifact version.""" + info = create_test_session + mock_artifact_service.add_artifact( + app_name=info["app_name"], + user_id=info["user_id"], + session_id=info["session_id"], + filename="report.txt", + artifact=types.Part(text="hello"), + custom_metadata={"foo": "bar"}, + mime_type="text/plain", + ) + + url = ( + f"/apps/{info['app_name']}/users/{info['user_id']}/sessions/" + f"{info['session_id']}/artifacts/report.txt/versions/0/metadata" + ) + response = test_app.get(url) + + assert response.status_code == 200 + data = response.json() + assert data["version"] == 0 + assert data["customMetadata"] == {"foo": "bar"} + assert data["mimeType"] == "text/plain" + + +def test_list_artifact_versions_metadata( + test_app, create_test_session, mock_artifact_service +): + """Test listing metadata for all versions of an artifact.""" + info = create_test_session + mock_artifact_service.add_artifact( + app_name=info["app_name"], + user_id=info["user_id"], + session_id=info["session_id"], + filename="report.txt", + artifact=types.Part(text="v0"), + ) + mock_artifact_service.add_artifact( + app_name=info["app_name"], + user_id=info["user_id"], + session_id=info["session_id"], + filename="report.txt", + artifact=types.Part(text="v1"), + custom_metadata={"foo": "bar"}, + ) + + url = ( + f"/apps/{info['app_name']}/users/{info['user_id']}/sessions/" + f"{info['session_id']}/artifacts/report.txt/versions/metadata" + ) + response = test_app.get(url) + + assert response.status_code == 200 + data = response.json() + assert isinstance(data, list) + assert len(data) == 2 + assert data[1]["version"] == 1 + assert data[1]["customMetadata"] == {"foo": "bar"} + + def test_create_eval_set(test_app, test_session_info): """Test creating an eval set.""" url = f"/apps/{test_session_info['app_name']}/eval_sets/test_eval_set_id"