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
@@ -290,6 +290,48 @@ def mock_artifact_service():
|
||||
})
|
||||
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(
|
||||
self, app_name, user_id, session_id, filename, version=None
|
||||
):
|
||||
@@ -322,6 +364,15 @@ def mock_artifact_service():
|
||||
return []
|
||||
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):
|
||||
"""Delete an artifact."""
|
||||
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"
|
||||
|
||||
|
||||
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):
|
||||
"""Test getting an eval set result that doesn't exist."""
|
||||
url = "/apps/test_app_name/eval_results/test_eval_result_id_not_found"
|
||||
|
||||
Reference in New Issue
Block a user