feat: GenAI Client(evals): Add scenario generation function to eval components

PiperOrigin-RevId: 872110328

Author: vertex-sdk-bot

tests/unit/vertexai/genai/replays/test_generate_user_scenarios.py

@@ -0,0 +1,114 @@
+# Copyright 2026 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.
+#
+# pylint: disable=protected-access,bad-continuation,missing-function-docstring
+
+from tests.unit.vertexai.genai.replays import pytest_helper
+from vertexai import types
+import pytest
+
+def test_gen_user_scenarios(client):
+    """Tests that generate_user_scenarios() correctly calls the API and parses the response."""
+    eval_dataset = client.evals.generate_user_scenarios(
+        agents={
+            "booking-agent": types.evals.AgentConfig(
+                agent_id="booking-agent",
+                agent_type="service_agent",
+                description="An agent capable of booking flights and hotels.",
+                instruction="You are a helpful travel assistant. Use tools to find flights.",
+                tools=[
+                    {
+                        "function_declarations": [
+                            {
+                                "name": "search_flights",
+                                "description": "Search for available flights.",
+                            }
+                        ]
+                    }
+                ],
+            )
+        },
+        user_scenario_generation_config=types.evals.UserScenarioGenerationConfig(
+            user_scenario_count=2,
+            simulation_instruction=(
+                "Generate scenarios where the user tries to book a flight but"
+                " changes their mind about the destination."
+            ),
+            environment_data="Today is Monday. Flights to Paris are available.",
+            model_name="gemini-2.5-flash",
+        ),
+        root_agent_id="booking-agent",
+    )
+    assert isinstance(eval_dataset, types.EvaluationDataset)
+    assert len(eval_dataset.eval_cases) == 2
+    assert (
+        eval_dataset.eval_cases[0].user_scenario.starting_prompt
+        == "I want to find a flight from New York to London."
+    )
+    assert (
+        eval_dataset.eval_cases[0].user_scenario.conversation_plan
+        == "Actually, I meant Paris, not London. Please search for flights to Paris."
+    )
+
+pytest_plugins = ("pytest_asyncio",)
+
+@pytest.mark.asyncio
+async def test_gen_user_scenarios_async(client):
+    """Tests that generate_user_scenarios() async correctly calls the API and parses the response."""
+    eval_dataset = await client.aio.evals.generate_user_scenarios(
+        agents={
+            "booking-agent": types.evals.AgentConfig(
+                agent_id="booking-agent",
+                agent_type="service_agent",
+                description="An agent capable of booking flights and hotels.",
+                instruction="You are a helpful travel assistant. Use tools to find flights.",
+                tools=[
+                    {
+                        "function_declarations": [
+                            {
+                                "name": "search_flights",
+                                "description": "Search for available flights.",
+                            }
+                        ]
+                    }
+                ],
+            )
+        },
+        user_scenario_generation_config=types.evals.UserScenarioGenerationConfig(
+            user_scenario_count=2,
+            simulation_instruction=(
+                "Generate scenarios where the user tries to book a flight but"
+                " changes their mind about the destination."
+            ),
+            environment_data="Today is Monday. Flights to Paris are available.",
+            model_name="gemini-2.5-flash",
+        ),
+        root_agent_id="booking-agent",
+    )
+    assert isinstance(eval_dataset, types.EvaluationDataset)
+    assert len(eval_dataset.eval_cases) == 2
+    assert (
+        eval_dataset.eval_cases[1].user_scenario.starting_prompt
+        == "Find me a flight from Boston to Rome for next month."
+    )
+    assert (
+        eval_dataset.eval_cases[1].user_scenario.conversation_plan
+        == "Wait, change of plans. I need to go to Milan instead, and it needs to be a round trip, returning two weeks after departure."
+    )
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method="evals.generate_user_scenarios",
+)

tests/unit/vertexai/genai/test_evals.py

@@ -20,6 +20,7 @@
 import re
 import statistics
 import sys
+import unittest
 from unittest import mock
 
 import google.auth.credentials
@@ -5500,3 +5501,70 @@ def read_file_contents_side_effect(src: str) -> str:
                 }
             ),
         )
+
+
+class TestEvalsGenerateUserScenarios(unittest.TestCase):
+    """Unit tests for the Evals generate_user_scenarios method."""
+
+    def setUp(self):
+        self.addCleanup(mock.patch.stopall)
+        self.mock_client = mock.MagicMock(spec=client.Client)
+        self.mock_client.vertexai = True
+        self.mock_api_client = mock.MagicMock()
+        self.mock_client._api_client = self.mock_api_client
+
+        self.mock_response = mock.MagicMock()
+        self.mock_response.body = json.dumps(
+            {
+                "userScenarios": [
+                    {"startingPrompt": "Prompt 1", "conversationPlan": "Plan 1"},
+                    {"startingPrompt": "Prompt 2", "conversationPlan": "Plan 2"},
+                ]
+            }
+        )
+        self.mock_api_client.request.return_value = self.mock_response
+
+    def test_generate_user_scenarios(self):
+        """Tests that generate_user_scenarios correctly calls the API and parses the response."""
+        evals_module = evals.Evals(api_client_=self.mock_api_client)
+
+        eval_dataset = evals_module.generate_user_scenarios(
+            agents={"agent_1": {}},
+            user_scenario_generation_config={"user_scenario_count": 2},
+            root_agent_id="agent_1",
+        )
+        assert isinstance(eval_dataset, vertexai_genai_types.EvaluationDataset)
+        assert len(eval_dataset.eval_cases) == 2
+        assert eval_dataset.eval_cases[0].user_scenario.starting_prompt == "Prompt 1"
+        assert eval_dataset.eval_cases[0].user_scenario.conversation_plan == "Plan 1"
+        assert eval_dataset.eval_cases[1].user_scenario.starting_prompt == "Prompt 2"
+        assert eval_dataset.eval_cases[1].user_scenario.conversation_plan == "Plan 2"
+
+        assert eval_dataset.eval_dataset_df is not None
+        assert len(eval_dataset.eval_dataset_df) == 2
+        assert eval_dataset.eval_dataset_df.iloc[0]["starting_prompt"] == "Prompt 1"
+
+        self.mock_api_client.request.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_async_generate_user_scenarios(self):
+        """Tests that async generate_user_scenarios correctly calls the API and parses the response."""
+
+        self.mock_api_client.async_request = mock.AsyncMock(
+            return_value=self.mock_response
+        )
+        async_evals_module = evals.AsyncEvals(api_client_=self.mock_api_client)
+
+        eval_dataset = await async_evals_module.generate_user_scenarios(
+            agents={"agent_1": {}},
+            user_scenario_generation_config={"user_scenario_count": 2},
+            root_agent_id="agent_1",
+        )
+        assert isinstance(eval_dataset, vertexai_genai_types.EvaluationDataset)
+        assert len(eval_dataset.eval_cases) == 2
+        assert eval_dataset.eval_cases[0].user_scenario.starting_prompt == "Prompt 1"
+
+        assert eval_dataset.eval_dataset_df is not None
+        assert len(eval_dataset.eval_dataset_df) == 2
+
+        self.mock_api_client.async_request.assert_called_once()