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

fix(main.py): fix typing for image gen response

Browse source
This commit is contained in:
Krrish Dholakia 2024-06-04 08:29:12 -07:00
parent 7b474ec267
commit 5d3674b63d
2 changed files with 6 additions and 4 deletions
Download patch file Download diff file Expand all files Collapse all files

4
litellm/llms/custom_httpx/http_handler.py
View file

@ -84,9 +84,9 @@ class AsyncHTTPHandler:
response.raise_for_status() response.raise_for_status()
return response return response
except httpx.HTTPStatusError as e: except httpx.HTTPStatusError as e:
raise raise e
except Exception as e: except Exception as e:
raise raise e
def __del__(self) -> None: def __del__(self) -> None:
try: try:

6
litellm/main.py
View file

@ -3724,7 +3724,7 @@ async def amoderation(input: str, model: str, api_key: Optional[str] = None, **k
##### Image Generation ####################### ##### Image Generation #######################
@client @client
async def aimage_generation(*args, **kwargs): async def aimage_generation(*args, **kwargs) -> ImageResponse:
""" """
Asynchronously calls the `image_generation` function with the given arguments and keyword arguments. Asynchronously calls the `image_generation` function with the given arguments and keyword arguments.
@ -3757,6 +3757,8 @@ async def aimage_generation(*args, **kwargs):
if isinstance(init_response, dict) or isinstance( if isinstance(init_response, dict) or isinstance(
init_response, ImageResponse init_response, ImageResponse
): ## CACHING SCENARIO ): ## CACHING SCENARIO
if isinstance(init_response, dict):
init_response = ImageResponse(**init_response)
response = init_response response = init_response
elif asyncio.iscoroutine(init_response): elif asyncio.iscoroutine(init_response):
response = await init_response response = await init_response
@ -3792,7 +3794,7 @@ def image_generation(
litellm_logging_obj=None, litellm_logging_obj=None,
custom_llm_provider=None, custom_llm_provider=None,
**kwargs, **kwargs,
): ) -> ImageResponse:
""" """
Maps the https://api.openai.com/v1/images/generations endpoint. Maps the https://api.openai.com/v1/images/generations endpoint.