phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-27 11:43:54 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

track image gen in spend logs

Browse source
This commit is contained in:
Ishaan Jaff 2024-09-02 17:36:25 -07:00
parent aa13977136
commit 811aa34a36
1 changed files with 22 additions and 5 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -104,18 +104,35 @@ class PassThroughEndpointLogging:
cache_hit=cache_hit, cache_hit=cache_hit,
) )
elif "predict" in url_route: elif "predict" in url_route:
from litellm.llms.vertex_ai_and_google_ai_studio.image_generation.image_generation_handler import (
VertexImageGeneration,
)
from litellm.llms.vertex_ai_and_google_ai_studio.vertex_embeddings.embedding_handler import ( from litellm.llms.vertex_ai_and_google_ai_studio.vertex_embeddings.embedding_handler import (
transform_vertex_response_to_openai, transform_vertex_response_to_openai,
) )
vertex_image_generation_class = VertexImageGeneration()
model = self.extract_model_from_url(url_route) model = self.extract_model_from_url(url_route)
_json_response = httpx_response.json() _json_response = httpx_response.json()
litellm_model_response = await transform_vertex_response_to_openai( litellm_model_response = litellm.ModelResponse()
response=_json_response, if vertex_image_generation_class.is_image_generation_response(
model=model, _json_response
model_response=litellm.EmbeddingResponse(), ):
) litellm_model_response = (
vertex_image_generation_class.process_image_generation_response(
_json_response,
model_response=litellm.ImageResponse(),
model=model,
)
)
else:
litellm_model_response = await transform_vertex_response_to_openai(
response=_json_response,
model=model,
model_response=litellm.EmbeddingResponse(),
)
litellm_model_response.model = model litellm_model_response.model = model
logging_obj.model = litellm_model_response.model logging_obj.model = litellm_model_response.model