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

PiperOrigin-RevId: 947191099

Author: jsondai

agentplatform/_genai/_evals_common.py

@@ -928,6 +928,217 @@ def _get_resolved_location(api_client: Any) -> Optional[str]:
     return getattr(api_client, "location", None)
 
 
+class _InteractionsRestClient:
+    """Minimal Interactions API client issued through the SDK api_client.
+
+    Calls go through `api_client.request()` (rather than the google.genai
+    `_gaos` client) so that the `ReplayApiClient` records and replays them.
+    Requests and responses are plain dicts.
+    """
+
+    def __init__(self, api_client: BaseApiClient):
+        self._api_client = api_client
+
+    def create(self, request_dict: dict[str, Any]) -> dict[str, Any]:
+        response = self._api_client.request("post", "interactions", request_dict)
+        return json.loads(response.body) if response.body else {}
+
+    def get(self, interaction_id: str) -> dict[str, Any]:
+        response = self._api_client.request("get", f"interactions/{interaction_id}", {})
+        return json.loads(response.body) if response.body else {}
+
+
+def _get_interactions_client(api_client: BaseApiClient) -> _InteractionsRestClient:
+    """Returns an Interactions API client bound to `api_client`.
+
+    The client issues calls through the SDK's existing `api_client` (a
+    `BaseApiClient`, or a `ReplayApiClient` in tests) so that replay recording
+    captures the interaction calls.
+
+    Args:
+        api_client: The API client used to issue interaction calls.
+
+    Returns:
+        An `_InteractionsRestClient`.
+    """
+    return _InteractionsRestClient(api_client)
+
+
+def _agent_data_response_text(agent_data: types.evals.AgentData) -> Optional[str]:
+    """Concatenates the text of all model-role events in an AgentData."""
+    text_parts: list[str] = []
+    for turn in agent_data.turns or []:
+        for event in turn.events or []:
+            content = event.content
+            if not content or content.role != _evals_constant.MODEL_AUTHOR:
+                continue
+            for part in content.parts or []:
+                if part.text:
+                    text_parts.append(part.text)
+    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],
+    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. 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`.
+        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:
+        The resolved interaction once it reaches a terminal state.
+
+    Raises:
+        TimeoutError: If the interaction does not complete within the timeout.
+    """
+    if interaction.get("status") in _INTERACTION_TERMINAL_STATES:
+        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)
+        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."
+    )
+
+
+def _run_gemini_agent_inference(
+    *,
+    api_client: BaseApiClient,
+    gemini_agent: str,
+    prompt_dataset: pd.DataFrame,
+) -> pd.DataFrame:
+    """Runs inference against a Gemini Agents API agent via the Interactions API.
+
+    For each prompt row, creates an interaction against `gemini_agent` and
+    collects the interaction id, response text, and agent data.
+
+    Args:
+        api_client: The API client used to issue interaction calls.
+        gemini_agent: The Gemini Agents API agent resource name.
+        prompt_dataset: The prompt DataFrame. The prompt is read from the
+          `request` column if present, otherwise from the `prompt` column.
+
+    Returns:
+        A DataFrame with columns prompt, response, interaction_id, agent_data.
+    """
+    prompt_column = (
+        "request" if "request" in prompt_dataset.columns else _evals_constant.PROMPT
+    )
+    if prompt_column not in prompt_dataset.columns:
+        raise ValueError(
+            "The eval dataset provided for Gemini agent inference must contain a"
+            f" '{_evals_constant.PROMPT}' or 'request' column."
+        )
+
+    interactions_client = _get_interactions_client(api_client)
+
+    agent_short_id = gemini_agent.split("/")[-1]
+    prompts: list[str] = []
+    responses: list[Optional[str]] = []
+    interaction_ids: list[Optional[str]] = []
+    agent_data: list[dict[str, Any]] = []
+    for prompt in tqdm(
+        prompt_dataset[prompt_column].tolist(), desc="Gemini Agent Inference"
+    ):
+        prompts.append(prompt)
+        try:
+            interaction = interactions_client.create(
+                {
+                    "agent": agent_short_id,
+                    "input": [{"type": "text", "text": prompt}],
+                    "store": True,
+                    "background": True,
+                }
+            )
+            interaction = _await_interaction(interactions_client, interaction)
+            agent_data_obj = _interaction_dict_to_agent_data(interaction)
+            _merge_text_parts_in_agent_data(agent_data_obj)
+            responses.append(_agent_data_response_text(agent_data_obj))
+            interaction_ids.append(interaction.get("id"))
+            agent_data.append(agent_data_obj.model_dump(mode="json", exclude_none=True))
+        except Exception as e:  # pylint: disable=broad-exception-caught
+            logger.warning(
+                "Gemini agent inference failed for a prompt (recording an empty"
+                " row and continuing): %s",
+                e,
+            )
+            responses.append(None)
+            interaction_ids.append(None)
+            agent_data.append({})
+
+    return pd.DataFrame(
+        {
+            _evals_constant.PROMPT: prompts,
+            _evals_constant.RESPONSE: responses,
+            _evals_constant.INTERACTION_ID: interaction_ids,
+            _evals_constant.AGENT_DATA: agent_data,
+        }
+    )
+
+
 def _normalize_interaction_resource(
     interaction: str, agent: str, location: Optional[str]
 ) -> str:
@@ -1838,6 +2049,7 @@ def _execute_inference(
     model: Optional[Union[Callable[[Any], Any], str]] = None,
     agent_engine: Optional[Union[str, types.AgentEngine]] = None,
     agent: Optional["LlmAgent"] = None,  # type: ignore # noqa: F821
+    gemini_agent: Optional[str] = None,
     dest: Optional[str] = None,
     config: Optional[genai_types.GenerateContentConfig] = None,
     prompt_template: Optional[Union[str, types.PromptTemplateOrDict]] = None,
@@ -1856,6 +2068,8 @@ def _execute_inference(
         agent_engine: The agent engine to use for inference. Can be a resource
           name string or an `AgentEngine` instance.
         agent: The local agent to use for inference. Can be an ADK agent instance.
+        gemini_agent: The Gemini Agents API agent resource name to run inference
+          against via the Interactions API.
         dest: The destination to save the inference results. Can be a string
           representing a file path or a GCS URI.
         config: The generation configuration for the model.
@@ -1873,9 +2087,10 @@ def _execute_inference(
     if location:
         api_client = _get_api_client_with_location(api_client, location)
 
-    if sum(x is not None for x in [model, agent_engine, agent]) != 1:
+    if sum(x is not None for x in [model, agent_engine, agent, gemini_agent]) != 1:
         raise ValueError(
-            "Exactly one of model, agent_engine, or agent must be provided."
+            "Exactly one of model, agent_engine, agent, or gemini_agent must be"
+            " provided."
         )
 
     prompt_dataset = _load_dataframe(api_client, src)
@@ -1888,7 +2103,24 @@ def _execute_inference(
 
         _apply_prompt_template(prompt_dataset, prompt_template)
 
-    if model:
+    if gemini_agent:
+        start_time = time.time()
+        logger.debug("Starting Gemini Agent inference process ...")
+        results_df = _run_gemini_agent_inference(
+            api_client=api_client,
+            gemini_agent=gemini_agent,
+            prompt_dataset=prompt_dataset,
+        )
+        end_time = time.time()
+        logger.info(
+            "Gemini Agent inference completed in %.2f seconds.",
+            end_time - start_time,
+        )
+        return types.EvaluationDataset(
+            eval_dataset_df=results_df,
+            candidate_name=gemini_agent.split("/")[-1],
+        )
+    elif model:
         start_time = time.time()
         logger.debug("Starting inference process ...")
         results_df = _run_inference_internal(

agentplatform/_genai/_evals_constant.py

@@ -56,6 +56,7 @@
 CONTENT = "content"
 PARTS = "parts"
 USER_AUTHOR = "user"
+MODEL_AUTHOR = "model"
 AGENT_DATA = "agent_data"
 INTERACTION_ID = "interaction_id"
 STARTING_PROMPT = "starting_prompt"