From 7cf6a95b62a12250cb1af489cf748d9c573dfbf7 Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Fri, 24 Apr 2026 16:07:11 -0700 Subject: [PATCH] fix(vertex passthrough): log :embedContent and :batchEmbedContents responses (#26146) * fix(vertex passthrough): log :embedContent and :batchEmbedContents responses * test(vertex passthrough): add unit tests for :embedContent and :batchEmbedContents logging * fix(vertex passthrough): extract input text from request body for embedContent token counting * fix(vertex passthrough): add embedContent and batchEmbedContents to TRACKED_VERTEX_ROUTES * fix(vertex passthrough): detect Google AI Studio URLs in embedContent handler * test(vertex passthrough): add unit test for Google AI Studio URL embedContent provider detection * style: black format vertex_passthrough_logging_handler --- .../vertex_passthrough_logging_handler.py | 87 +++++++++ .../pass_through_endpoints/success_handler.py | 2 + .../test_llm_pass_through_endpoints.py | 170 ++++++++++++++++++ 3 files changed, 259 insertions(+) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index d709956df5c..86dd23e12ca 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -130,6 +130,14 @@ class VertexPassthroughLoggingHandler: "kwargs": kwargs, } + elif "embedContent" in url_route or "batchEmbedContents" in url_route: + return VertexPassthroughLoggingHandler._handle_embed_content_response( + httpx_response=httpx_response, + logging_obj=logging_obj, + url_route=url_route, + kwargs=kwargs, + request_body=request_body, + ) elif "predict" in url_route: return VertexPassthroughLoggingHandler._handle_predict_response( httpx_response=httpx_response, @@ -322,6 +330,85 @@ class VertexPassthroughLoggingHandler: "kwargs": kwargs, } + @staticmethod + def _extract_embed_content_input(request_body: Optional[dict], batch: bool) -> str: + """Extract raw input text from an :embedContent or :batchEmbedContents request body for token counting.""" + if not request_body: + return "" + if batch: + texts = [] + for req in request_body.get("requests", []): + for part in req.get("content", {}).get("parts", []): + texts.append(part.get("text", "")) + return " ".join(texts) + else: + parts = request_body.get("content", {}).get("parts", []) + return " ".join(part.get("text", "") for part in parts) + + @staticmethod + def _handle_embed_content_response( + httpx_response: httpx.Response, + logging_obj: LiteLLMLoggingObj, + url_route: str, + kwargs: dict, + request_body: Optional[dict] = None, + ) -> PassThroughEndpointLoggingTypedDict: + """Handle Vertex :embedContent and :batchEmbedContents endpoint responses.""" + from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import ( + process_embed_content_response, + process_response as process_batch_embed_response, + ) + + model = VertexPassthroughLoggingHandler.extract_model_from_url(url_route) + response_json = httpx_response.json() + is_batch = "batchEmbedContents" in url_route + + input_text = VertexPassthroughLoggingHandler._extract_embed_content_input( + request_body=request_body, batch=is_batch + ) + + model_response = litellm.EmbeddingResponse() + if is_batch: + litellm_embedding_response = process_batch_embed_response( + input=input_text, + model_response=model_response, + model=model, + _predictions=response_json, + ) + else: + litellm_embedding_response = process_embed_content_response( + input=input_text, + model_response=model_response, + model=model, + response_json=response_json, + ) + + custom_llm_provider = ( + VertexPassthroughLoggingHandler._get_custom_llm_provider_from_url(url_route) + ) + + litellm_embedding_response.model = model + logging_obj.model = model + logging_obj.model_call_details["model"] = model + logging_obj.model_call_details["custom_llm_provider"] = custom_llm_provider + logging_obj.custom_llm_provider = custom_llm_provider + + response_cost = litellm.completion_cost( + completion_response=litellm_embedding_response, + model=model, + custom_llm_provider=custom_llm_provider, + ) + + kwargs["response_cost"] = response_cost + kwargs["model"] = model + kwargs["custom_llm_provider"] = custom_llm_provider + logging_obj.model_call_details["response_cost"] = response_cost + + return { + "result": litellm_embedding_response, + "kwargs": kwargs, + } + @staticmethod def _handle_logging_vertex_collected_chunks( litellm_logging_obj: LiteLLMLoggingObj, diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index c14378ce616..0bc0183aa7c 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -46,6 +46,8 @@ class PassThroughEndpointLogging: "search", "batchPredictionJobs", "predictLongRunning", + "embedContent", + "batchEmbedContents", ] # Anthropic 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 06748e3f477..c9be00afed0 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 @@ -928,6 +928,176 @@ class TestVertexAIPassThroughHandler: assert "model" in result["kwargs"] assert result["kwargs"]["model"] == "textembedding-gecko@001" + def test_vertex_passthrough_handler_embed_content_response(self): + """ + Test that vertex_passthrough_handler correctly handles :embedContent responses + and invokes cost/logging callbacks (regression for silent drop bug). + """ + 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-embed" + mock_logging_obj.model_call_details = {} + + url_route = "/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-embedding-001: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.0002 + + 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 + ), "result must not be None — logging callbacks need a non-null response" + assert "kwargs" in result + assert result["kwargs"].get("response_cost") == 0.0002 + assert result["kwargs"].get("model") == "gemini-embedding-001" + assert result["kwargs"].get("custom_llm_provider") == "vertex_ai" + assert mock_logging_obj.model_call_details.get("response_cost") == 0.0002 + mock_completion_cost.assert_called_once() + + def test_vertex_passthrough_handler_batch_embed_contents_response(self): + """ + Test that vertex_passthrough_handler correctly handles :batchEmbedContents responses. + """ + 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, + ) + + batch_embed_response_data = { + "embeddings": [ + {"values": [0.1, 0.2, 0.3]}, + {"values": [0.4, 0.5, 0.6]}, + ] + } + + mock_httpx_response = Mock() + mock_httpx_response.json.return_value = batch_embed_response_data + mock_httpx_response.status_code = 200 + + mock_logging_obj = Mock(spec=LiteLLMLoggingObj) + mock_logging_obj.litellm_call_id = "test-call-id-batch" + mock_logging_obj.model_call_details = {} + + url_route = "/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-embedding-001:batchEmbedContents" + + 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.0003 + + 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 + ), "result must not be None for batchEmbedContents" + assert result["kwargs"].get("response_cost") == 0.0003 + assert result["kwargs"].get("model") == "gemini-embedding-001" + 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: """