scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-06 02:48:13 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(vertex passthrough): log :embedContent and :batchEmbedContents responses

Browse source
This commit is contained in:
Ishaan Jaffer 2026-04-20 20:06:30 -07:00
parent 850fe595ac
commit c498f5979d
No known key found for this signature in database
1 changed files with 61 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

61
litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py
View file

@ -130,6 +130,13 @@ 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,
)
elif "predict" in url_route:
return VertexPassthroughLoggingHandler._handle_predict_response(
httpx_response=httpx_response,
@ -322,6 +329,60 @@ class VertexPassthroughLoggingHandler:
"kwargs": kwargs,
}
@staticmethod
def _handle_embed_content_response(
httpx_response: httpx.Response,
logging_obj: LiteLLMLoggingObj,
url_route: str,
kwargs: dict,
) -> 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()
model_response = litellm.EmbeddingResponse()
if "batchEmbedContents" in url_route:
litellm_embedding_response = process_batch_embed_response(
input="",
model_response=model_response,
model=model,
_predictions=response_json,
)
else:
litellm_embedding_response = process_embed_content_response(
input="",
model_response=model_response,
model=model,
response_json=response_json,
)
litellm_embedding_response.model = model
logging_obj.model = model
logging_obj.model_call_details["model"] = model
logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai"
logging_obj.custom_llm_provider = "vertex_ai"
response_cost = litellm.completion_cost(
completion_response=litellm_embedding_response,
model=model,
custom_llm_provider="vertex_ai",
)
kwargs["response_cost"] = response_cost
kwargs["model"] = model
kwargs["custom_llm_provider"] = "vertex_ai"
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,