feat: GenAI Client(evals) - exponential backoff for Gemini agent interaction polling

PiperOrigin-RevId: 947191099

Author: jsondai

agentplatform/_genai/_evals_common.py

@@ -988,28 +988,65 @@ def _agent_data_response_text(agent_data: types.evals.AgentData) -> Optional[str
     return "".join(text_parts) or None
 
 
+def _agent_resource_to_agent_info(
+    agent: str, api_client: BaseApiClient
+) -> "types.evals.AgentInfo":
+    """Builds an `AgentInfo` from a Gemini Agents API agent resource name.
+
+    Fetches the agent through the SDK's `api_client` (so replay recording is
+    preserved) via `_fetch_agent_config_dict` and derives a single-agent
+    `AgentInfo`: the agent's short name is the agents-map key and
+    `root_agent_id`.
+
+    Args:
+        agent: The Gemini Agents API agent resource name
+          (`projects/{p}/locations/{l}/agents/{name}`).
+        api_client: The API client used to fetch the agent.
+
+    Returns:
+        An `AgentInfo` describing the fetched agent.
+    """
+    agent_config = _fetch_agent_config_dict(api_client, agent)
+    short_name = agent_config.agent_id
+    return types.evals.AgentInfo(  # pytype: disable=missing-parameter
+        name=short_name,
+        agents={short_name: agent_config},
+        root_agent_id=short_name,
+    )
+
+
 _INTERACTION_TERMINAL_STATES = frozenset(
     ["completed", "failed", "cancelled", "incomplete", "budget_exceeded"]
 )
 
+_INITIAL_POLL_INTERVAL_SECONDS = 2.0
+_MAX_POLL_INTERVAL_SECONDS = 30.0
+_POLL_BACKOFF_MULTIPLIER = 2.0
+
 
 def _await_interaction(
     interactions_client: "_InteractionsRestClient",
     interaction: dict[str, Any],
-    poll_interval_seconds: float = 2.0,
+    initial_poll_interval_seconds: float = _INITIAL_POLL_INTERVAL_SECONDS,
+    max_poll_interval_seconds: float = _MAX_POLL_INTERVAL_SECONDS,
+    poll_backoff_multiplier: float = _POLL_BACKOFF_MULTIPLIER,
     timeout_seconds: float = 600.0,
 ) -> dict[str, Any]:
     """Polls a background interaction until it reaches a terminal state.
 
     Gemini agent interactions must run in the background (`background=True`), so
     `create` returns before the model output is ready. This polls
     `interactions.get` until the interaction reaches a terminal state and then
-    returns the resolved interaction.
+    returns the resolved interaction. The delay between polls grows
+    exponentially (capped at `max_poll_interval_seconds`) to avoid hitting rate
+    limits when evaluating large datasets.
 
     Args:
         interactions_client: The interactions client used to poll.
         interaction: The interaction returned by `create`.
-        poll_interval_seconds: Delay between poll attempts.
+        initial_poll_interval_seconds: Delay before the first poll.
+        max_poll_interval_seconds: Upper bound for the poll interval.
+        poll_backoff_multiplier: Factor the interval grows by after each poll.
         timeout_seconds: Maximum time to wait before raising.
 
     Returns:
@@ -1022,11 +1059,15 @@ def _await_interaction(
         return interaction
     interaction_id = interaction.get("id")
     deadline = time.monotonic() + timeout_seconds
+    poll_interval = initial_poll_interval_seconds
     while time.monotonic() < deadline:
-        time.sleep(poll_interval_seconds)
+        time.sleep(poll_interval)
         interaction = interactions_client.get(interaction_id)
         if interaction.get("status") in _INTERACTION_TERMINAL_STATES:
             return interaction
+        poll_interval = min(
+            poll_interval * poll_backoff_multiplier, max_poll_interval_seconds
+        )
     raise TimeoutError(
         f"Interaction {interaction_id} did not complete within"
         f" {timeout_seconds} seconds."

agentplatform/_genai/evals.py

@@ -2956,28 +2956,54 @@ def create_evaluation_set(
     def generate_conversation_scenarios(
         self,
         *,
-        agent_info: evals_types.AgentInfoOrDict,
+        agent_info: Optional[evals_types.AgentInfoOrDict] = None,
+        agent: Optional[str] = None,
         config: evals_types.UserScenarioGenerationConfigOrDict,
         allow_cross_region_model: Optional[bool] = None,
     ) -> types.EvaluationDataset:
         """Generates an evaluation dataset with user scenarios,
            which helps to generate conversations between a simulated user
            and the agent under test.
 
+        Exactly one of `agent_info` or `agent` must be provided. When `agent` is
+        a Gemini Agents API agent resource name, the agent is fetched and an
+        `AgentInfo` is derived from it.
+
         Args:
-            agent_info: The agent info to generate user scenarios for.
+            agent_info: The agent info to generate user scenarios for. Mutually
+                exclusive with `agent`.
+            agent: A Gemini Agents API agent resource name
+                (`projects/{p}/locations/{l}/agents/{name}`). When provided, the
+                agent is fetched and its configuration is used to build the agent
+                info. Mutually exclusive with `agent_info`.
             config: Configuration for generating user scenarios.
             allow_cross_region_model: Opt-in flag to authorize cross-region
                 routing for model inference.
 
         Returns:
             An EvaluationDataset containing the generated user scenarios.
         """
-        parsed_agent_info = (
-            evals_types.AgentInfo.model_validate(agent_info)
-            if isinstance(agent_info, dict)
-            else agent_info
-        )
+        if agent is not None and agent_info is not None:
+            raise ValueError(
+                "Only one of `agent` or `agent_info` may be provided, not both."
+            )
+        if agent is None and agent_info is None:
+            raise ValueError("One of `agent` or `agent_info` must be provided.")
+        if agent is not None:
+            if not _evals_common._is_gemini_agent_resource(agent):
+                raise ValueError(
+                    "`agent` must be a Gemini Agents API agent resource name of the"
+                    " form projects/{project}/locations/{location}/agents/{agent}."
+                )
+            parsed_agent_info = _evals_common._agent_resource_to_agent_info(
+                agent, self._api_client
+            )
+        else:
+            parsed_agent_info = (
+                evals_types.AgentInfo.model_validate(agent_info)
+                if isinstance(agent_info, dict)
+                else agent_info
+            )
         response = self._generate_user_scenarios(
             agents=parsed_agent_info.agents,
             root_agent_id=parsed_agent_info.root_agent_id,
@@ -4769,28 +4795,54 @@ async def create_evaluation_set(
     async def generate_conversation_scenarios(
         self,
         *,
-        agent_info: evals_types.AgentInfoOrDict,
+        agent_info: Optional[evals_types.AgentInfoOrDict] = None,
+        agent: Optional[str] = None,
         config: evals_types.UserScenarioGenerationConfigOrDict,
         allow_cross_region_model: Optional[bool] = None,
     ) -> types.EvaluationDataset:
         """Generates an evaluation dataset with user scenarios,
            which helps to generate conversations between a simulated user
            and the agent under test.
 
+        Exactly one of `agent_info` or `agent` must be provided. When `agent` is
+        a Gemini Agents API agent resource name, the agent is fetched and an
+        `AgentInfo` is derived from it.
+
         Args:
-            agent_info: The agent info to generate user scenarios for.
+            agent_info: The agent info to generate user scenarios for. Mutually
+                exclusive with `agent`.
+            agent: A Gemini Agents API agent resource name
+                (`projects/{p}/locations/{l}/agents/{name}`). When provided, the
+                agent is fetched and its configuration is used to build the agent
+                info. Mutually exclusive with `agent_info`.
             config: Configuration for generating user scenarios.
             allow_cross_region_model: Opt-in flag to authorize cross-region
                 routing for model inference.
 
         Returns:
             An EvaluationDataset containing the generated user scenarios.
         """
-        parsed_agent_info = (
-            evals_types.AgentInfo.model_validate(agent_info)
-            if isinstance(agent_info, dict)
-            else agent_info
-        )
+        if agent is not None and agent_info is not None:
+            raise ValueError(
+                "Only one of `agent` or `agent_info` may be provided, not both."
+            )
+        if agent is None and agent_info is None:
+            raise ValueError("One of `agent` or `agent_info` must be provided.")
+        if agent is not None:
+            if not _evals_common._is_gemini_agent_resource(agent):
+                raise ValueError(
+                    "`agent` must be a Gemini Agents API agent resource name of the"
+                    " form projects/{project}/locations/{location}/agents/{agent}."
+                )
+            parsed_agent_info = _evals_common._agent_resource_to_agent_info(
+                agent, self._api_client
+            )
+        else:
+            parsed_agent_info = (
+                evals_types.AgentInfo.model_validate(agent_info)
+                if isinstance(agent_info, dict)
+                else agent_info
+            )
         response = await self._generate_user_scenarios(
             agents=parsed_agent_info.agents,
             root_agent_id=parsed_agent_info.root_agent_id,

tests/unit/agentplatform/genai/replays/test_generate_conversation_scenarios.py

@@ -103,6 +103,30 @@ async def test_gen_conversation_scenarios_async(client):
     assert eval_dataset.eval_cases[1].user_scenario.conversation_plan
 
 
+def test_scenarios_from_gemini_agent(client):
+    """Tests generate_conversation_scenarios() against a Gemini Agents API agent.
+
+    Fetches the agent via google.genai agents.get, derives an AgentInfo from it,
+    and generates user scenarios from the derived config.
+    """
+    eval_dataset = client.evals.generate_conversation_scenarios(
+        agent=("projects/model-evaluation-dev/locations/global/agents/test-agent-eval"),
+        config=types.evals.UserScenarioGenerationConfig(
+            count=2,
+            generation_instruction=(
+                "Generate scenarios where the user tries to book a flight but"
+                " changes their mind about the destination."
+            ),
+            environment_context="Today is Monday. Flights to Paris are available.",
+            model_name="gemini-2.5-flash",
+        ),
+    )
+    assert isinstance(eval_dataset, types.EvaluationDataset)
+    assert len(eval_dataset.eval_cases) == 2
+    assert eval_dataset.eval_cases[0].user_scenario.starting_prompt
+    assert eval_dataset.eval_cases[0].user_scenario.conversation_plan
+
+
 pytestmark = pytest_helper.setup(
     file=__file__,
     globals_for_file=globals(),

tests/unit/agentplatform/genai/test_evals.py

@@ -4292,6 +4292,65 @@ def test_run_inference_non_gemini_agent_routes_to_agent_engine(
         mock_get_interactions_client.assert_not_called()
 
 
+@pytest.mark.usefixtures("google_auth_mock")
+class TestAwaitInteraction:
+    """Unit tests for _await_interaction polling behavior."""
+
+    def test_await_interaction_returns_immediately_when_terminal(self):
+        client = mock.Mock()
+        interaction = {"id": "i1", "status": "completed"}
+        result = _evals_common._await_interaction(client, interaction)
+        assert result is interaction
+        client.get.assert_not_called()
+
+    def test_await_interaction_polls_with_exponential_backoff(self):
+        client = mock.Mock()
+        client.get.side_effect = [
+            {"id": "i1", "status": "in_progress"},
+            {"id": "i1", "status": "in_progress"},
+            {"id": "i1", "status": "completed"},
+        ]
+        sleeps = []
+        with mock.patch.object(_evals_common.time, "sleep", sleeps.append):
+            result = _evals_common._await_interaction(
+                client,
+                {"id": "i1", "status": "in_progress"},
+                initial_poll_interval_seconds=2.0,
+                max_poll_interval_seconds=30.0,
+                poll_backoff_multiplier=2.0,
+            )
+        assert result["status"] == "completed"
+        assert sleeps == [2.0, 4.0, 8.0]
+
+    def test_await_interaction_caps_poll_interval(self):
+        client = mock.Mock()
+        client.get.side_effect = [
+            {"id": "i1", "status": "in_progress"},
+            {"id": "i1", "status": "completed"},
+        ]
+        sleeps = []
+        with mock.patch.object(_evals_common.time, "sleep", sleeps.append):
+            _evals_common._await_interaction(
+                client,
+                {"id": "i1", "status": "in_progress"},
+                initial_poll_interval_seconds=25.0,
+                max_poll_interval_seconds=30.0,
+                poll_backoff_multiplier=2.0,
+            )
+        assert sleeps == [25.0, 30.0]
+
+    def test_await_interaction_raises_on_timeout(self):
+        client = mock.Mock()
+        client.get.return_value = {"id": "i1", "status": "in_progress"}
+        with mock.patch.object(_evals_common.time, "sleep"):
+            with pytest.raises(TimeoutError, match="did not complete"):
+                _evals_common._await_interaction(
+                    client,
+                    {"id": "i1", "status": "in_progress"},
+                    timeout_seconds=-1.0,
+                )
+
+
 @pytest.mark.usefixtures("google_auth_mock")
 class TestEvalsMetricHandlers:
     """Unit tests for utility functions in _evals_metric_handlers."""
@@ -9133,6 +9192,68 @@ async def test_async_generate_conversation_scenarios(self):
         request_body = call_args[0][2]  # Third positional arg is the request dict
         assert request_body.get("allowCrossRegionModel") is True
 
+    @mock.patch.object(_evals_common, "_fetch_agent_config_dict")
+    def test_generate_conversation_scenarios_from_gemini_agent(
+        self, mock_fetch_agent_config
+    ):
+        mock_fetch_agent_config.return_value = (
+            agentplatform_genai_types.evals.AgentConfig(
+                agent_id="test-agent",
+                instruction="You are a helpful travel assistant.",
+                description="An agent that books flights.",
+                tools=[genai_types.Tool(google_search=genai_types.GoogleSearch())],
+            )
+        )
+
+        evals_module = evals.Evals(api_client_=self.mock_api_client)
+
+        with mock.patch.object(
+            evals_module, "_generate_user_scenarios"
+        ) as mock_generate_user_scenarios:
+            mock_generate_user_scenarios.return_value = self.mock_response
+            evals_module.generate_conversation_scenarios(
+                agent=_TEST_GEMINI_AGENT,
+                config={"count": 2},
+            )
+
+        mock_fetch_agent_config.assert_called_once_with(
+            self.mock_api_client, _TEST_GEMINI_AGENT
+        )
+        call_kwargs = mock_generate_user_scenarios.call_args.kwargs
+        assert call_kwargs["root_agent_id"] == "test-agent"
+        agents = call_kwargs["agents"]
+        assert "test-agent" in agents
+        derived_config = agents["test-agent"]
+        assert derived_config.instruction == "You are a helpful travel assistant."
+        assert derived_config.description == "An agent that books flights."
+        assert derived_config.tools is not None
+        assert derived_config.tools[0].google_search is not None
+
+    def test_generate_conversation_scenarios_agent_and_agent_info_raises(self):
+        evals_module = evals.Evals(api_client_=self.mock_api_client)
+        with pytest.raises(ValueError, match="not both"):
+            evals_module.generate_conversation_scenarios(
+                agent=_TEST_GEMINI_AGENT,
+                agent_info=agentplatform_genai_types.evals.AgentInfo(
+                    agents={"agent_1": {}},
+                    root_agent_id="agent_1",
+                ),
+                config={"count": 2},
+            )
+
+    def test_generate_conversation_scenarios_no_agent_raises(self):
+        evals_module = evals.Evals(api_client_=self.mock_api_client)
+        with pytest.raises(ValueError, match="must be provided"):
+            evals_module.generate_conversation_scenarios(config={"count": 2})
+
+    def test_generate_conversation_scenarios_non_gemini_agent_raises(self):
+        evals_module = evals.Evals(api_client_=self.mock_api_client)
+        with pytest.raises(ValueError, match="Gemini Agents API"):
+            evals_module.generate_conversation_scenarios(
+                agent=_TEST_AGENT_ENGINE,
+                config={"count": 2},
+            )
+
 
 class TestTransformDataframe:
     """Unit tests for the _transform_dataframe function."""