fix: GenAI Client(evals) - patch for vulnerability in visualization

jsondai · copybara-github · commit cd797d8f3097 · 2025-12-08T14:07:10.000-08:00
PiperOrigin-RevId: 840870138
diff --git a/tests/unit/vertexai/genai/test_evals.py b/tests/unit/vertexai/genai/test_evals.py
@@ -13,9 +13,11 @@
 # limitations under the License.
 #
 # pylint: disable=protected-access,bad-continuation,
+import base64
 import importlib
 import json
 import os
+import re
 import statistics
 import sys
 from unittest import mock
@@ -291,8 +293,11 @@ def test_display_evaluation_result_with_agent_trace_prefixes(self, mock_is_ipyth
 
         mock_display_module.HTML.assert_called_once()
         html_content = mock_display_module.HTML.call_args[0][0]
-        assert "my_function" in html_content
-        assert "this is model response" in html_content
+        match = re.search(r'atob\("([^"]+)"\)', html_content)
+        assert match
+        decoded_json = base64.b64decode(match.group(1)).decode("utf-8")
+        assert "my_function" in decoded_json
+        assert "this is model response" in decoded_json
 
         del sys.modules["IPython"]
         del sys.modules["IPython.display"]
diff --git a/vertexai/_genai/_evals_visualization.py b/vertexai/_genai/_evals_visualization.py
@@ -14,6 +14,7 @@
 #
 """Visualization utilities for GenAI Evaluation SDK."""
 
+import base64
 import json
 import logging
 from typing import Any, Optional
@@ -80,6 +81,7 @@ def stringify_cell(cell: Any) -> Optional[str]:
 
 def _get_evaluation_html(eval_result_json: str) -> str:
     """Returns a self-contained HTML for single evaluation visualization."""
+    payload_b64 = base64.b64encode(eval_result_json.encode("utf-8")).decode("utf-8")
     return f"""
 <!DOCTYPE html>
 <html>
@@ -254,7 +256,7 @@ def _get_evaluation_html(eval_result_json: str) -> str:
         <div id="details-section"></div>
     </div>
     <script>
-        var vizData_vertex_eval_sdk = {eval_result_json};
+        var vizData_vertex_eval_sdk = JSON.parse(atob("{payload_b64}"));
         function formatDictVals(obj) {{
             if (typeof obj === 'string') return obj;
             if (obj === undefined || obj === null) return '';
@@ -556,6 +558,7 @@ def _get_evaluation_html(eval_result_json: str) -> str:
 
 def _get_comparison_html(eval_result_json: str) -> str:
     """Returns a self-contained HTML for a side-by-side eval comparison."""
+    payload_b64 = base64.b64encode(eval_result_json.encode("utf-8")).decode("utf-8")
     return f"""
 <!DOCTYPE html>
 <html>
@@ -616,7 +619,7 @@ def _get_comparison_html(eval_result_json: str) -> str:
         <div id="details-section"></div>
     </div>
     <script>
-        var vizData_vertex_eval_sdk = {eval_result_json};
+        var vizData_vertex_eval_sdk = JSON.parse(atob("{payload_b64}"));
         function renderSummary(summaryMetrics, metadata) {{
             const container = document.getElementById('summary-section');
             if (!summaryMetrics || summaryMetrics.length === 0) {{ container.innerHTML = '<h2>Summary Metrics</h2><p>No summary metrics.</p>'; return; }}
@@ -696,6 +699,7 @@ def _get_comparison_html(eval_result_json: str) -> str:
 
 def _get_inference_html(dataframe_json: str) -> str:
     """Returns a self-contained HTML for displaying inference results."""
+    payload_b64 = base64.b64encode(dataframe_json.encode("utf-8")).decode("utf-8")
     return f"""
 <!DOCTYPE html>
 <html>
@@ -746,7 +750,7 @@ def _get_inference_html(dataframe_json: str) -> str:
         <div id="results-table"></div>
     </div>
     <script>
-        var vizData_vertex_eval_sdk = {dataframe_json};
+        var vizData_vertex_eval_sdk = JSON.parse(atob("{payload_b64}"));
         var container_vertex_eval_sdk = document.getElementById('results-table');
 
         function renderRubrics(cellValue) {{
