feat: GenAI Client(evals): Add user 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,118 @@
+# 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()

vertexai/_genai/_evals_utils.py

@@ -19,6 +19,7 @@
 import os
 from typing import Any, Optional, Union
 
+from google.genai import types as genai_types
 from google.genai._api_client import BaseApiClient
 from google.genai._common import get_value_by_path as getv
 from google.genai._common import set_value_by_path as setv
@@ -335,3 +336,37 @@ class EvalDataConverter(abc.ABC):
     def convert(self, raw_data: Any) -> "types.EvaluationDataset":
         """Converts a loaded raw dataset into an EvaluationDataset."""
         raise NotImplementedError()
+
+
+def _postprocess_user_scenarios_response(
+    response: types.GenerateUserScenariosResponse,
+) -> types.EvaluationDataset:
+    """Postprocesses the response from generating user scenarios."""
+    eval_cases = []
+    data_for_df = []
+    if hasattr(response, "user_scenarios") and response.user_scenarios:
+        for scenario in response.user_scenarios:
+            prompt_content = None
+            if scenario.starting_prompt:
+                prompt_content = genai_types.Content(
+                    parts=[genai_types.Part(text=scenario.starting_prompt)]
+                )
+            eval_case = types.EvalCase(
+                prompt=prompt_content,
+                user_scenario=scenario,
+            )
+            eval_cases.append(eval_case)
+            data_for_df.append(
+                {
+                    "starting_prompt": scenario.starting_prompt,
+                    "conversation_plan": scenario.conversation_plan,
+                }
+            )
+    eval_dataset_df = None
+    if pd is not None:
+        eval_dataset_df = pd.DataFrame(data_for_df)
+    else:
+        logger.warning("Pandas is not installed. eval_dataset_df will be None.")
+    return types.EvaluationDataset(
+        eval_cases=eval_cases, eval_dataset_df=eval_dataset_df
+    )