From 5d3674b63d23181a7c48b7ac0790c88969bb4b52 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Tue, 4 Jun 2024 08:29:12 -0700
Subject: [PATCH] fix(main.py): fix typing for image gen response

---
 litellm/llms/custom_httpx/http_handler.py | 4 ++--
 litellm/main.py                           | 6 ++++--
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py
index b186d4e1152..b91aaee2ae4 100644
--- a/litellm/llms/custom_httpx/http_handler.py
+++ b/litellm/llms/custom_httpx/http_handler.py
@@ -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:
diff --git a/litellm/main.py b/litellm/main.py
index f1d47427f42..f011adaade4 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -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.