From 16489d32d124f27d2bc7ff45f5d46cd6e5e5a3b4 Mon Sep 17 00:00:00 2001
From: Jason Dai <jsndai@google.com>
Date: Tue, 6 Oct 2026 13:25:28 -0700
Subject: [PATCH] chore: GenAI Client(evals) - Add public alias for
 extract_dataset_rows.

PiperOrigin-RevId: 994589437
---
 agentplatform/_genai/_evals_visualization.py |  3 ++
 tests/unit/agentplatform/genai/test_evals.py | 30 ++++++++++++++++++++
 2 files changed, 33 insertions(+)

diff --git a/agentplatform/_genai/_evals_visualization.py b/agentplatform/_genai/_evals_visualization.py
index b6d335e7c5..c529b6f14e 100644
--- a/agentplatform/_genai/_evals_visualization.py
+++ b/agentplatform/_genai/_evals_visualization.py
@@ -232,6 +232,9 @@ def _extract_dataset_rows(dataset: types.EvaluationDataset) -> list[dict[str, An
     return processed_rows
 
 
+extract_dataset_rows = _extract_dataset_rows
+
+
 def get_evaluation_html(eval_result_json: str) -> str:
     """Returns a self-contained HTML for single evaluation visualization."""
     payload_b64 = _encode_to_base64(eval_result_json)
diff --git a/tests/unit/agentplatform/genai/test_evals.py b/tests/unit/agentplatform/genai/test_evals.py
index 28807dfc67..22e368fd56 100644
--- a/tests/unit/agentplatform/genai/test_evals.py
+++ b/tests/unit/agentplatform/genai/test_evals.py
@@ -844,6 +844,36 @@ def test_get_inference_html(self):
         assert payload_b64 in html
         assert "DOMPurify" in html
 
+    def test_extract_dataset_rows(self):
+        dataset = common_types.EvaluationDataset(
+            eval_cases=[
+                common_types.EvalCase(
+                    prompt=genai_types.Content(
+                        parts=[genai_types.Part(text="What is 2+2?")]
+                    ),
+                    responses=[
+                        common_types.ResponseCandidate(
+                            response=genai_types.Content(
+                                parts=[genai_types.Part(text="4")]
+                            )
+                        )
+                    ],
+                    reference=common_types.ResponseCandidate(
+                        response=genai_types.Content(
+                            parts=[genai_types.Part(text="Four")]
+                        )
+                    ),
+                )
+            ]
+        )
+
+        rows = _evals_visualization.extract_dataset_rows(dataset)
+
+        assert len(rows) == 1
+        assert rows[0]["prompt_display_text"] == "What is 2+2?"
+        assert rows[0]["response_display_text"] == "4"
+        assert rows[0]["reference"] == "Four"
+
     def test_display_loss_clusters_response_no_ipython(self):
         """Tests graceful fallback when not in IPython."""
         from agentplatform._genai import _evals_visualization