mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
feat: Adds LLM-Backed User Simulator
Details: - Adds the `LlmBackedUserSimulator` which uses an LLM to generate user prompts until it decides that the conversation is complete. - Adds unit tests for the new functionality. PiperOrigin-RevId: 823557910
This commit is contained in:
Keyur Joshi
committed by
Copybara-Service
Copybara-Service
parent
97a224fe46
commit
54c4ecc733
@@ -0,0 +1,249 @@
|
||||
# Copyright 2025 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from google.adk.evaluation import conversation_scenarios
|
||||
from google.adk.evaluation.llm_backed_user_simulator import LlmBackedUserSimulator
|
||||
from google.adk.evaluation.llm_backed_user_simulator import LlmBackedUserSimulatorConfig
|
||||
from google.adk.evaluation.user_simulator import Status
|
||||
from google.adk.events.event import Event
|
||||
from google.genai import types
|
||||
import pytest
|
||||
|
||||
_INPUT_EVENTS = [
|
||||
Event(
|
||||
author="user",
|
||||
content=types.Content(
|
||||
parts=[types.Part(text="Can you help me?")], role="user"
|
||||
),
|
||||
invocation_id="inv1",
|
||||
),
|
||||
Event(
|
||||
author="helpful_assistant",
|
||||
content=types.Content(
|
||||
parts=[
|
||||
types.Part(
|
||||
text="I'll get the user's name and greet them first.",
|
||||
thought=True,
|
||||
),
|
||||
types.Part(
|
||||
function_call=types.FunctionCall(name="get_user_name")
|
||||
),
|
||||
types.Part(
|
||||
function_response=types.FunctionResponse(
|
||||
name="get_user_name",
|
||||
response={"name": "John Doe"},
|
||||
)
|
||||
),
|
||||
types.Part(text="Hi John, what can I do for you?"),
|
||||
],
|
||||
role="model",
|
||||
),
|
||||
invocation_id="inv1",
|
||||
),
|
||||
]
|
||||
|
||||
_INPUT_EVENTS_LONG = _INPUT_EVENTS + [
|
||||
Event(
|
||||
author="user",
|
||||
content=types.Content(
|
||||
parts=[types.Part(text="I need to book a flight.")], role="user"
|
||||
),
|
||||
invocation_id="inv2",
|
||||
),
|
||||
Event(
|
||||
author="helpful_assistant",
|
||||
content=types.Content(
|
||||
parts=[
|
||||
types.Part(
|
||||
text="Sure, what is your departure date and destination?",
|
||||
),
|
||||
],
|
||||
role="model",
|
||||
),
|
||||
invocation_id="inv2",
|
||||
),
|
||||
]
|
||||
|
||||
_EXPECTED_REWRITTEN_DIALOGUE = """user: Can you help me?
|
||||
|
||||
helpful_assistant: Hi John, what can I do for you?"""
|
||||
|
||||
_EXPECTED_REWRITTEN_DIALOGUE_LONG = _EXPECTED_REWRITTEN_DIALOGUE + """
|
||||
|
||||
user: I need to book a flight.
|
||||
|
||||
helpful_assistant: Sure, what is your departure date and destination?"""
|
||||
|
||||
|
||||
class TestHelperMethods:
|
||||
"""Test cases for LlmBackedUserSimulator helper methods."""
|
||||
|
||||
def test_convert_conversation_to_user_sim_pov(self):
|
||||
"""Tests _convert_conversation_to_user_sim_pov method."""
|
||||
rewritten_dialogue = LlmBackedUserSimulator._summarize_conversation(
|
||||
_INPUT_EVENTS
|
||||
)
|
||||
assert rewritten_dialogue == _EXPECTED_REWRITTEN_DIALOGUE
|
||||
rewritten_dialogue = LlmBackedUserSimulator._summarize_conversation(
|
||||
_INPUT_EVENTS_LONG
|
||||
)
|
||||
assert rewritten_dialogue == _EXPECTED_REWRITTEN_DIALOGUE_LONG
|
||||
|
||||
|
||||
async def to_async_iter(items):
|
||||
for item in items:
|
||||
yield item
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_llm_agent(mocker):
|
||||
"""Provides a mock LLM agent."""
|
||||
mock_llm_registry_cls = mocker.patch(
|
||||
"google.adk.evaluation.llm_backed_user_simulator.LLMRegistry"
|
||||
)
|
||||
mock_llm_registry = mocker.MagicMock()
|
||||
mock_llm_registry_cls.return_value = mock_llm_registry
|
||||
mock_agent = mocker.MagicMock()
|
||||
mock_llm_registry.resolve.return_value.return_value = mock_agent
|
||||
return mock_agent
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def conversation_scenario():
|
||||
"""Provides a test conversation scenario."""
|
||||
return conversation_scenarios.ConversationScenario(
|
||||
starting_prompt="Hello", conversation_plan="test plan"
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def simulator(mock_llm_agent, conversation_scenario):
|
||||
"""Provides an LlmBackedUserSimulator instance for testing."""
|
||||
config = LlmBackedUserSimulatorConfig(
|
||||
model="test-model",
|
||||
model_configuration=types.GenerateContentConfig(),
|
||||
)
|
||||
sim = LlmBackedUserSimulator(
|
||||
config=config, conversation_scenario=conversation_scenario
|
||||
)
|
||||
sim._invocation_count = 1 # Bypass starting prompt by default for tests
|
||||
return sim
|
||||
|
||||
|
||||
class TestLlmBackedUserSimulator:
|
||||
"""Test cases for LlmBackedUserSimulator main methods."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_llm_response_return_value(
|
||||
self, simulator, mock_llm_agent, mocker
|
||||
):
|
||||
"""Tests that _get_llm_response returns the full response correctly."""
|
||||
mock_llm_response = mocker.MagicMock()
|
||||
mock_llm_response.content = types.Content(
|
||||
parts=[
|
||||
types.Part(text="some thought", thought=True),
|
||||
types.Part(text="Hello world!"),
|
||||
]
|
||||
)
|
||||
mock_llm_agent.generate_content_async.return_value = to_async_iter(
|
||||
[mock_llm_response]
|
||||
)
|
||||
response = await simulator._get_llm_response(rewritten_dialogue="")
|
||||
assert response == "Hello world!"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_next_user_message_first_invocation(
|
||||
self, simulator, mock_llm_agent, conversation_scenario
|
||||
):
|
||||
"""Tests that the first invocation returns the starting prompt."""
|
||||
simulator._invocation_count = 0 # override testing default
|
||||
next_user_message = await simulator.get_next_user_message(events=[])
|
||||
|
||||
expected_user_message = types.Content(
|
||||
parts=[types.Part(text=conversation_scenario.starting_prompt)],
|
||||
role="user",
|
||||
)
|
||||
assert next_user_message.status == Status.SUCCESS
|
||||
assert next_user_message.user_message == expected_user_message
|
||||
mock_llm_agent.generate_content_async.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_turn_limit_reached(self, conversation_scenario):
|
||||
"""Tests get_next_user_message when the turn limit is reached."""
|
||||
config = LlmBackedUserSimulatorConfig(
|
||||
max_allowed_invocations=1,
|
||||
)
|
||||
simulator = LlmBackedUserSimulator(
|
||||
config=config, conversation_scenario=conversation_scenario
|
||||
)
|
||||
simulator._invocation_count = 1
|
||||
|
||||
next_user_message = await simulator.get_next_user_message(
|
||||
events=_INPUT_EVENTS
|
||||
)
|
||||
|
||||
assert next_user_message.status == Status.TURN_LIMIT_REACHED
|
||||
assert next_user_message.user_message is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_signal_detected(self, simulator, mock_llm_agent, mocker):
|
||||
"""Tests get_next_user_message when the stop signal is detected."""
|
||||
mock_llm_response = mocker.MagicMock()
|
||||
mock_llm_response.content = types.Content(
|
||||
parts=[types.Part(text="Thanks! Bye!</finished>")]
|
||||
)
|
||||
mock_llm_agent.generate_content_async.return_value = to_async_iter(
|
||||
[mock_llm_response]
|
||||
)
|
||||
|
||||
next_user_message = await simulator.get_next_user_message(
|
||||
events=_INPUT_EVENTS
|
||||
)
|
||||
|
||||
assert next_user_message.status == Status.STOP_SIGNAL_DETECTED
|
||||
assert next_user_message.user_message is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_message_generated(self, simulator, mock_llm_agent):
|
||||
"""Tests get_next_user_message when no message is generated."""
|
||||
mock_llm_agent.generate_content_async.return_value = to_async_iter([])
|
||||
|
||||
with pytest.raises(RuntimeError, match="Failed to generate a user message"):
|
||||
await simulator.get_next_user_message(events=_INPUT_EVENTS)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_next_user_message_success(
|
||||
self, simulator, mock_llm_agent, mocker
|
||||
):
|
||||
"""Tests get_next_user_message when the user message is generated successfully."""
|
||||
mock_llm_response = mocker.MagicMock()
|
||||
mock_llm_response.content = types.Content(
|
||||
parts=[types.Part(text="I need to book a flight.")]
|
||||
)
|
||||
mock_llm_agent.generate_content_async.return_value = to_async_iter(
|
||||
[mock_llm_response]
|
||||
)
|
||||
|
||||
next_user_message = await simulator.get_next_user_message(
|
||||
events=_INPUT_EVENTS
|
||||
)
|
||||
|
||||
expected_user_message = types.Content(
|
||||
parts=[types.Part(text="I need to book a flight.")], role="user"
|
||||
)
|
||||
|
||||
assert next_user_message.status == Status.SUCCESS
|
||||
assert next_user_message.user_message == expected_user_message
|
||||
Reference in New Issue
Block a user