feat: implement new methods in in-memory artifact service

* save_artifact with custom_metadata
* list_artifact_versions
* get_artifact_version

PiperOrigin-RevId: 820321444
This commit is contained in:
Google Team Member
2025-10-16 13:07:09 -07:00
committed by Copybara-Service
parent 1e1d63f34c
commit 5a543c00df
3 changed files with 259 additions and 18 deletions
+17 -7
View File
@@ -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(