feat: GenAI Client(evals) - add user-facing generate_loss_clusters with LRO polling and replay tests

PiperOrigin-RevId: 893874547

Author: vertex-sdk-bot

tests/unit/vertexai/genai/replays/test_generate_loss_clusters.py

@@ -0,0 +1,76 @@
+# 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_loss_clusters(client):
+  """Tests that generate_loss_clusters() correctly calls the API and returns LossClusters."""
+  eval_result = types.EvaluationResult()
+  loss_clusters = client.evals.generate_loss_clusters(
+      eval_result=eval_result,
+      config=types.LossAnalysisConfig(
+          metric="multi_turn_task_success_v1",
+          candidate="travel-agent",
+      ),
+  )
+  assert type(loss_clusters).__name__ == "LossClusters"
+  assert len(loss_clusters.results) == 1
+  result = loss_clusters.results[0]
+  assert result.config.metric == "multi_turn_task_success_v1"
+  assert result.config.candidate == "travel-agent"
+  assert len(result.clusters) == 2
+  assert result.clusters[0].cluster_id == "cluster-1"
+  assert result.clusters[0].taxonomy_entry.l1_category == "Tool Calling"
+  assert (
+      result.clusters[0].taxonomy_entry.l2_category == "Missing Tool Invocation"
+  )
+  assert result.clusters[0].item_count == 3
+  assert result.clusters[1].cluster_id == "cluster-2"
+  assert result.clusters[1].taxonomy_entry.l1_category == "Hallucination"
+  assert result.clusters[1].item_count == 2
+
+
+pytest_plugins = ("pytest_asyncio",)
+
+
+@pytest.mark.asyncio
+async def test_gen_loss_clusters_async(client):
+  """Tests that generate_loss_clusters() async correctly calls the API and returns LossClusters."""
+  eval_result = types.EvaluationResult()
+  loss_clusters = await client.aio.evals.generate_loss_clusters(
+      eval_result=eval_result,
+      config=types.LossAnalysisConfig(
+          metric="multi_turn_task_success_v1",
+          candidate="travel-agent",
+      ),
+  )
+  assert type(loss_clusters).__name__ == "LossClusters"
+  assert len(loss_clusters.results) == 1
+  result = loss_clusters.results[0]
+  assert result.config.metric == "multi_turn_task_success_v1"
+  assert len(result.clusters) == 2
+  assert result.clusters[0].cluster_id == "cluster-1"
+  assert result.clusters[1].cluster_id == "cluster-2"
+
+
+pytestmark = pytest_helper.setup(
+    file=__file__,
+    globals_for_file=globals(),
+    test_method="evals.generate_loss_clusters",
+)

tests/unit/vertexai/genai/test_evals.py

@@ -29,13 +29,16 @@
 from google.cloud.aiplatform import initializer as aiplatform_initializer
 from vertexai import _genai
 from vertexai._genai import _evals_data_converters
+from vertexai._genai import _evals_utils
 from vertexai._genai import _evals_metric_handlers
 from vertexai._genai import _evals_visualization
 from vertexai._genai import _evals_metric_loaders
 from vertexai._genai import _gcs_utils
 from vertexai._genai import _observability_data_converter
+from vertexai._genai import _transformers
 from vertexai._genai import evals
 from vertexai._genai import types as vertexai_genai_types
+from vertexai._genai.types import common as common_types
 from google.genai import client
 from google.genai import errors as genai_errors
 from google.genai import types as genai_types
