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 87

track /embedding in spendLogs

Browse source
This commit is contained in:
Ishaan Jaff 2024-09-02 17:05:53 -07:00
parent 42b95c5979
commit eb9e4f1787
2 changed files with 54 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

30
litellm/llms/vertex_ai_and_google_ai_studio/vertex_embeddings/embedding_handler.py
View file

@ -281,3 +281,33 @@ async def async_embedding(
)
setattr(model_response, "usage", usage)
return model_response
async def transform_vertex_response_to_openai(
response: dict, model: str, model_response: litellm.EmbeddingResponse
) -> litellm.EmbeddingResponse:
_predictions = response["predictions"]
embedding_response = []
input_tokens: int = 0
for idx, element in enumerate(_predictions):
embedding = element["embeddings"]
embedding_response.append(
{
"object": "embedding",
"index": idx,
"embedding": embedding["values"],
}
)
input_tokens += embedding["statistics"]["token_count"]
model_response.object = "list"
model_response.data = embedding_response
model_response.model = model
usage = Usage(
prompt_tokens=input_tokens, completion_tokens=0, total_tokens=input_tokens
)
setattr(model_response, "usage", usage)
return model_response

24
litellm/proxy/pass_through_endpoints/success_handler.py
View file

@ -97,6 +97,30 @@ class PassThroughEndpointLogging:
logging_obj.model = litellm_model_response.model
logging_obj.model_call_details["model"] = logging_obj.model
await logging_obj.async_success_handler(
result=litellm_model_response,
start_time=start_time,
end_time=end_time,
cache_hit=cache_hit,
)
elif "predict" in url_route:
from litellm.llms.vertex_ai_and_google_ai_studio.vertex_embeddings.embedding_handler import (
transform_vertex_response_to_openai,
)
model = self.extract_model_from_url(url_route)
_json_response = httpx_response.json()
litellm_model_response = await transform_vertex_response_to_openai(
response=_json_response,
model=model,
model_response=litellm.EmbeddingResponse(),
)
litellm_model_response.model = model
logging_obj.model = litellm_model_response.model
logging_obj.model_call_details["model"] = logging_obj.model
await logging_obj.async_success_handler(
result=litellm_model_response,
start_time=start_time,