diff --git a/tests/unit/vertexai/genai/replays/test_generate_user_scenarios.py b/tests/unit/vertexai/genai/replays/test_generate_user_scenarios.py new file mode 100644 index 0000000000..afa54271f8 --- /dev/null +++ b/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", +) diff --git a/tests/unit/vertexai/genai/test_evals.py b/tests/unit/vertexai/genai/test_evals.py index 7c4ed635f5..9f880c4103 100644 --- a/tests/unit/vertexai/genai/test_evals.py +++ b/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() diff --git a/vertexai/_genai/_evals_utils.py b/vertexai/_genai/_evals_utils.py index 36c94f02a4..4a338d34b0 100644 --- a/vertexai/_genai/_evals_utils.py +++ b/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 + ) diff --git a/vertexai/_genai/evals.py b/vertexai/_genai/evals.py index dba63496fb..bbf5a56c56 100644 --- a/vertexai/_genai/evals.py +++ b/vertexai/_genai/evals.py @@ -32,6 +32,7 @@ from . import _evals_utils from . import _transformers as t from . import types +from .types import evals as evals_types try: from google.adk.agents import LlmAgent @@ -407,6 +408,33 @@ def _GenerateInstanceRubricsRequest_to_vertex( return to_object +def _GenerateUserScenariosParameters_to_vertex( + from_object: Union[dict[str, Any], object], + parent_object: Optional[dict[str, Any]] = None, +) -> dict[str, Any]: + to_object: dict[str, Any] = {} + if getv(from_object, ["location"]) is not None: + setv(to_object, ["location"], getv(from_object, ["location"])) + + if getv(from_object, ["agents"]) is not None: + setv(to_object, ["agents"], getv(from_object, ["agents"])) + + if getv(from_object, ["root_agent_id"]) is not None: + setv(to_object, ["rootAgentId"], getv(from_object, ["root_agent_id"])) + + if getv(from_object, ["user_scenario_generation_config"]) is not None: + setv( + to_object, + ["userScenarioGenerationConfig"], + getv(from_object, ["user_scenario_generation_config"]), + ) + + if getv(from_object, ["config"]) is not None: + setv(to_object, ["config"], getv(from_object, ["config"])) + + return to_object + + def _GetEvaluationItemParameters_to_vertex( from_object: Union[dict[str, Any], object], parent_object: Optional[dict[str, Any]] = None, @@ -962,6 +990,67 @@ def _generate_rubrics( self._api_client._verify_response(return_value) return return_value + def _generate_user_scenarios( + self, + *, + location: Optional[str] = None, + agents: Optional[dict[str, evals_types.AgentConfigOrDict]] = None, + root_agent_id: Optional[str] = None, + user_scenario_generation_config: Optional[ + evals_types.UserScenarioGenerationConfigOrDict + ] = None, + config: Optional[types.GenerateUserScenariosConfigOrDict] = None, + ) -> types.GenerateUserScenariosResponse: + """ + Generates user scenarios for agent evaluation. + """ + + parameter_model = types._GenerateUserScenariosParameters( + location=location, + agents=agents, + root_agent_id=root_agent_id, + user_scenario_generation_config=user_scenario_generation_config, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError("This method is only supported in the Vertex AI client.") + else: + request_dict = _GenerateUserScenariosParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = ":generateUserScenarios".format_map(request_url_dict) + else: + path = ":generateUserScenarios" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = self._api_client.request("post", path, request_dict, http_options) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.GenerateUserScenariosResponse._from_response( + response=response_dict, kwargs=parameter_model.model_dump() + ) + + self._api_client._verify_response(return_value) + return return_value + def _get_evaluation_run( self, *, name: str, config: Optional[types.GetEvaluationRunConfigOrDict] = None ) -> types.EvaluationRun: @@ -1594,7 +1683,7 @@ def create_evaluation_run( metrics: list[types.EvaluationRunMetricOrDict], name: Optional[str] = None, display_name: Optional[str] = None, - agent_info: Optional[types.evals.AgentInfoOrDict] = None, + agent_info: Optional[evals_types.AgentInfoOrDict] = None, inference_configs: Optional[ dict[str, types.EvaluationRunInferenceConfigOrDict] ] = None, @@ -1626,9 +1715,9 @@ def create_evaluation_run( "At most one of agent_info or inference_configs can be provided." ) agent_info_pydantic = ( - types.evals.AgentInfo.model_validate(agent_info) + evals_types.AgentInfo.model_validate(agent_info) if isinstance(agent_info, dict) - else (agent_info or types.evals.AgentInfo()) + else (agent_info or evals_types.AgentInfo()) ) resolved_dataset = _evals_common._resolve_dataset( self._api_client, dataset, dest, agent_info_pydantic @@ -1793,6 +1882,36 @@ def create_evaluation_set( config=config, ) + @_common.experimental_warning( + "The Vertex SDK GenAI evals.generate_user_scenarios module is experimental, " + "and may change in future versions." + ) + def generate_user_scenarios( + self, + *, + agents: dict[str, evals_types.AgentConfigOrDict], + user_scenario_generation_config: evals_types.UserScenarioGenerationConfigOrDict, + root_agent_id: str, + ) -> types.EvaluationDataset: + """Generates an evaluation dataset with user scenarios, + which helps to generate conversations between a simulated user + and the agent under test. + + Args: + agents: A map of agent ID to AgentConfig. + user_scenario_generation_config: Configuration for generating user scenarios. + root_agent_id: The ID of the root agent. + + Returns: + An EvaluationDataset containing the generated user scenarios. + """ + response = self._generate_user_scenarios( + agents=agents, + user_scenario_generation_config=user_scenario_generation_config, + root_agent_id=root_agent_id, + ) + return _evals_utils._postprocess_user_scenarios_response(response) + class AsyncEvals(_api_module.BaseModule): @@ -2128,6 +2247,69 @@ async def _generate_rubrics( self._api_client._verify_response(return_value) return return_value + async def _generate_user_scenarios( + self, + *, + location: Optional[str] = None, + agents: Optional[dict[str, evals_types.AgentConfigOrDict]] = None, + root_agent_id: Optional[str] = None, + user_scenario_generation_config: Optional[ + evals_types.UserScenarioGenerationConfigOrDict + ] = None, + config: Optional[types.GenerateUserScenariosConfigOrDict] = None, + ) -> types.GenerateUserScenariosResponse: + """ + Generates user scenarios for agent evaluation. + """ + + parameter_model = types._GenerateUserScenariosParameters( + location=location, + agents=agents, + root_agent_id=root_agent_id, + user_scenario_generation_config=user_scenario_generation_config, + config=config, + ) + + request_url_dict: Optional[dict[str, str]] + if not self._api_client.vertexai: + raise ValueError("This method is only supported in the Vertex AI client.") + else: + request_dict = _GenerateUserScenariosParameters_to_vertex(parameter_model) + request_url_dict = request_dict.get("_url") + if request_url_dict: + path = ":generateUserScenarios".format_map(request_url_dict) + else: + path = ":generateUserScenarios" + + query_params = request_dict.get("_query") + if query_params: + path = f"{path}?{urlencode(query_params)}" + # TODO: remove the hack that pops config. + request_dict.pop("config", None) + + http_options: Optional[types.HttpOptions] = None + if ( + parameter_model.config is not None + and parameter_model.config.http_options is not None + ): + http_options = parameter_model.config.http_options + + request_dict = _common.convert_to_dict(request_dict) + request_dict = _common.encode_unserializable_types(request_dict) + + response = await self._api_client.async_request( + "post", path, request_dict, http_options + ) + + response_dict = {} if not response.body else json.loads(response.body) + + return_value = types.GenerateUserScenariosResponse._from_response( + response=response_dict, kwargs=parameter_model.model_dump() + ) + + self._api_client._verify_response(return_value) + return return_value + async def _get_evaluation_run( self, *, name: str, config: Optional[types.GetEvaluationRunConfigOrDict] = None ) -> types.EvaluationRun: @@ -2429,7 +2611,7 @@ async def create_evaluation_run( metrics: list[types.EvaluationRunMetricOrDict], name: Optional[str] = None, display_name: Optional[str] = None, - agent_info: Optional[types.evals.AgentInfo] = None, + agent_info: Optional[evals_types.AgentInfo] = None, inference_configs: Optional[ dict[str, types.EvaluationRunInferenceConfigOrDict] ] = None, @@ -2461,9 +2643,9 @@ async def create_evaluation_run( "At most one of agent_info or inference_configs can be provided." ) agent_info_pydantic = ( - types.evals.AgentInfo.model_validate(agent_info) + evals_types.AgentInfo.model_validate(agent_info) if isinstance(agent_info, dict) - else (agent_info or types.evals.AgentInfo()) + else (agent_info or evals_types.AgentInfo()) ) resolved_dataset = _evals_common._resolve_dataset( self._api_client, dataset, dest, agent_info_pydantic @@ -2634,3 +2816,33 @@ async def create_evaluation_set( config=config, ) return result + + @_common.experimental_warning( + "The Vertex SDK GenAI evals.generate_user_scenarios module is experimental, " + "and may change in future versions." + ) + async def generate_user_scenarios( + self, + *, + agents: dict[str, evals_types.AgentConfigOrDict], + user_scenario_generation_config: evals_types.UserScenarioGenerationConfigOrDict, + root_agent_id: str, + ) -> types.EvaluationDataset: + """Generates an evaluation dataset with user scenarios, + which helps to generate conversations between a simulated user + and the agent under test. + + Args: + agents: A map of agent ID to AgentConfig. + user_scenario_generation_config: Configuration for generating user scenarios. + root_agent_id: The ID of the root agent. + + Returns: + An EvaluationDataset containing the generated user scenarios. + """ + response = await self._generate_user_scenarios( + agents=agents, + user_scenario_generation_config=user_scenario_generation_config, + root_agent_id=root_agent_id, + ) + return _evals_utils._postprocess_user_scenarios_response(response) diff --git a/vertexai/_genai/types/__init__.py b/vertexai/_genai/types/__init__.py index 5eff899c3a..d3a933f84d 100644 --- a/vertexai/_genai/types/__init__.py +++ b/vertexai/_genai/types/__init__.py @@ -50,6 +50,7 @@ from .common import _ExecuteCodeAgentEngineSandboxRequestParameters from .common import _GenerateAgentEngineMemoriesRequestParameters from .common import _GenerateInstanceRubricsRequest +from .common import _GenerateUserScenariosParameters from .common import _GetAgentEngineGenerateMemoriesOperationParameters from .common import _GetAgentEngineMemoryOperationParameters from .common import _GetAgentEngineMemoryRequestParameters @@ -474,6 +475,12 @@ from .common import GenerateMemoriesResponseGeneratedMemoryDict from .common import GenerateMemoriesResponseGeneratedMemoryOrDict from .common import GenerateMemoriesResponseOrDict +from .common import GenerateUserScenariosConfig +from .common import GenerateUserScenariosConfigDict +from .common import GenerateUserScenariosConfigOrDict +from .common import GenerateUserScenariosResponse +from .common import GenerateUserScenariosResponseDict +from .common import GenerateUserScenariosResponseOrDict from .common import GetAgentEngineConfig from .common import GetAgentEngineConfigDict from .common import GetAgentEngineConfigOrDict @@ -1525,6 +1532,12 @@ "GenerateInstanceRubricsResponse", "GenerateInstanceRubricsResponseDict", "GenerateInstanceRubricsResponseOrDict", + "GenerateUserScenariosConfig", + "GenerateUserScenariosConfigDict", + "GenerateUserScenariosConfigOrDict", + "GenerateUserScenariosResponse", + "GenerateUserScenariosResponseDict", + "GenerateUserScenariosResponseOrDict", "GetEvaluationRunConfig", "GetEvaluationRunConfigDict", "GetEvaluationRunConfigOrDict", @@ -2198,6 +2211,7 @@ "_CreateEvaluationSetParameters", "_EvaluateInstancesRequestParameters", "_GenerateInstanceRubricsRequest", + "_GenerateUserScenariosParameters", "_GetEvaluationRunParameters", "_GetEvaluationSetParameters", "_GetEvaluationItemParameters", diff --git a/vertexai/_genai/types/common.py b/vertexai/_genai/types/common.py index d9e37aff44..a4ab38e698 100644 --- a/vertexai/_genai/types/common.py +++ b/vertexai/_genai/types/common.py @@ -2685,6 +2685,10 @@ class EvalCase(_common.BaseModel): default=None, description="""This field is experimental and may change in future versions. The agent data of the agent under evaluation.""", ) + user_scenario: Optional[evals_types.UserScenario] = Field( + default=None, + description="""This field is experimental and may change in future versions. The user scenario for the evaluation case.""", + ) # Allow extra fields to support custom metric prompts and stay backward compatible. model_config = ConfigDict(frozen=True, extra="allow") @@ -2722,6 +2726,9 @@ class EvalCaseDict(TypedDict, total=False): agent_data: Optional[evals_types.AgentData] """This field is experimental and may change in future versions. The agent data of the agent under evaluation.""" + user_scenario: Optional[evals_types.UserScenario] + """This field is experimental and may change in future versions. The user scenario for the evaluation case.""" + EvalCaseOrDict = Union[EvalCase, EvalCaseDict] @@ -5094,6 +5101,84 @@ class GenerateInstanceRubricsResponseDict(TypedDict, total=False): ] +class GenerateUserScenariosConfig(_common.BaseModel): + + http_options: Optional[genai_types.HttpOptions] = Field( + default=None, description="""Used to override HTTP request options.""" + ) + + +class GenerateUserScenariosConfigDict(TypedDict, total=False): + + http_options: Optional[genai_types.HttpOptionsDict] + """Used to override HTTP request options.""" + + +GenerateUserScenariosConfigOrDict = Union[ + GenerateUserScenariosConfig, GenerateUserScenariosConfigDict +] + + +class _GenerateUserScenariosParameters(_common.BaseModel): + """Parameters for GenerateUserScenarios.""" + + location: Optional[str] = Field(default=None, description="""""") + agents: Optional[dict[str, evals_types.AgentConfig]] = Field( + default=None, description="""""" + ) + root_agent_id: Optional[str] = Field(default=None, description="""""") + user_scenario_generation_config: Optional[ + evals_types.UserScenarioGenerationConfig + ] = Field(default=None, description="""""") + config: Optional[GenerateUserScenariosConfig] = Field( + default=None, description="""""" + ) + + +class _GenerateUserScenariosParametersDict(TypedDict, total=False): + """Parameters for GenerateUserScenarios.""" + + location: Optional[str] + """""" + + agents: Optional[dict[str, evals_types.AgentConfig]] + """""" + + root_agent_id: Optional[str] + """""" + + user_scenario_generation_config: Optional[evals_types.UserScenarioGenerationConfig] + """""" + + config: Optional[GenerateUserScenariosConfigDict] + """""" + + +_GenerateUserScenariosParametersOrDict = Union[ + _GenerateUserScenariosParameters, _GenerateUserScenariosParametersDict +] + + +class GenerateUserScenariosResponse(_common.BaseModel): + """Response message for DataFoundryService.GenerateUserScenarios.""" + + user_scenarios: Optional[list[evals_types.UserScenario]] = Field( + default=None, description="""""" + ) + + +class GenerateUserScenariosResponseDict(TypedDict, total=False): + """Response message for DataFoundryService.GenerateUserScenarios.""" + + user_scenarios: Optional[list[evals_types.UserScenario]] + """""" + + +GenerateUserScenariosResponseOrDict = Union[ + GenerateUserScenariosResponse, GenerateUserScenariosResponseDict +] + + class GetEvaluationRunConfig(_common.BaseModel): """Config for get evaluation run.""" diff --git a/vertexai/_genai/types/evals.py b/vertexai/_genai/types/evals.py index 970286cd1d..9b30fa13ef 100644 --- a/vertexai/_genai/types/evals.py +++ b/vertexai/_genai/types/evals.py @@ -796,6 +796,48 @@ class UserScenarioDict(TypedDict, total=False): UserScenarioOrDict = Union[UserScenario, UserScenarioDict] +class UserScenarioGenerationConfig(_common.BaseModel): + """User scenario generation configuration.""" + + user_scenario_count: Optional[int] = Field( + default=None, + description="""The number of user scenarios to generate. The maximum number of scenarios that can be generated is 100.""", + ) + simulation_instruction: Optional[str] = Field( + default=None, + description="""Simulation instruction to guide the user scenario generation.""", + ) + environment_data: Optional[str] = Field( + default=None, + description="""Environment data to drive simulation. For example, for a QA agent, this could be the docs queried by the tools.""", + ) + model_name: Optional[str] = Field( + default=None, + description="""The model name to use for user scenario generation.""", + ) + + +class UserScenarioGenerationConfigDict(TypedDict, total=False): + """User scenario generation configuration.""" + + user_scenario_count: Optional[int] + """The number of user scenarios to generate. The maximum number of scenarios that can be generated is 100.""" + + simulation_instruction: Optional[str] + """Simulation instruction to guide the user scenario generation.""" + + environment_data: Optional[str] + """Environment data to drive simulation. For example, for a QA agent, this could be the docs queried by the tools.""" + + model_name: Optional[str] + """The model name to use for user scenario generation.""" + + +UserScenarioGenerationConfigOrDict = Union[ + UserScenarioGenerationConfig, UserScenarioGenerationConfigDict +] + + class UserSimulatorConfig(_common.BaseModel): """Configuration for a user simulator that uses an LLM to generate messages."""