Function test_auto_to_single

tests/unittests/workflow/test_agent_transfer.py:101–150

(is_resumable: bool)

@pytest.mark.parametrize('is_resumable', [True, False])
def test_auto_to_single(is_resumable: bool):
  response = [
      transfer_call_part('sub_agent_1'),
      'response1',
      'response2',
  ]
  mock_model = testing_utils.MockModel.create(responses=response)
  # root (auto) - sub_agent_1 (single)
  sub_agent_1 = Agent(
      name='sub_agent_1',
      model=mock_model,
      disallow_transfer_to_parent=True,
      disallow_transfer_to_peers=True,
  )
  root_agent = Agent(
      name='root_agent', model=mock_model, sub_agents=[sub_agent_1]
  )
  app = App(
      name='test_app',
      root_agent=root_agent,
      resumability_config=ResumabilityConfig(is_resumable=is_resumable),
  )
  runner = testing_utils.InMemoryRunner(app=app)

  if not is_resumable:
    # Asserts the responses.
    assert testing_utils.simplify_events(runner.run('test1')) == [
        ('root_agent', transfer_call_part('sub_agent_1')),
        ('root_agent', TRANSFER_RESPONSE_PART),
        ('sub_agent_1', 'response1'),
    ]

    # root_agent should still be the current agent, because sub_agent_1 is
    # single.
    assert testing_utils.simplify_events(runner.run('test2')) == [
        ('root_agent', 'response2'),
    ]
  else:
    assert testing_utils.simplify_resumable_app_events(runner.run('test1')) == [
        ('root_agent', transfer_call_part('sub_agent_1')),
        ('root_agent', TRANSFER_RESPONSE_PART),
        ('root_agent', END_OF_AGENT),
        ('sub_agent_1', 'response1'),
        ('sub_agent_1', END_OF_AGENT),
    ]
    # Same session, different invocation.
    assert testing_utils.simplify_resumable_app_events(runner.run('test2')) == [
        ('root_agent', 'response2'),
        ('root_agent', END_OF_AGENT),
    ]


@pytest.mark.parametrize('is_resumable', [True, False])