mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
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
This commit is contained in:
ishaan-berri
GitHub
parent
082a8faf46
commit
7cf6a95b62
3 changed files with 259 additions and 0 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -46,6 +46,8 @@ class PassThroughEndpointLogging:
|
|||
"search",
|
||||
"batchPredictionJobs",
|
||||
"predictLongRunning",
|
||||
"embedContent",
|
||||
"batchEmbedContents",
|
||||
]
|
||||
|
||||
# Anthropic
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue