Function test_update_artifacts

tests/unittests/tools/test_agent_tool.py:316–371

The agent tool can read and write artifacts.

()

@mark.asyncio
async def test_update_artifacts():
  """The agent tool can read and write artifacts."""

  async def before_tool_agent(callback_context: CallbackContext):
    # Artifact 1 should be available in the tool agent.
    artifact = await callback_context.load_artifact('artifact_1')
    await callback_context.save_artifact(
        'artifact_2', Part.from_text(text=artifact.text + ' 2')
    )

  tool_agent = SequentialAgent(
      name='tool_agent',
      before_agent_callback=before_tool_agent,
  )

  async def before_main_agent(callback_context: CallbackContext):
    await callback_context.save_artifact(
        'artifact_1', Part.from_text(text='test')
    )

  async def after_main_agent(callback_context: CallbackContext):
    # Artifact 2 should be available after the tool agent.
    artifact_2 = await callback_context.load_artifact('artifact_2')
    await callback_context.save_artifact(
        'artifact_3', Part.from_text(text=artifact_2.text + ' 3')
    )

  mock_model = testing_utils.MockModel.create(
      responses=[function_call_no_schema, 'response2']
  )
  root_agent = Agent(
      name='root_agent',
      before_agent_callback=before_main_agent,
      after_agent_callback=after_main_agent,
      tools=[AgentTool(agent=tool_agent)],
      model=mock_model,
  )

  runner = testing_utils.InMemoryRunner(root_agent)
  runner.run('test1')

  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')