@@ -218,6 +221,133 @@ def test_get_api_client_with_none_location(
         mock_vertexai_client.assert_not_called()
 
 
+class TestTransformers:
+    """Unit tests for transformers."""
+
+    def test_t_inline_results(self):
+        eval_result = common_types.EvaluationResult(
+            eval_case_results=[
+                common_types.EvalCaseResult(
+                    eval_case_index=0,
+                    response_candidate_results=[
+                        common_types.ResponseCandidateResult(
+                            response_index=0,
+                            metric_results={
+                                "tool_use_quality": common_types.EvalCaseMetricResult(
+                                    score=0.0,
+                                    explanation="Failed tool use",
+                                )
+                            },
+                        )
+                    ],
+                )
+            ],
+            evaluation_dataset=[
+                common_types.EvaluationDataset(
+                    eval_cases=[
+                        common_types.EvalCase(
+                            prompt=genai_types.Content(
+                                parts=[genai_types.Part(text="test prompt")]
+                            )
+                        )
+                    ]
+                )
+            ],
+            metadata=common_types.EvaluationRunMetadata(candidate_names=["gemini-pro"]),
+        )
+
+        payload = _transformers.t_inline_results([eval_result])
+
+        assert len(payload) == 1
+        assert payload[0]["metric"] == "tool_use_quality"
+        assert payload[0]["request"]["prompt"]["text"] == "test prompt"
+        assert len(payload[0]["candidate_results"]) == 1
+        assert payload[0]["candidate_results"][0]["candidate"] == "gemini-pro"
+        assert payload[0]["candidate_results"][0]["score"] == 0.0
+
+
+class TestLossClusters:
+    """Unit tests for LossClusters and postprocessing."""
+
+    def test_postprocess_loss_clusters_response(self):
+        response = common_types.GenerateLossClustersResponse(
+            analysis_time="2026-04-01T10:00:00Z",
+            results=[
+                common_types.LossAnalysisResult(
+                    config=common_types.LossAnalysisConfig(
+                        metric="multi_turn_task_success_v1",
+                        candidate="travel-agent",
+                    ),
+                    analysis_time="2026-04-01T10:00:00Z",
+                    clusters=[
+                        common_types.LossCluster(
+                            cluster_id="cluster-1",
+                            taxonomy_entry=common_types.LossTaxonomyEntry(
+                                l1_category="Tool Calling",
+                                l2_category="Missing Tool Invocation",
+                                description="The agent failed to invoke a required tool.",
+                            ),
+                            item_count=3,
+                        ),
+                        common_types.LossCluster(
+                            cluster_id="cluster-2",
+                            taxonomy_entry=common_types.LossTaxonomyEntry(
+                                l1_category="Hallucination",
+                                l2_category="Hallucination of Action",
+                                description="Verbally confirmed action without tool.",
+                            ),
+                            item_count=2,
+                        ),
+                    ],
+                )
+            ],
+        )
+        loss_clusters = _evals_utils._postprocess_loss_clusters_response(response)
+        assert isinstance(loss_clusters, _evals_utils.LossClusters)
+        assert len(loss_clusters.results) == 1
+        assert loss_clusters.analysis_time == "2026-04-01T10:00:00Z"
+        result = loss_clusters.results[0]
+        assert result.config.metric == "multi_turn_task_success_v1"
+        assert len(result.clusters) == 2
+        assert result.clusters[0].cluster_id == "cluster-1"
+        assert result.clusters[0].item_count == 3
+        assert result.clusters[1].cluster_id == "cluster-2"
+
+    def test_loss_clusters_show_no_clusters(self, capsys):
+        response = common_types.GenerateLossClustersResponse(results=[])
+        loss_clusters = _evals_utils._postprocess_loss_clusters_response(response)
+        loss_clusters.show()
+        captured = capsys.readouterr()
+        assert "No loss clusters found" in captured.out
+
+    def test_loss_clusters_show_with_clusters(self, capsys):
+        response = common_types.GenerateLossClustersResponse(
+            results=[
+                common_types.LossAnalysisResult(
+                    config=common_types.LossAnalysisConfig(
+                        metric="test_metric",
+                        candidate="test-candidate",
+                    ),
+                    clusters=[
+                        common_types.LossCluster(
+                            cluster_id="c1",
+                            taxonomy_entry=common_types.LossTaxonomyEntry(
+                                l1_category="Cat1",
+                                l2_category="SubCat1",
+                            ),
+                            item_count=5,
+                        ),
+                    ],
+                )
+            ],
+        )
+        loss_clusters = _evals_utils._postprocess_loss_clusters_response(response)
+        loss_clusters.show()
+        captured = capsys.readouterr()
+        assert "Cat1" in captured.out
+        assert "SubCat1" in captured.out
+
+
 class TestEvals:
     """Unit tests for the GenAI client."""
 
@@ -1754,7 +1884,7 @@ def test_run_inference_with_litellm_openai_request_format(
         self,
         mock_api_client_fixture,
     ):
-        """Tests inference with LiteLLM where the row contains an chat completion request body."""
+        """Tests inference with LiteLLM where the row contains a chat completion request body."""
         with mock.patch(
             "vertexai._genai._evals_common.litellm"
         ) as mock_litellm, mock.patch(

vertexai/_genai/_evals_utils.py

@@ -15,9 +15,11 @@
 """Utility functions for evals."""
 
 import abc
+import asyncio
+import json
 import logging
 import os
-import json
+import time
 from typing import Any, Optional, Union
 
 from google.genai._api_client import BaseApiClient
@@ -366,6 +368,120 @@ def _postprocess_user_scenarios_response(
     )
 
 
+class LossClusters:
+    """User-friendly wrapper for loss analysis results with visualization."""
+
+    def __init__(
+        self,
+        response: types.GenerateLossClustersResponse,
+    ):
+        self._response = response
+        self.results = response.results or []
+        self.analysis_time = response.analysis_time
+
+    def show(self) -> None:
+        """Displays the loss clusters as a formatted pandas DataFrame."""
+        rows = []
+        for result in self.results:
+            metric = result.config.metric if result.config else None
+            candidate = result.config.candidate if result.config else None
+            for cluster in result.clusters or []:
+                entry = cluster.taxonomy_entry
+                row = {
+                    "metric": metric,
+                    "candidate": candidate,
+                    "cluster_id": cluster.cluster_id,
+                    "l1_category": entry.l1_category if entry else None,
+                    "l2_category": entry.l2_category if entry else None,
+                    "description": entry.description if entry else None,
+                    "item_count": cluster.item_count,
+                }
+                rows.append(row)
+
+        if not rows:
+            print("No loss clusters found.")  # pylint: disable=print-function
+            return
+
+        df = pd.DataFrame(rows)
+        try:
+            from IPython.display import display  # pylint: disable=g-import-not-at-top
+
+            display(df)
+        except ImportError:
+            print(df.to_string())  # pylint: disable=print-function
+
+
+def _postprocess_loss_clusters_response(
+    response: types.GenerateLossClustersResponse,
+) -> LossClusters:
+    """Wraps the GenerateLossClustersResponse in a LossClusters object."""
+    return LossClusters(response=response)
+
+
+def _poll_operation(
+    api_client: BaseApiClient,
+    operation_name: str,
+    poll_interval_seconds: float = 5.0,
+) -> types.GenerateLossClustersOperation:
+    """Polls a long-running operation until completion.
+
+    Args:
+        api_client: The API client to use for polling.
+        operation_name: The full resource name of the operation.
+        poll_interval_seconds: Time between polls.
+
+    Returns:
+        The completed operation.
+    """
+    start_time = time.time()
+    while True:
+        response = api_client.request("get", operation_name, {}, None)
+        response_dict = {} if not response.body else json.loads(response.body)
+        operation = types.GenerateLossClustersOperation._from_response(
+            response=response_dict, kwargs={}
+        )
+        if operation.done:
+            return operation
+        elapsed = int(time.time() - start_time)
+        logger.info(
+            "Loss analysis operation still running... Elapsed time: %d seconds",
+            elapsed,
+        )
+        time.sleep(poll_interval_seconds)
+
+
+async def _poll_operation_async(
+    api_client: BaseApiClient,
+    operation_name: str,
+    poll_interval_seconds: float = 5.0,
+) -> types.GenerateLossClustersOperation:
+    """Polls a long-running operation until completion (async).
+
+    Args:
+        api_client: The API client to use for polling.
+        operation_name: The full resource name of the operation.
+        poll_interval_seconds: Time between polls.
+
+    Returns:
+        The completed operation.
+    """
+    start_time = time.time()
+    while True:
+        response = await api_client.async_request("get", operation_name, {}, None)
+        response_dict = {} if not response.body else json.loads(response.body)
+        operation = types.GenerateLossClustersOperation._from_response(
+            response=response_dict, kwargs={}
+        )
+        if operation.done:
+            return operation
+        elapsed = int(time.time() - start_time)
+        logger.info(
+            "Loss analysis operation still running... Elapsed time: %d seconds",
+            elapsed,
+        )
+        await asyncio.sleep(poll_interval_seconds)
+
+
 def _validate_dataset_agent_data(
     dataset: types.EvaluationDataset,
     inference_configs: Optional[dict[str, Any]] = None,