From ff55a48c59f5c6a656df6aa167b09770c8878219 Mon Sep 17 00:00:00 2001 From: Jason Dai Date: Wed, 17 Sep 2025 13:22:04 -0700 Subject: [PATCH] fix: GenAI Client(evals) - Improve LiteLLM model completion response parsing PiperOrigin-RevId: 808266861 --- tests/unit/vertexai/genai/test_evals.py | 118 ++++++++++++++++++++++-- vertexai/_genai/_evals_common.py | 59 +++++++++++- 2 files changed, 168 insertions(+), 9 deletions(-) diff --git a/tests/unit/vertexai/genai/test_evals.py b/tests/unit/vertexai/genai/test_evals.py index 5510b11357..95fe411e63 100644 --- a/tests/unit/vertexai/genai/test_evals.py +++ b/tests/unit/vertexai/genai/test_evals.py @@ -1002,11 +1002,8 @@ def test_run_inference_with_litellm_string_prompt_format( assert call_kwargs["model"] == "gpt-4o" assert call_kwargs["messages"] == expected_messages assert "response" in result_dataset.eval_dataset_df.columns - response_content = json.loads(result_dataset.eval_dataset_df["response"][0]) - assert ( - response_content["choices"][0]["message"]["content"] - == "LiteLLM is a library..." - ) + response_content = result_dataset.eval_dataset_df["response"][0] + assert response_content == "LiteLLM is a library..." def test_run_inference_with_litellm_openai_request_format( self, @@ -1077,8 +1074,8 @@ def test_run_inference_with_litellm_openai_request_format( assert call_kwargs["model"] == "gpt-4o" assert call_kwargs["messages"] == expected_messages assert "response" in result_dataset.eval_dataset_df.columns - response_content = json.loads(result_dataset.eval_dataset_df["response"][0]) - assert response_content["choices"][0]["message"]["content"] == "Hello there" + response_content = result_dataset.eval_dataset_df["response"][0] + assert response_content == "Hello there" def test_run_inference_with_unsupported_model_string( self, @@ -1106,6 +1103,113 @@ def test_run_inference_with_litellm_import_error(self, mock_api_client_fixture): ): evals_module.run_inference(model="gpt-4o", src=prompt_df) + @mock.patch.object(_evals_common, "_run_litellm_inference") + @mock.patch.object(_evals_common, "_is_gemini_model") + @mock.patch.object(_evals_common, "_is_litellm_model") + @mock.patch.object(_evals_common, "_is_litellm_vertex_maas_model") + @mock.patch.object(_evals_utils, "EvalDatasetLoader") + def test_run_inference_with_litellm_parsing( + self, + mock_eval_dataset_loader, + mock_is_litellm_vertex_maas_model, + mock_is_litellm_model, + mock_is_gemini_model, + mock_run_litellm_inference, + ): + """Tests the parsing logic for LiteLLM responses within _run_inference_internal.""" + mock_is_gemini_model.return_value = False + mock_is_litellm_model.return_value = True + mock_is_litellm_vertex_maas_model.return_value = False + + mock_df = pd.DataFrame( + { + "prompt": [ + "prompt1", + "prompt2", + "prompt3", + "prompt4", + "prompt5", + "prompt6", + ] + } + ) + mock_eval_dataset_loader.return_value.load.return_value = mock_df.to_dict( + orient="records" + ) + + raw_responses = [ + { # Successful response + "choices": [{"message": {"content": "LiteLLM response 1"}}] + }, + {"error": "LiteLLM call failed: API key error"}, # Response with error key + {"id": "test-id-3"}, # Missing choices + {"choices": [{"index": 0}]}, # Missing message + {"choices": [{"message": {"role": "assistant"}}]}, # Missing content + "Invalid JSON string", # Non-dict response + ] + mock_run_litellm_inference.return_value = raw_responses + # fmt: off + with mock.patch("vertexai._genai._evals_common.litellm") as mock_litellm: + # fmt: on + mock_litellm.utils.get_valid_models.return_value = ["gpt-4o"] + inference_result = self.client.evals.run_inference( + model="gpt-4o", + src=mock_df, + ) + + expected_responses = [ + "LiteLLM response 1", + json.dumps({"error": "LiteLLM call failed: API key error"}), + json.dumps( + { + "error": "LiteLLM response missing 'choices'", + "details": {"id": "test-id-3"}, + } + ), + json.dumps( + { + "error": "LiteLLM response missing 'message' in first choice", + "details": {"choices": [{"index": 0}]}, + } + ), + json.dumps( + { + "error": "LiteLLM response missing 'content' in message", + "details": {"choices": [{"message": {"role": "assistant"}}]}, + } + ), + json.dumps( + { + "error": "Invalid LiteLLM response format", + "details": "Invalid JSON string", + } + ), + ] + + expected_df = pd.DataFrame( + { + "prompt": [ + "prompt1", + "prompt2", + "prompt3", + "prompt4", + "prompt5", + "prompt6", + ], + "response": expected_responses, + } + ) + + pd.testing.assert_frame_equal( + inference_result.eval_dataset_df.reset_index(drop=True), + expected_df.reset_index(drop=True), + check_dtype=False, + ) + mock_run_litellm_inference.assert_called_once() + _, call_kwargs = mock_run_litellm_inference.call_args + assert call_kwargs["model"] == "gpt-4o" + pd.testing.assert_frame_equal(call_kwargs["prompt_dataset"], mock_df) + class TestMetricPromptBuilder: """Unit tests for the MetricPromptBuilder class.""" diff --git a/vertexai/_genai/_evals_common.py b/vertexai/_genai/_evals_common.py index c9f4277552..7a998db1e2 100644 --- a/vertexai/_genai/_evals_common.py +++ b/vertexai/_genai/_evals_common.py @@ -496,8 +496,63 @@ def _run_inference_internal( raw_responses = _run_litellm_inference( model=processed_model_id, prompt_dataset=prompt_dataset ) - responses = [json.dumps(resp) for resp in raw_responses] - + processed_llm_responses = [] + for response_dict in raw_responses: + if not isinstance(response_dict, dict): + processed_llm_responses.append( + json.dumps( + { + "error": "Invalid LiteLLM response format", + "details": str(response_dict), + } + ) + ) + continue + + if "error" in response_dict: + processed_llm_responses.append(json.dumps(response_dict)) + continue + + if ( + "choices" in response_dict + and isinstance(response_dict["choices"], list) + and len(response_dict["choices"]) > 0 + ): + first_choice = response_dict["choices"][0] + if "message" in first_choice and isinstance( + first_choice["message"], dict + ): + message = first_choice["message"] + if "content" in message and isinstance(message["content"], str): + processed_llm_responses.append(message["content"]) + else: + processed_llm_responses.append( + json.dumps( + { + "error": "LiteLLM response missing 'content' in message", + "details": response_dict, + } + ) + ) + else: + processed_llm_responses.append( + json.dumps( + { + "error": "LiteLLM response missing 'message' in first choice", + "details": response_dict, + } + ) + ) + else: + processed_llm_responses.append( + json.dumps( + { + "error": "LiteLLM response missing 'choices'", + "details": response_dict, + } + ) + ) + responses = processed_llm_responses else: raise TypeError( f"Unsupported model type: {type(model)}. Expecting string (model"