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

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

@ -104,13 +104,30 @@ class PassThroughEndpointLogging:
cache_hit=cache_hit,
)
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 (
transform_vertex_response_to_openai,
)
vertex_image_generation_class = VertexImageGeneration()
model = self.extract_model_from_url(url_route)
_json_response = httpx_response.json()
litellm_model_response = litellm.ModelResponse()
if vertex_image_generation_class.is_image_generation_response(
_json_response
):
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,