From 9d7942eb352d9387fdf5779931c4bdc23e3fc710 Mon Sep 17 00:00:00 2001 From: Tim Elfrink Date: Mon, 15 Sep 2025 21:43:07 +0200 Subject: [PATCH] Fix: Vertex AI Gemini labels field provider-aware filtering (#14563) * Add comprehensive tests for Vertex AI Gemini labels provider filtering - Test Google GenAI endpoints exclude labels even when explicitly provided - Test Vertex AI endpoints include labels when provided - Cover provider detection logic for different endpoint URLs - Verify metadata-to-labels conversion only happens for Vertex AI - Ensure edge cases are handled properly (null/empty api_base) * Fix Vertex AI Gemini labels field provider-aware filtering - Add _is_google_genai_endpoint() function to detect Google GenAI vs Vertex AI endpoints - Update _transform_request_body() to accept api_base parameter - Only include labels field for Vertex AI endpoints (not Google GenAI) - Pass api_base through sync/async transform functions - Maintain backward compatibility with existing usage - Fixes issue where Google GenAI requests failed with unsupported labels field * Refactor labels filtering to use custom_llm_provider instead of URL parsing Replace URL-based endpoint detection with custom_llm_provider parameter checking for cleaner, more reliable provider identification. Changes: - Remove _is_google_genai_endpoint() helper function - Update labels condition to use custom_llm_provider != "gemini" - Remove api_base parameter from _transform_request_body() - Simplify sync/async transform function signatures - Update tests to reflect new parameter structure - Remove obsolete test_provider_detection test This approach aligns with existing codebase patterns where custom_llm_provider="gemini" identifies Google AI Studio endpoints that don't support labels, while vertex_ai/vertex_ai_beta identify Vertex AI endpoints that do support labels. * Use LlmProviders.GEMINI constant instead of hardcoded string --- .../llms/vertex_ai/gemini/transformation.py | 6 +- .../test_vertex_ai_gemini_transformation.py | 84 ++++++++++++++++++- 2 files changed, 88 insertions(+), 2 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 327b269d1d..c59e3bb24e 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -28,6 +28,7 @@ from litellm.types.files import ( get_file_type_from_extension, is_gemini_1_5_accepted_file_type, ) +from litellm.types.utils import LlmProviders from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionAssistantMessage, @@ -492,7 +493,8 @@ def _transform_request_body( data["generationConfig"] = generation_config if cached_content is not None: data["cachedContent"] = cached_content - if labels is not None: + # Only add labels for Vertex AI endpoints (not Google GenAI/AI Studio) and only if non-empty + if labels and custom_llm_provider != LlmProviders.GEMINI: data["labels"] = labels except Exception as e: raise e @@ -647,3 +649,5 @@ def _transform_system_message( return SystemInstructions(parts=system_content_blocks), messages return None, messages + + diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py index d6d3325857..4da2976e1f 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py @@ -1,4 +1,7 @@ -from litellm.llms.vertex_ai.gemini.transformation import check_if_part_exists_in_parts +from litellm.llms.vertex_ai.gemini.transformation import ( + check_if_part_exists_in_parts, + _transform_request_body, +) def test_check_if_part_exists_in_parts(): @@ -73,3 +76,82 @@ def test_check_if_part_exists_in_parts_camel_case_snake_case(): } assert check_if_part_exists_in_parts(parts_mixed, part_mixed_casing) + + +# Tests for issue #14556: Labels field provider-aware filtering +def test_google_genai_excludes_labels(): + """Test that Google GenAI/AI Studio endpoints exclude labels when custom_llm_provider='gemini'""" + messages = [{"role": "user", "content": "test"}] + optional_params = {"labels": {"project": "test", "team": "ai"}} + litellm_params = {} + + result = _transform_request_body( + messages=messages, + model="gemini-2.5-pro", + optional_params=optional_params, + custom_llm_provider="gemini", + litellm_params=litellm_params, + cached_content=None, + ) + + # Google GenAI/AI Studio should NOT include labels + assert "labels" not in result + assert "contents" in result + + +def test_vertex_ai_includes_labels(): + """Test that Vertex AI endpoints include labels when custom_llm_provider='vertex_ai'""" + messages = [{"role": "user", "content": "test"}] + optional_params = {"labels": {"project": "test", "team": "ai"}} + litellm_params = {} + + result = _transform_request_body( + messages=messages, + model="gemini-2.5-pro", + optional_params=optional_params, + custom_llm_provider="vertex_ai", + litellm_params=litellm_params, + cached_content=None, + ) + + # Vertex AI SHOULD include labels + assert "labels" in result + assert result["labels"] == {"project": "test", "team": "ai"} + + + +def test_metadata_to_labels_vertex_only(): + """Test that metadata->labels conversion only happens for Vertex AI""" + messages = [{"role": "user", "content": "test"}] + optional_params = {} + litellm_params = { + "metadata": { + "requester_metadata": { + "user": "john_doe", + "project": "test-project" + } + } + } + + # Google GenAI/AI Studio should not include labels from metadata + result = _transform_request_body( + messages=messages, + model="gemini-2.5-pro", + optional_params=optional_params.copy(), + custom_llm_provider="gemini", + litellm_params=litellm_params.copy(), + cached_content=None, + ) + assert "labels" not in result + + # Vertex AI should include labels from metadata + result = _transform_request_body( + messages=messages, + model="gemini-2.5-pro", + optional_params=optional_params.copy(), + custom_llm_provider="vertex_ai", + litellm_params=litellm_params.copy(), + cached_content=None, + ) + assert "labels" in result + assert result["labels"] == {"user": "john_doe", "project": "test-project"}