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

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

View file

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

View file

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