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

PiperOrigin-RevId: 945262438

Author: jsondai

agentplatform/_genai/_evals_common.py

@@ -32,6 +32,7 @@
 import agentplatform
 from google.genai import types as genai_types
 from google.genai._api_client import BaseApiClient
+from google.genai._gaos.google_genai import GeminiNextGenInteractions
 from google.genai.models import Models
 import pandas as pd
 from tqdm import tqdm
@@ -575,6 +576,269 @@ def _is_gemini_agent_resource(agent: str) -> bool:
     )
 
 
+def _get_resolved_location(api_client: Any) -> Optional[str]:
+    """Returns the location configured on the API client."""
+    return getattr(api_client, "location", None)
+
+
+def _get_interactions_client(api_client: BaseApiClient) -> Any:
+    """Returns a google.genai interactions client bound to `api_client`.
+
+    The interactions client is constructed around the SDK's existing
+    `api_client` (a `BaseApiClient`, or a `ReplayApiClient` in tests) so that
+    replay recording captures the interaction calls. This mirrors how
+    `google.genai.client.Client.interactions` is built internally.
+
+    Args:
+        api_client: The API client to bind the interactions client to.
+
+    Returns:
+        A google.genai interactions client.
+    """
+    return GeminiNextGenInteractions(api_client)
+
+
+def _interaction_to_agent_data(
+    gemini_agent: str, prompt: str, interaction: Any
+) -> dict[str, Any]:
+    """Maps a Gemini Agents interaction into the canonical AgentData schema.
+
+    Populates the AgentConfig for the agent plus a single ConversationTurn with
+    the user input event and the final model output event. The interaction's
+    full step trajectory is not expanded event-by-event; the response text is
+    used as the model output content.
+
+    Args:
+        gemini_agent: The Gemini Agents API agent resource name.
+        prompt: The user prompt for the interaction.
+        interaction: The interaction returned by the Interactions API.
+
+    Returns:
+        A dict representation of the AgentData for the interaction.
+    """
+    agent_id = gemini_agent.split("/")[-1]
+    agents = {
+        agent_id: types.evals.AgentConfig(  # pytype: disable=missing-parameter
+            agent_id=agent_id,
+        )
+    }
+    events = [
+        types.evals.AgentEvent(
+            author=_evals_constant.USER_AUTHOR,
+            content=genai_types.Content(
+                role=_evals_constant.USER_AUTHOR,
+                parts=[genai_types.Part(text=prompt)],
+            ),
+        )
+    ]
+    output_text = getattr(interaction, "output_text", None)
+    if output_text:
+        events.append(
+            types.evals.AgentEvent(
+                author=agent_id,
+                content=genai_types.Content(
+                    role=_evals_constant.MODEL_AUTHOR,
+                    parts=[genai_types.Part(text=output_text)],
+                ),
+            )
+        )
+    turns = [
+        types.evals.ConversationTurn(
+            turn_index=0,
+            turn_id=getattr(interaction, "id", None) or "turn_0",
+            events=events,
+        )
+    ]
+    return types.evals.AgentData(
+        agents=agents,
+        turns=turns,
+    ).model_dump(mode="json", exclude_none=True)
+
+
+_INTERACTION_TERMINAL_STATES = frozenset(
+    ["completed", "failed", "cancelled", "incomplete", "budget_exceeded"]
+)
+
+
+def _await_interaction(
+    interactions_client: Any,
+    interaction: Any,
+    poll_interval_seconds: float = 2.0,
+    timeout_seconds: float = 600.0,
+) -> 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 getattr(interaction, "status", None) in _INTERACTION_TERMINAL_STATES:
+        return interaction
+    interaction_id = getattr(interaction, "id", None)
+    deadline = time.monotonic() + timeout_seconds
+    while time.monotonic() < deadline:
+        time.sleep(poll_interval_seconds)
+        interaction = interactions_client.get(interaction_id)
+        if getattr(interaction, "status", None) 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)
+
+    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"
+    ):
+        interaction = interactions_client.create(
+            agent=gemini_agent,
+            input=prompt,
+            store=True,
+            background=True,
+        )
+        interaction = _await_interaction(interactions_client, interaction)
+        prompts.append(prompt)
+        responses.append(getattr(interaction, "output_text", None))
+        interaction_ids.append(getattr(interaction, "id", None))
+        agent_data.append(_interaction_to_agent_data(gemini_agent, prompt, interaction))
+
+    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"
+            " evaluating interaction ids. Got: %s" % 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("Missing `interaction_id` value for row %d." % i)
+        eval_cases.append(
+            types.EvalCase(
+                eval_case_id="eval_case_%s" % 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,
@@ -1406,6 +1670,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,
@@ -1424,6 +1689,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.
@@ -1441,9 +1708,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)
@@ -1456,7 +1724,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(
@@ -1591,6 +1876,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.
 
@@ -1640,6 +1927,13 @@ 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:
+            parsed_evaluation_datasets.append(interaction_dataset)
+            continue
+
         if dataset_schema:
             current_schema = _evals_data_converters.EvalDatasetSchema(dataset_schema)
         else:
@@ -1797,6 +2091,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,
@@ -1877,6 +2172,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,
     }
 )