diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index c7839584024..d52aa43be05 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -1012,6 +1012,62 @@ class TestVertexAIPassThroughHandler: assert result["kwargs"].get("custom_llm_provider") == "vertex_ai" mock_completion_cost.assert_called_once() + def test_vertex_passthrough_handler_embed_content_google_ai_studio_url(self): + """ + Test that _handle_embed_content_response sets custom_llm_provider=gemini + when the URL is a generativelanguage.googleapis.com (Google AI Studio) endpoint. + """ + import datetime + from unittest.mock import Mock, patch + + from litellm.litellm_core_utils.litellm_logging import ( + Logging as LiteLLMLoggingObj, + ) + from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import ( + VertexPassthroughLoggingHandler, + ) + + embed_content_response_data = { + "embedding": { + "values": [0.1, 0.2, 0.3, 0.4, 0.5], + } + } + + mock_httpx_response = Mock() + mock_httpx_response.json.return_value = embed_content_response_data + mock_httpx_response.status_code = 200 + + mock_logging_obj = Mock(spec=LiteLLMLoggingObj) + mock_logging_obj.litellm_call_id = "test-call-id-gemini-studio" + mock_logging_obj.model_call_details = {} + + # Google AI Studio URL (not Vertex AI) + url_route = "https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-2-preview:embedContent" + + start_time = datetime.datetime.now() + end_time = datetime.datetime.now() + + with patch("litellm.completion_cost") as mock_completion_cost: + mock_completion_cost.return_value = 0.0001 + + result = VertexPassthroughLoggingHandler.vertex_passthrough_handler( + httpx_response=mock_httpx_response, + logging_obj=mock_logging_obj, + url_route=url_route, + result="test-result", + start_time=start_time, + end_time=end_time, + cache_hit=False, + ) + + assert result is not None + assert result["result"] is not None + assert result["kwargs"].get("custom_llm_provider") == "gemini", ( + "Google AI Studio embedContent URLs must set custom_llm_provider=gemini, not vertex_ai" + ) + assert result["kwargs"].get("model") == "gemini-embedding-2-preview" + mock_completion_cost.assert_called_once() + class TestVertexAIDiscoveryPassThroughHandler: """