mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
fix: Add endpoints to get/list artifact version metadata
This change introduces new FastAPI endpoints in adk_web_server.py and corresponding client methods in adk_web_server_client.py to allow fetching metadata for artifact versions without downloading the artifact content Close #3710 Co-authored-by: George Weale <gweale@google.com> PiperOrigin-RevId: 868217569
This commit is contained in:
committed by
Copybara-Service
co-authored by
George Weale
parent
7c7d25a4a6
commit
e0b9712a49
@@ -1347,6 +1347,24 @@ class AdkWebServer:
|
|||||||
raise HTTPException(status_code=404, detail="Artifact not found")
|
raise HTTPException(status_code=404, detail="Artifact not found")
|
||||||
return artifact
|
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(
|
@app.get(
|
||||||
"/apps/{app_name}/users/{user_id}/sessions/{session_id}/artifacts/{artifact_name}/versions/{version_id}",
|
"/apps/{app_name}/users/{user_id}/sessions/{session_id}/artifacts/{artifact_name}/versions/{version_id}",
|
||||||
response_model_exclude_none=True,
|
response_model_exclude_none=True,
|
||||||
@@ -1416,6 +1434,31 @@ class AdkWebServer:
|
|||||||
)
|
)
|
||||||
return artifact_version
|
return artifact_version
|
||||||
|
|
||||||
|
@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(
|
@app.get(
|
||||||
"/apps/{app_name}/users/{user_id}/sessions/{session_id}/artifacts",
|
"/apps/{app_name}/users/{user_id}/sessions/{session_id}/artifacts",
|
||||||
response_model_exclude_none=True,
|
response_model_exclude_none=True,
|
||||||
|
|||||||
@@ -27,6 +27,7 @@ from typing import Optional
|
|||||||
|
|
||||||
import httpx
|
import httpx
|
||||||
|
|
||||||
|
from ...artifacts.base_artifact_service import ArtifactVersion
|
||||||
from ...events.event import Event
|
from ...events.event import Event
|
||||||
from ...sessions.session import Session
|
from ...sessions.session import Session
|
||||||
from ..adk_web_server import RunAgentRequest
|
from ..adk_web_server import RunAgentRequest
|
||||||
@@ -268,3 +269,38 @@ class AdkWebServerClient:
|
|||||||
yield Event.model_validate(event_data)
|
yield Event.model_validate(event_data)
|
||||||
else:
|
else:
|
||||||
logger.debug("Non data line received: %s", line)
|
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()]
|
||||||
|
|||||||
@@ -17,6 +17,7 @@ from unittest.mock import AsyncMock
|
|||||||
from unittest.mock import MagicMock
|
from unittest.mock import MagicMock
|
||||||
from unittest.mock import patch
|
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.adk_web_server import RunAgentRequest
|
||||||
from google.adk.cli.conformance.adk_web_server_client import AdkWebServerClient
|
from google.adk.cli.conformance.adk_web_server_client import AdkWebServerClient
|
||||||
from google.adk.events.event import Event
|
from google.adk.events.event import Event
|
||||||
@@ -262,6 +263,84 @@ async def test_run_agent_raises_on_streamed_error():
|
|||||||
pass
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
@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
|
@pytest.mark.asyncio
|
||||||
async def test_close():
|
async def test_close():
|
||||||
client = AdkWebServerClient()
|
client = AdkWebServerClient()
|
||||||
|
|||||||
@@ -290,6 +290,48 @@ def mock_artifact_service():
|
|||||||
})
|
})
|
||||||
return version
|
return version
|
||||||
|
|
||||||
|
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:
|
||||||
|
"""Synchronous helper for tests to add artifacts."""
|
||||||
|
key = _artifact_key(app_name, user_id, session_id, filename)
|
||||||
|
entries = artifacts.setdefault(key, [])
|
||||||
|
version = len(entries)
|
||||||
|
artifact_version = ArtifactVersion(
|
||||||
|
version=version,
|
||||||
|
canonical_uri=(
|
||||||
|
canonical_uri
|
||||||
|
or _canonical_uri(
|
||||||
|
app_name, user_id, session_id, filename, version
|
||||||
|
)
|
||||||
|
),
|
||||||
|
custom_metadata=custom_metadata or {},
|
||||||
|
)
|
||||||
|
if mime_type:
|
||||||
|
artifact_version.mime_type = mime_type
|
||||||
|
elif 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
|
||||||
|
|
||||||
|
entries.append({
|
||||||
|
"version": version,
|
||||||
|
"artifact": artifact,
|
||||||
|
"metadata": artifact_version,
|
||||||
|
})
|
||||||
|
return version
|
||||||
|
|
||||||
async def load_artifact(
|
async def load_artifact(
|
||||||
self, app_name, user_id, session_id, filename, version=None
|
self, app_name, user_id, session_id, filename, version=None
|
||||||
):
|
):
|
||||||
@@ -322,6 +364,15 @@ def mock_artifact_service():
|
|||||||
return []
|
return []
|
||||||
return [entry["version"] for entry in artifacts[key]]
|
return [entry["version"] for entry in artifacts[key]]
|
||||||
|
|
||||||
|
async def list_artifact_versions(
|
||||||
|
self, app_name, user_id, session_id, filename
|
||||||
|
):
|
||||||
|
"""List all artifact versions with metadata."""
|
||||||
|
key = _artifact_key(app_name, user_id, session_id, filename)
|
||||||
|
if key not in artifacts:
|
||||||
|
return []
|
||||||
|
return [entry["metadata"] for entry in artifacts[key]]
|
||||||
|
|
||||||
async def delete_artifact(self, app_name, user_id, session_id, filename):
|
async def delete_artifact(self, app_name, user_id, session_id, filename):
|
||||||
"""Delete an artifact."""
|
"""Delete an artifact."""
|
||||||
key = _artifact_key(app_name, user_id, session_id, filename)
|
key = _artifact_key(app_name, user_id, session_id, filename)
|
||||||
@@ -1193,6 +1244,69 @@ def test_save_artifact_returns_500_on_unexpected_error(
|
|||||||
assert response.json()["detail"] == "unexpected failure"
|
assert response.json()["detail"] == "unexpected failure"
|
||||||
|
|
||||||
|
|
||||||
|
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_get_eval_set_result_not_found(test_app):
|
def test_get_eval_set_result_not_found(test_app):
|
||||||
"""Test getting an eval set result that doesn't exist."""
|
"""Test getting an eval set result that doesn't exist."""
|
||||||
url = "/apps/test_app_name/eval_results/test_eval_result_id_not_found"
|
url = "/apps/test_app_name/eval_results/test_eval_result_id_not_found"
|
||||||
|
|||||||
Reference in New Issue
Block a user