Update genesis/providers.py

genesis/providers.py

@@ -1,81 +1,110 @@
 import os
 import requests
+import logging
+import time
 from typing import Optional
 
+# Setup logging
+logging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s")
+
 # Load API keys from environment
 DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")
 GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
 OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
 HF_API_KEY = os.getenv("HF_API_KEY")
 
+# Retry decorator
+def retry_request(func):
+    def wrapper(*args, **kwargs):
+        retries = 3
+        delay = 2
+        for attempt in range(1, retries + 1):
+            try:
+                return func(*args, **kwargs)
+            except Exception as e:
+                logging.warning(f"{func.__name__} failed (attempt {attempt}): {e}")
+                if attempt < retries:
+                    time.sleep(delay)
+        return None
+    return wrapper
+
+
 # === DeepSeek Summarization ===
+@retry_request
 def run_deepseek_summary(prompt: str) -> str:
     if not DEEPSEEK_API_KEY:
+        logging.error("Missing DeepSeek API key.")
         return "[DeepSeek] API key missing."
-    try:
-        resp = requests.post(
-            "https://api.deepseek.com/v1/chat/completions",
-            headers={
-                "Authorization": f"Bearer {DEEPSEEK_API_KEY}",
-                "Content-Type": "application/json"
-            },
-            json={
-                "model": "deepseek-chat",
-                "messages": [{"role": "user", "content": prompt}],
-                "temperature": 0.3
-            }
-        )
-        resp.raise_for_status()
-        return resp.json()["choices"][0]["message"]["content"].strip()
-    except Exception as e:
-        return f"[DeepSeek] Failed: {str(e)}"
+    
+    resp = requests.post(
+        "https://api.deepseek.com/v1/chat/completions",
+        headers={
+            "Authorization": f"Bearer {DEEPSEEK_API_KEY}",
+            "Content-Type": "application/json"
+        },
+        json={
+            "model": "deepseek-chat",
+            "messages": [{"role": "user", "content": prompt}],
+            "temperature": 0.3
+        },
+        timeout=60
+    )
+    resp.raise_for_status()
+    return resp.json()["choices"][0]["message"]["content"].strip()
 
 
 # === Gemini Polishing ===
+@retry_request
 def run_gemini_polish(text: str) -> str:
     if not GEMINI_API_KEY:
+        logging.warning("No Gemini API key found, skipping polish.")
         return text
-    try:
-        resp = requests.post(
-            f"https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:generateContent?key={GEMINI_API_KEY}",
-            headers={"Content-Type": "application/json"},
-            json={
-                "contents": [{"parts": [{"text": f"Polish and refine the following scientific text:\n\n{text}"}]}]
-            }
-        )
-        resp.raise_for_status()
-        candidates = resp.json().get("candidates", [])
-        if candidates and "content" in candidates[0]:
-            return candidates[0]["content"]["parts"][0]["text"].strip()
-        return text
-    except Exception as e:
-        return f"[Gemini] Failed: {str(e)}"
+    
+    resp = requests.post(
+        f"https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:generateContent?key={GEMINI_API_KEY}",
+        headers={"Content-Type": "application/json"},
+        json={
+            "contents": [{"parts": [{"text": f"Polish and refine this scientific text:\n\n{text}"}]}]
+        },
+        timeout=60
+    )
+    resp.raise_for_status()
+    candidates = resp.json().get("candidates", [])
+    if candidates and "content" in candidates[0]:
+        return candidates[0]["content"]["parts"][0]["text"].strip()
+    return text
 
 
 # === Gemini Image Generation ===
+@retry_request
 def run_gemini_image(prompt: str) -> Optional[str]:
     if not GEMINI_API_KEY:
+        logging.warning("No Gemini API key found, skipping image generation.")
         return None
+    
     try:
         resp = requests.post(
             f"https://generativelanguage.googleapis.com/v1beta/models/gemini-pro-vision:generateImage?key={GEMINI_API_KEY}",
             headers={"Content-Type": "application/json"},
-            json={"prompt": prompt}
+            json={"prompt": prompt},
+            timeout=60
         )
         resp.raise_for_status()
         data = resp.json()
         if "imageUrl" in data:
             return data["imageUrl"]
-        return None
     except Exception as e:
-        print(f"[Gemini Image] Failed: {e}")
-        return None
+        logging.error(f"[Gemini Image] Failed: {e}")
+    return None
 
 
 # === OpenAI Image Generation ===
+@retry_request
 def run_openai_image(prompt: str) -> Optional[str]:
     if not OPENAI_API_KEY:
+        logging.warning("No OpenAI API key found, skipping OpenAI image.")
         return None
+    
     try:
         resp = requests.post(
             "https://api.openai.com/v1/images/generations",
@@ -83,31 +112,52 @@ def run_openai_image(prompt: str) -> Optional[str]:
                 "Authorization": f"Bearer {OPENAI_API_KEY}",
                 "Content-Type": "application/json"
             },
-            json={"model": "gpt-image-1", "prompt": prompt, "size": "1024x1024"}
+            json={"model": "gpt-image-1", "prompt": prompt, "size": "1024x1024"},
+            timeout=60
         )
         resp.raise_for_status()
         return resp.json()["data"][0]["url"]
     except Exception as e:
-        print(f"[OpenAI Image] Failed: {e}")
+        logging.error(f"[OpenAI Image] Failed: {e}")
         return None
 
 
-# === Hugging Face Fallback Image Generation ===
+# === Hugging Face Fallback Image ===
+@retry_request
 def run_hf_image(prompt: str) -> Optional[str]:
     if not HF_API_KEY:
+        logging.warning("No Hugging Face API key found, skipping HF image.")
         return None
+    
     try:
         resp = requests.post(
             "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-2",
             headers={"Authorization": f"Bearer {HF_API_KEY}"},
-            json={"inputs": prompt}
+            json={"inputs": prompt},
+            timeout=120
         )
         resp.raise_for_status()
-        # HF returns binary image; you need to save or serve it yourself
         output_file = "hf_image.png"
         with open(output_file, "wb") as f:
             f.write(resp.content)
+        logging.info(f"Hugging Face image saved to {output_file}")
         return output_file
     except Exception as e:
-        print(f"[HF Image] Failed: {e}")
+        logging.error(f"[HF Image] Failed: {e}")
         return None
+
+
+# === Master Image Generator ===
+def generate_image_with_fallback(prompt: str) -> Optional[str]:
+    """
+    Try Gemini → OpenAI → Hugging Face in order.
+    """
+    img_url = run_gemini_image(prompt)
+    if img_url:
+        return img_url
+
+    img_url = run_openai_image(prompt)
+    if img_url:
+        return img_url
+
+    return run_hf_image(prompt)