Skip to content

Commit b28737e

Browse files
jsondaicopybara-github
jsondai
authored and
copybara-github
committed
fix: GenAI Client(evals) - Treat NaN and array cells as missing in eval DataFrames
PiperOrigin-RevId: 996139461
1 parent 4d4ac8a commit b28737e

3 files changed

Lines changed: 127 additions & 13 deletions

File tree

‎agentplatform/_genai/_evals_common.py‎

Lines changed: 14 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -123,6 +123,11 @@ def _temp_logger_level(logger_name: str, level: int) -> None: # type: ignore[mi
123123
logger_instance.setLevel(original_level)
124124

125125

126+
def _is_missing(value: Any) -> bool:
127+
"""Returns True for None or a scalar NaN, such as a missing DataFrame cell."""
128+
return bool(pd.api.types.is_scalar(value) and pd.isna(value))
129+
130+
126131
def _get_api_client_with_location(
127132
api_client: BaseApiClient, location: Optional[str]
128133
) -> BaseApiClient:
@@ -1585,7 +1590,7 @@ def _build_interaction_id_dataset(
15851590
eval_cases = []
15861591
for i, row in enumerate(loaded_data):
15871592
interaction = row.get(_evals_constant.INTERACTION_ID)
1588-
if not interaction:
1593+
if _is_missing(interaction) or not interaction:
15891594
raise ValueError(f"Missing `interaction_id` value for row {i}.")
15901595
eval_cases.append(
15911596
types.EvalCase(
@@ -1895,12 +1900,12 @@ def _execute_inference_concurrently(
18951900
max_workers = AGENT_MAX_WORKERS if runtime or agent else MAX_WORKERS
18961901
with tqdm(total=len(prompt_dataset), desc=progress_desc) as pbar:
18971902
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
1898-
for index, row in prompt_dataset.iterrows():
1903+
for index, (_, row) in enumerate(prompt_dataset.iterrows()):
18991904
try:
19001905
if (
19011906
has_agent_data
19021907
and AGENT_DATA in row.index
1903-
and row.get(AGENT_DATA) is not None
1908+
and not _is_missing(row.get(AGENT_DATA))
19041909
):
19051910
agent_data_obj = row[AGENT_DATA]
19061911
if isinstance(agent_data_obj, dict):
@@ -2125,7 +2130,7 @@ def _run_litellm_inference(
21252130

21262131
with tqdm(total=len(prompt_dataset), desc=f"LiteLLM Inference ({model})") as pbar:
21272132
with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
2128-
for index, row in prompt_dataset.iterrows():
2133+
for index, (_, row) in enumerate(prompt_dataset.iterrows()):
21292134
messages = _convert_prompt_row_to_litellm_messages(row)
21302135
future = executor.submit(
21312136
_call_litellm_completion, model=model, messages=messages
@@ -3866,7 +3871,7 @@ def _drop_empty_columns(df: "pd.DataFrame") -> "pd.DataFrame":
38663871
def is_empty(x: Any) -> bool:
38673872
if isinstance(x, (list, dict)):
38683873
return not x
3869-
return pd.isna(x) # type: ignore[no-any-return]
3874+
return _is_missing(x)
38703875

38713876
cols_to_drop = [col for col in df.columns if df[col].apply(is_empty).all()]
38723877
return df.drop(columns=cols_to_drop)
@@ -4165,6 +4170,8 @@ def _create_evaluation_set_from_dataframe(
41654170
if _evals_constant.RESPONSE in row or agent_data_obj or intermediate_events:
41664171
# Resolve the oneof conflict: prioritize agent_data over flat text
41674172
response_text = row.get(_evals_constant.RESPONSE) or None
4173+
if _is_missing(response_text):
4174+
response_text = None
41684175

41694176
if agent_data_obj and response_text:
41704177
logger.info(
@@ -4219,15 +4226,15 @@ def _create_evaluation_set_from_dataframe(
42194226
prompt = types.EvaluationPrompt(
42204227
prompt_template_data=types.PromptTemplateData(values=values)
42214228
)
4222-
elif _evals_constant.PROMPT in row:
4229+
elif not _is_missing(row.get(_evals_constant.PROMPT)):
42234230
prompt = types.EvaluationPrompt(text=row[_evals_constant.PROMPT])
42244231

42254232
eval_item_requests.append(
42264233
types.EvaluationItemRequest(
42274234
prompt=prompt or None,
42284235
golden_response=(
42294236
types.CandidateResponse(text=row[_evals_constant.REFERENCE])
4230-
if _evals_constant.REFERENCE in row
4237+
if not _is_missing(row.get(_evals_constant.REFERENCE))
42314238
else None
42324239
),
42334240
candidate_responses=(

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

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,16 @@ def test_drop_empty_columns_none_empty(self):
102102
result_df = _evals_common._drop_empty_columns(df)
103103
assert list(result_df.columns) == ["col1", "col2"]
104104

105+
def test_drop_empty_columns_with_array_cells(self):
106+
df = pd.DataFrame(
107+
{
108+
"col1": [None, None],
109+
"col2": [pd.Series([1, 2]).to_numpy(), pd.Series([3, 4]).to_numpy()],
110+
}
111+
)
112+
result_df = _evals_common._drop_empty_columns(df)
113+
assert list(result_df.columns) == ["col2"]
114+
105115

106116
def _create_content_dump(text: str) -> dict[str, list[genai_types.Content]]:
107117
return {
@@ -4816,6 +4826,61 @@ def test_has_tool_call_with_agent_event(self):
48164826
assert _evals_metric_handlers._has_tool_call(events)
48174827

48184828

4829+
class TestExecuteInferenceConcurrently:
4830+
"""Unit tests for row handling in the concurrent inference helpers."""
4831+
4832+
def test_missing_agent_data_falls_back_to_prompt(self):
4833+
completed_trace = {
4834+
"turns": [
4835+
{
4836+
"events": [
4837+
{"author": "user", "content": {"parts": [{"text": "hi"}]}},
4838+
{"author": "agent", "content": {"parts": [{"text": "hey"}]}},
4839+
]
4840+
}
4841+
]
4842+
}
4843+
prompt_dataset = pd.DataFrame(
4844+
[{"prompt": "p0", "agent_data": completed_trace}, {"prompt": "p1"}]
4845+
)
4846+
inference_fn = mock.Mock(return_value={"response": "r1"})
4847+
4848+
responses = _evals_common._execute_inference_concurrently(
4849+
api_client=mock.Mock(),
4850+
prompt_dataset=prompt_dataset,
4851+
progress_desc="Agent Run",
4852+
inference_fn=inference_fn,
4853+
runtime=mock.Mock(),
4854+
)
4855+
4856+
assert [event["author"] for event in responses[0]] == ["user", "agent"]
4857+
assert responses[1] == {"response": "r1"}
4858+
assert inference_fn.call_args.kwargs["contents"] == "p1"
4859+
4860+
def test_responses_follow_row_positions(self):
4861+
prompt_dataset = pd.DataFrame({"prompt": ["p0", "p1"]}, index=[1, 0])
4862+
4863+
responses = _evals_common._execute_inference_concurrently(
4864+
api_client=None,
4865+
prompt_dataset=prompt_dataset,
4866+
progress_desc="Custom Inference",
4867+
model_or_fn=lambda contents: f"response to {contents}",
4868+
)
4869+
4870+
assert responses == ["response to p0", "response to p1"]
4871+
4872+
@mock.patch.object(_evals_common, "_call_litellm_completion")
4873+
def test_litellm_responses_follow_row_positions(self, mock_call_litellm_completion):
4874+
mock_call_litellm_completion.side_effect = lambda model, messages: {
4875+
"content": messages[0]["content"]
4876+
}
4877+
prompt_dataset = pd.DataFrame({"prompt": ["p0", "p1"]}, index=[1, 0])
4878+
4879+
responses = _evals_common._run_litellm_inference("gpt-4o", prompt_dataset)
4880+
4881+
assert responses == [{"content": "p0"}, {"content": "p1"}]
4882+
4883+
48194884
@pytest.mark.usefixtures("google_auth_mock")
48204885
class TestRunAgent:
48214886
"""Unit tests for the _run_agent function."""
@@ -10308,6 +10373,33 @@ def test_create_evaluation_set_with_intermediate_events(
1030810373
]
1030910374
assert candidate_response["events"] == expected_events
1031010375

10376+
@mock.patch.object(_evals_common, "evals")
10377+
@mock.patch.object(_evals_common, "_gcs_utils")
10378+
def test_create_evaluation_set_treats_nan_cells_as_missing(
10379+
self, mock_gcs_utils, mock_evals_module
10380+
):
10381+
eval_df = pd.DataFrame(
10382+
[{"prompt": "p0", "response": "r0"}, {"reference": "ref1"}]
10383+
)
10384+
10385+
_evals_common._create_evaluation_set_from_dataframe(
10386+
api_client=self.mock_api_client,
10387+
gcs_dest_prefix="gs://bucket/prefix",
10388+
eval_df=eval_df,
10389+
candidate_name="test-candidate",
10390+
)
10391+
10392+
upload = mock_gcs_utils.GcsUtils.return_value.upload_json_to_prefix
10393+
first, second = (call.kwargs["data"] for call in upload.call_args_list)
10394+
assert first == {
10395+
"prompt": {"text": "p0"},
10396+
"candidateResponses": [{"candidate": "test-candidate", "text": "r0"}],
10397+
}
10398+
assert second == {
10399+
"goldenResponse": {"text": "ref1"},
10400+
"candidateResponses": [{"candidate": "test-candidate"}],
10401+
}
10402+
1031110403
@mock.patch.object(_evals_common, "evals")
1031210404
@mock.patch.object(_evals_common, "_gcs_utils")
1031310405
def test_create_evaluation_set_with_user_scenario(
@@ -11799,6 +11891,14 @@ def test_build_interaction_id_dataset_requires_agent(self):
1179911891
[{"interaction_id": "abc123"}], None, "global"
1180011892
)
1180111893

11894+
def test_build_interaction_id_dataset_rejects_nan_interaction_id(self):
11895+
with pytest.raises(ValueError, match="Missing `interaction_id` value"):
11896+
_evals_common._build_interaction_id_dataset(
11897+
[{"interaction_id": "abc123"}, {"interaction_id": float("nan")}],
11898+
_TEST_GEMINI_AGENT,
11899+
"global",
11900+
)
11901+
1180211902
def test_build_interaction_id_dataset_rejects_non_gemini_agent(self):
1180311903
with pytest.raises(ValueError, match="Gemini Agents API resource name"):
1180411904
_evals_common._build_interaction_id_dataset(

‎vertexai/_genai/_evals_common.py‎

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,11 @@ def _temp_logger_level(logger_name: str, level: int) -> None: # type: ignore[mi
9393
logger_instance.setLevel(original_level)
9494

9595

96+
def _is_missing(value: Any) -> bool:
97+
"""Returns True for None or a scalar NaN, such as a missing DataFrame cell."""
98+
return bool(pd.api.types.is_scalar(value) and pd.isna(value))
99+
100+
96101
def _get_api_client_with_location(
97102
api_client: BaseApiClient, location: Optional[str]
98103
) -> BaseApiClient:
@@ -684,12 +689,12 @@ def _execute_inference_concurrently(
684689
max_workers = AGENT_MAX_WORKERS if agent_engine or agent else MAX_WORKERS
685690
with tqdm(total=len(prompt_dataset), desc=progress_desc) as pbar:
686691
with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
687-
for index, row in prompt_dataset.iterrows():
692+
for index, (_, row) in enumerate(prompt_dataset.iterrows()):
688693
try:
689694
if (
690695
has_agent_data
691696
and AGENT_DATA in row.index
692-
and row.get(AGENT_DATA) is not None
697+
and not _is_missing(row.get(AGENT_DATA))
693698
):
694699
agent_data_obj = row[AGENT_DATA]
695700
if isinstance(agent_data_obj, dict):
@@ -913,7 +918,7 @@ def _run_litellm_inference(
913918

914919
with tqdm(total=len(prompt_dataset), desc=f"LiteLLM Inference ({model})") as pbar:
915920
with concurrent.futures.ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor:
916-
for index, row in prompt_dataset.iterrows():
921+
for index, (_, row) in enumerate(prompt_dataset.iterrows()):
917922
messages = _convert_prompt_row_to_litellm_messages(row)
918923
future = executor.submit(
919924
_call_litellm_completion, model=model, messages=messages
@@ -2645,7 +2650,7 @@ def _drop_empty_columns(df: "pd.DataFrame") -> "pd.DataFrame":
26452650
def is_empty(x: Any) -> bool:
26462651
if isinstance(x, (list, dict)):
26472652
return not x
2648-
return pd.isna(x) # type: ignore[no-any-return]
2653+
return _is_missing(x)
26492654

26502655
cols_to_drop = [col for col in df.columns if df[col].apply(is_empty).all()]
26512656
return df.drop(columns=cols_to_drop)
@@ -2944,6 +2949,8 @@ def _create_evaluation_set_from_dataframe(
29442949
if _evals_constant.RESPONSE in row or agent_data_obj or intermediate_events:
29452950
# Resolve the oneof conflict: prioritize agent_data over flat text
29462951
response_text = row.get(_evals_constant.RESPONSE) or None
2952+
if _is_missing(response_text):
2953+
response_text = None
29472954

29482955
if agent_data_obj and response_text:
29492956
logger.info(
@@ -2998,15 +3005,15 @@ def _create_evaluation_set_from_dataframe(
29983005
prompt = types.EvaluationPrompt(
29993006
prompt_template_data=types.PromptTemplateData(values=values)
30003007
)
3001-
elif _evals_constant.PROMPT in row:
3008+
elif not _is_missing(row.get(_evals_constant.PROMPT)):
30023009
prompt = types.EvaluationPrompt(text=row[_evals_constant.PROMPT])
30033010

30043011
eval_item_requests.append(
30053012
types.EvaluationItemRequest(
30063013
prompt=prompt or None,
30073014
golden_response=(
30083015
types.CandidateResponse(text=row[_evals_constant.REFERENCE])
3009-
if _evals_constant.REFERENCE in row
3016+
if not _is_missing(row.get(_evals_constant.REFERENCE))
30103017
else None
30113018
),
30123019
candidate_responses=(

0 commit comments

Comments
 (0)