phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix(router.py): add support for async image generation endpoints

Browse source
This commit is contained in:
Krrish Dholakia 2023-12-21 14:38:44 +05:30
parent a4aa645cf6
commit be68796eba
6 changed files with 109 additions and 13 deletions
Download patch file Download diff file Expand all files Collapse all files

8
litellm/utils.py
View file

@ -551,6 +551,8 @@ class ImageResponse(OpenAIObject):
data: Optional[list] = None
usage: Optional[dict] = None
def __init__(self, created=None, data=None, response_ms=None):
if response_ms:
_response_ms = response_ms
@ -565,8 +567,10 @@ class ImageResponse(OpenAIObject):
created = created
else:
created = None
super().__init__(data=data, created=created)
self.usage = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}
def __contains__(self, key):
# Define custom behavior for the 'in' operator
@ -1668,6 +1672,8 @@ def client(original_function):
return result
elif "aembedding" in kwargs and kwargs["aembedding"] == True:
return result
elif "aimg_generation" in kwargs and kwargs["aimg_generation"] == True:
return result
### POST-CALL RULES ###
post_call_processing(original_response=result, model=model or None)