Skip to content

Commit 8a00d43

Browse files
jsondaicopybara-github
jsondai
authored and
copybara-github
committed
fix: GenAI Client(evals) - patch for vulnerability in visualization
PiperOrigin-RevId: 844909702
1 parent df0976e commit 8a00d43

File tree

2 files changed

+100
-18
lines changed

2 files changed

+100
-18
lines changed

‎tests/unit/vertexai/genai/test_evals.py‎

Lines changed: 69 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -13,9 +13,11 @@
1313
# limitations under the License.
1414
#
1515
# pylint: disable=protected-access,bad-continuation,
16+
import base64
1617
import importlib
1718
import json
1819
import os
20+
import re
1921
import statistics
2022
import sys
2123
from unittest import mock
@@ -291,8 +293,72 @@ def test_display_evaluation_result_with_agent_trace_prefixes(self, mock_is_ipyth
291293

292294
mock_display_module.HTML.assert_called_once()
293295
html_content = mock_display_module.HTML.call_args[0][0]
294-
assert "my_function" in html_content
295-
assert "this is model response" in html_content
296+
match = re.search(r'atob\("([^"]+)"\)', html_content)
297+
assert match
298+
decoded_json = base64.b64decode(match.group(1)).decode("utf-8")
299+
assert "my_function" in decoded_json
300+
assert "this is model response" in decoded_json
301+
302+
del sys.modules["IPython"]
303+
del sys.modules["IPython.display"]
304+
305+
@mock.patch(
306+
"vertexai._genai._evals_visualization._is_ipython_env",
307+
return_value=True,
308+
)
309+
def test_display_evaluation_result_with_non_ascii_character(self, mock_is_ipython):
310+
"""Tests that non-ASCII characters are handled correctly."""
311+
mock_display_module = mock.MagicMock()
312+
mock_ipython_module = mock.MagicMock()
313+
mock_ipython_module.display = mock_display_module
314+
sys.modules["IPython"] = mock_ipython_module
315+
sys.modules["IPython.display"] = mock_display_module
316+
317+
dataset_df = pd.DataFrame(
318+
[
319+
{
320+
"prompt": "Test prompt with emoji 😊",
321+
"response": "Test response with emoji 😊",
322+
},
323+
]
324+
)
325+
eval_dataset = vertexai_genai_types.EvaluationDataset(
326+
eval_dataset_df=dataset_df
327+
)
328+
eval_result = vertexai_genai_types.EvaluationResult(
329+
evaluation_dataset=[eval_dataset],
330+
eval_case_results=[
331+
vertexai_genai_types.EvalCaseResult(
332+
eval_case_index=0,
333+
response_candidate_results=[
334+
vertexai_genai_types.ResponseCandidateResult(
335+
response_index=0, metric_results={}
336+
)
337+
],
338+
)
339+
],
340+
)
341+
342+
_evals_visualization.display_evaluation_result(eval_result)
343+
344+
mock_display_module.HTML.assert_called_once()
345+
html_content = mock_display_module.HTML.call_args[0][0]
346+
# Verify that the new decoding logic is present in the HTML
347+
assert "new TextDecoder().decode" in html_content
348+
349+
match = re.search(r'atob\("([^"]+)"\)', html_content)
350+
assert match
351+
decoded_json = base64.b64decode(match.group(1)).decode("utf-8")
352+
353+
# JSON serialization escapes non-ASCII characters (e.g. \uXXXX), so we
354+
# parse it back to check for the actual characters.
355+
parsed_json = json.loads(decoded_json)
356+
assert "Test prompt with emoji 😊" in json.dumps(
357+
parsed_json, ensure_ascii=False
358+
)
359+
assert "Test response with emoji 😊" in json.dumps(
360+
parsed_json, ensure_ascii=False
361+
)
296362

297363
del sys.modules["IPython"]
298364
del sys.modules["IPython.display"]
@@ -1290,7 +1356,7 @@ def test_run_inference_with_agent_engine_with_response_column_raises_error(
12901356
) in str(excinfo.value)
12911357

12921358
@mock.patch.object(_evals_utils, "EvalDatasetLoader")
1293-
@mock.patch("vertexai._genai._evals_common.InMemorySessionService")
1359+
@mock.patch("vertexai._genai._evals_common.InMemorySessionService") # fmt: skip
12941360
@mock.patch("vertexai._genai._evals_common.Runner")
12951361
@mock.patch("vertexai._genai._evals_common.LlmAgent")
12961362
def test_run_inference_with_local_agent(

‎vertexai/_genai/_evals_visualization.py‎

Lines changed: 31 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -14,8 +14,10 @@
1414
#
1515
"""Visualization utilities for GenAI Evaluation SDK."""
1616

17+
import base64
1718
import json
1819
import logging
20+
import textwrap
1921
from typing import Any, Optional
2022

2123
import pandas as pd
@@ -78,9 +80,16 @@ def stringify_cell(cell: Any) -> Optional[str]:
7880
return df_copy
7981

8082

83+
def _encode_to_base64(data: str) -> str:
84+
"""Encodes a string to a web-safe Base64 string."""
85+
return base64.b64encode(data.encode("utf-8")).decode("utf-8")
86+
87+
8188
def _get_evaluation_html(eval_result_json: str) -> str:
8289
"""Returns a self-contained HTML for single evaluation visualization."""
83-
return f"""
90+
payload_b64 = _encode_to_base64(eval_result_json)
91+
return textwrap.dedent(
92+
f"""
8493
<!DOCTYPE html>
8594
<html>
8695
<head>
@@ -249,12 +258,11 @@ def _get_evaluation_html(eval_result_json: str) -> str:
249258
<body>
250259
<div class="container">
251260
<h1>Evaluation Report</h1>
252-
<div id="summary-section"></div>
253-
<div id="agent-info-section"></div>
254-
<div id="details-section"></div>
255-
</div>
261+
< <div id="summary-section"></div>
262+
<div id="agent-info-section"></div>
263+
<div id="details-section"></div>
256264
<script>
257-
var vizData_vertex_eval_sdk = {eval_result_json};
265+
var vizData_vertex_eval_sdk = JSON.parse(new TextDecoder().decode(Uint8Array.from(atob("{payload_b64}"), c => c.charCodeAt(0))));
258266
function formatDictVals(obj) {{
259267
if (typeof obj === 'string') return obj;
260268
if (obj === undefined || obj === null) return '';
@@ -552,11 +560,14 @@ def _get_evaluation_html(eval_result_json: str) -> str:
552560
</body>
553561
</html>
554562
"""
563+
)
555564

556565

557566
def _get_comparison_html(eval_result_json: str) -> str:
558567
"""Returns a self-contained HTML for a side-by-side eval comparison."""
559-
return f"""
568+
payload_b64 = _encode_to_base64(eval_result_json)
569+
return textwrap.dedent(
570+
f"""
560571
<!DOCTYPE html>
561572
<html>
562573
<head>
@@ -612,11 +623,10 @@ def _get_comparison_html(eval_result_json: str) -> str:
612623
<body>
613624
<div class="container">
614625
<h1>Eval Comparison Report</h1>
615-
<div id="summary-section"></div>
616-
<div id="details-section"></div>
617-
</div>
626+
< <div id="summary-section"></div>
627+
<div id="details-section"></div>
618628
<script>
619-
var vizData_vertex_eval_sdk = {eval_result_json};
629+
var vizData_vertex_eval_sdk = JSON.parse(new TextDecoder().decode(Uint8Array.from(atob("{payload_b64}"), c => c.charCodeAt(0))));
620630
function renderSummary(summaryMetrics, metadata) {{
621631
const container = document.getElementById('summary-section');
622632
if (!summaryMetrics || summaryMetrics.length === 0) {{ container.innerHTML = '<h2>Summary Metrics</h2><p>No summary metrics.</p>'; return; }}
@@ -692,11 +702,14 @@ def _get_comparison_html(eval_result_json: str) -> str:
692702
</body>
693703
</html>
694704
"""
705+
)
695706

696707

697708
def _get_inference_html(dataframe_json: str) -> str:
698709
"""Returns a self-contained HTML for displaying inference results."""
699-
return f"""
710+
payload_b64 = _encode_to_base64(dataframe_json)
711+
return textwrap.dedent(
712+
f"""
700713
<!DOCTYPE html>
701714
<html>
702715
<head>
@@ -741,12 +754,12 @@ def _get_inference_html(dataframe_json: str) -> str:
741754
</style>
742755
</head>
743756
<body>
744-
<div class="container">
757+
< <div class="container">
745758
<h1>Evaluation Dataset</h1>
746759
<div id="results-table"></div>
747760
</div>
748761
<script>
749-
var vizData_vertex_eval_sdk = {dataframe_json};
762+
var vizData_vertex_eval_sdk = JSON.parse(new TextDecoder().decode(Uint8Array.from(atob("{payload_b64}"), c => c.charCodeAt(0))));
750763
var container_vertex_eval_sdk = document.getElementById('results-table');
751764
752765
function renderRubrics(cellValue) {{
@@ -822,6 +835,7 @@ def _get_inference_html(dataframe_json: str) -> str:
822835
</body>
823836
</html>
824837
"""
838+
)
825839

826840

827841
def _extract_text_and_raw_json(content: Any) -> dict[str, str]:
@@ -1086,12 +1100,14 @@ def _get_status_html(status: str, error_message: Optional[str] = None) -> str:
10861100
</p>
10871101
"""
10881102

1089-
return f"""
1103+
return textwrap.dedent(
1104+
f"""
10901105
<div>
10911106
<p><b>Status:</b> {status}</p>
10921107
{error_html}
10931108
</div>
10941109
"""
1110+
)
10951111

10961112

10971113
def display_evaluation_run_status(eval_run_obj: "types.EvaluationRun") -> None:

0 commit comments

Comments
 (0)