scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 88

fix(vertex passthrough): extract input text from request body for embedContent token counting

Browse source
This commit is contained in:
Ishaan Jaffer 2026-04-20 21:07:24 -07:00
parent bbcf39f4dc
commit 0f69ffe3fe
No known key found for this signature in database
1 changed files with 25 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -136,6 +136,7 @@ class VertexPassthroughLoggingHandler:
logging_obj=logging_obj,
url_route=url_route,
kwargs=kwargs,
request_body=request_body,
)
elif "predict" in url_route:
return VertexPassthroughLoggingHandler._handle_predict_response(
@ -329,12 +330,28 @@ 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 (
@ -344,18 +361,23 @@ class VertexPassthroughLoggingHandler:
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 "batchEmbedContents" in url_route:
if is_batch:
litellm_embedding_response = process_batch_embed_response(
input="",
input=input_text,
model_response=model_response,
model=model,
_predictions=response_json,
)
else:
litellm_embedding_response = process_embed_content_response(
input="",
input=input_text,
model_response=model_response,
model=model,
response_json=response_json,