feat: GenAI Client(evals) - run_inference support for Gemini Agents API agents

PiperOrigin-RevId: 945262438

Author: jsondai

agentplatform/_genai/_evals_common.py

@@ -923,6 +923,260 @@ def _fetch_agent_config_dict(
     )
 
 
+def _get_resolved_location(api_client: Any) -> Optional[str]:
+    """Returns the location configured on the API client."""
+    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
+
+
+_INTERACTION_TERMINAL_STATES = frozenset(
+    ["completed", "failed", "cancelled", "incomplete", "budget_exceeded"]
+)
+
+
+def _await_interaction(
+    interactions_client: "_InteractionsRestClient",
+    interaction: dict[str, Any],
+    poll_interval_seconds: float = 2.0,
+    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.
+
+    Args:
+        interactions_client: The interactions client used to poll.
+        interaction: The interaction returned by `create`.
+        poll_interval_seconds: Delay between poll attempts.
+        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
+    while time.monotonic() < deadline:
+        time.sleep(poll_interval_seconds)
+        interaction = interactions_client.get(interaction_id)
+        if interaction.get("status") in _INTERACTION_TERMINAL_STATES:
+            return interaction
+    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:
+    """Normalizes an interaction id into a full resource name.
+
+    A bare interaction id is expanded to
+    `projects/{project}/locations/{location}/interactions/{id}` using the
+    project and location parsed from the agent resource name. Fully-qualified
+    interaction resource names are returned unchanged.
+    """
+    if interaction.startswith("projects/"):
+        return interaction
+    parts = agent.split("/")
+    project = parts[1]
+    agent_location = parts[3] if len(parts) > 3 else (location or "global")
+    return f"projects/{project}/locations/{agent_location}/interactions/{interaction}"
+
+
+def _build_interaction_id_dataset(
+    loaded_data: list[dict[str, Any]],
+    agent: Optional[str],
+    location: Optional[str],
+) -> Optional[types.EvaluationDataset]:
+    """Builds an EvaluationDataset from rows that carry an `interaction_id`.
+
+    When the dataset contains an `interaction_id` column, each row is turned
+    into an EvalCase whose `interactions_data_source` references the interaction
+    and the Gemini agent. The backend resolves the interaction trace and agent
+    config; no client-side prompt/response is required. Returns None if the
+    data does not contain interaction ids.
+    """
+    has_interaction_id = bool(loaded_data) and any(
+        _evals_constant.INTERACTION_ID in row for row in loaded_data
+    )
+    if not has_interaction_id:
+        if agent:
+            raise ValueError(
+                "An `agent` was provided but the dataset does not contain an"
+                " `interaction_id` column. The `agent` argument is only used to"
+                " resolve an `interaction_id` dataset column (so the backend can"
+                " fetch the interaction trace and Agent config). To evaluate"
+                " with an agent, provide a dataset with an `interaction_id`"
+                " column; otherwise omit `agent`."
+            )
+        return None
+
+    if not agent:
+        raise ValueError(
+            "An `agent` resource name is required when the dataset contains an"
+            " `interaction_id` column, so the backend can resolve the Agent"
+            " config for each interaction."
+        )
+    if not _is_gemini_agent_resource(agent):
+        raise ValueError(
+            "`agent` must be a Gemini Agents API resource name of the form"
+            " projects/{project}/locations/{location}/agents/{agent} when"
+            f" evaluating interaction ids. Got: {agent}"
+        )
+
+    gemini_agent_config = types.GeminiAgentConfig(gemini_agent=agent)
+    eval_cases = []
+    for i, row in enumerate(loaded_data):
+        interaction = row.get(_evals_constant.INTERACTION_ID)
+        if not interaction:
+            raise ValueError(f"Missing `interaction_id` value for row {i}.")
+        eval_cases.append(
+            types.EvalCase(
+                eval_case_id=f"eval_case_{i}",
+                interactions_data_source=types.InteractionsDataSource(
+                    interaction=_normalize_interaction_resource(
+                        str(interaction), agent, location
+                    ),
+                    gemini_agent_config=gemini_agent_config,
+                ),
+            )
+        )
+    return types.EvaluationDataset(eval_cases=eval_cases)
+
+
 def _add_evaluation_run_labels(
     labels: Optional[dict[str, str]] = None,
     agent: Optional[str] = None,
@@ -1754,6 +2008,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,
@@ -1772,6 +2027,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.
@@ -1789,9 +2046,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)
@@ -1804,7 +2062,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(
@@ -1939,6 +2214,8 @@ def _resolve_dataset_inputs(
     dataset_schema: Optional[Literal["GEMINI", "FLATTEN", "OPENAI"]],
     loader: "_evals_utils.EvalDatasetLoader",
     agent_info: Optional[types.evals.AgentInfo] = None,
+    agent: Optional[str] = None,
+    api_client: Any = None,
 ) -> tuple[types.EvaluationDataset, int]:
     """Loads and processes single or multiple datasets for evaluation.
 
@@ -1988,6 +2265,21 @@ def _resolve_dataset_inputs(
         ds_source_for_loader = _get_dataset_source(ds_item)
         current_loaded_data = loader.load(ds_source_for_loader)
 
+        interaction_dataset = _build_interaction_id_dataset(
+            current_loaded_data, agent, _get_resolved_location(api_client)
+        )
+        if interaction_dataset is not None:
+            if dataset_schema:
+                raise ValueError(
+                    "`dataset_schema` is not supported for datasets with an"
+                    " `interaction_id` column. The interaction trace and agent"
+                    " config are resolved by the backend, so no client-side"
+                    " schema conversion is applied. Omit `dataset_schema` when"
+                    " evaluating an interaction_id dataset."
+                )
+            parsed_evaluation_datasets.append(interaction_dataset)
+            continue
+
         if dataset_schema:
             current_schema = _evals_data_converters.EvalDatasetSchema(dataset_schema)
         else:
@@ -2145,6 +2437,7 @@ def _execute_evaluation(  # type: ignore[no-untyped-def]
     api_client: Any,
     dataset: Union[types.EvaluationDataset, list[types.EvaluationDataset]],
     metrics: list[types.Metric],
+    agent: Optional[str] = None,
     dataset_schema: Optional[Literal["GEMINI", "FLATTEN", "OPENAI"]] = None,
     dest: Optional[str] = None,
     location: Optional[str] = None,
@@ -2225,6 +2518,8 @@ def _execute_evaluation(  # type: ignore[no-untyped-def]
         dataset_schema=dataset_schema,
         loader=loader,
         agent_info=validated_agent_info,
+        agent=agent,
+        api_client=api_client,
     )
 
     resolved_metrics = _resolve_metrics(metrics, api_client)

agentplatform/_genai/_evals_constant.py

@@ -56,7 +56,9 @@
 CONTENT = "content"
 PARTS = "parts"
 USER_AUTHOR = "user"
+MODEL_AUTHOR = "model"
 AGENT_DATA = "agent_data"
+INTERACTION_ID = "interaction_id"
 STARTING_PROMPT = "starting_prompt"
 CONVERSATION_PLAN = "conversation_plan"
 HISTORY = "history"
@@ -74,5 +76,6 @@
         STARTING_PROMPT,
         CONVERSATION_PLAN,
         AGENT_DATA,
+        INTERACTION_ID,
     }
 )