better error handling for gemini

defog-ai · Dec 13, 2024 · ba0148e · ba0148e
1 parent 2aa9110
commit ba0148e
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 14 deletions.
diff --git a/defog_utils/utils_llm.py b/defog_utils/utils_llm.py
@@ -367,13 +367,16 @@ def chat_gemini(
         # use Pydantic classes for response_format
         generation_config.response_mime_type = 'application/json'
         generation_config.response_schema = response_format
-
-    response = client.models.generate_content(
-        model=model,
-        contents=message,
-        config=generation_config,
-    )
-    content = response.text
+
+    try:
+        response = client.models.generate_content(
+            model=model,
+            contents=message,
+            config=generation_config,
+        )
+        content = response.text
+    except Exception as e:
+        raise Exception(f"An error occurred: {e}")
 
     if response_format:
         # convert the content into Pydantic class
@@ -427,12 +430,15 @@ async def chat_gemini_async(
         generation_config.response_mime_type = 'application/json'
         generation_config.response_schema = response_format
 
-    response = await client.aio.models.generate_content(
-        model=model,
-        contents=message,
-        config=generation_config,
-    )
-    content = response.text
+    try:
+        response = await client.aio.models.generate_content(
+            model=model,
+            contents=message,
+            config=generation_config,
+        )
+        content = response.text
+    except Exception as e:
+        raise Exception(f"An error occurred: {e}")
 
     if response_format:
         # convert the content into Pydantic class

diff --git a/defog_utils/utils_multi_llm.py b/defog_utils/utils_multi_llm.py
@@ -96,7 +96,8 @@ async def chat_async(
             )
         except Exception as e:
             delay = base_delay * (2 ** attempt)  # Exponential backoff
-            print(f"Attempt {attempt + 1} failed. Retrying in {delay} seconds...")
+            print(f"Attempt {attempt + 1} failed. Retrying in {delay} seconds...", flush=True)
+            print(f"Error: {e}", flush=True)
             await asyncio.sleep(delay)
 
     # If we get here, all attempts failed