Update genesis/providers.py

genesis/providers.py

@@ -1,4 +1,3 @@
-# genesis/providers.py
 import os
 import requests
 from typing import List, Dict
@@ -22,14 +21,15 @@ def run_deepseek_summary(prompt: str) -> str:
         }
         r = requests.post(url, headers=headers, json=payload, timeout=60)
         r.raise_for_status()
-        return r.json()["choices"][0]["message"]["content"]
+        data = r.json()
+        return data["choices"][0]["message"]["content"]
     except Exception as e:
         print(f"[DeepSeek] Failed: {e}")
         return prompt
 
 # -------- Gemini Polish --------
 def run_gemini_polish(text: str) -> str:
-    """Polish the summary using Gemini for clarity."""
+    """Polish the summary using Gemini."""
     if not GEMINI_API_KEY:
         return text
     try:
@@ -37,47 +37,56 @@ def run_gemini_polish(text: str) -> str:
         payload = {"contents": [{"parts": [{"text": f"Polish and clarify this research report:\n\n{text}"}]}]}
         r = requests.post(url, json=payload, timeout=30)
         r.raise_for_status()
-        return r.json()["candidates"][0]["content"]["parts"][0]["text"]
+        data = r.json()
+        return data["candidates"][0]["content"]["parts"][0]["text"]
     except Exception as e:
         print(f"[Gemini] Failed: {e}")
         return text
 
-# -------- Image Generation (OpenAI → HF Fallback) --------
-def run_openai_image(query: str) -> str:
-    """Generate a research diagram using OpenAI image API, fallback to HF Stable Diffusion."""
-    # Try OpenAI first
-    if OPENAI_API_KEY:
-        try:
-            url = "https://api.openai.com/v1/images/generations"
-            headers = {"Authorization": f"Bearer {OPENAI_API_KEY}"}
-            payload = {"model": "gpt-image-1", "prompt": f"Scientific diagram about: {query}", "size": "1024x1024"}
-            r = requests.post(url, headers=headers, json=payload, timeout=60)
-            r.raise_for_status()
-            return r.json()["data"][0]["url"]
-        except Exception as e:
-            print(f"[OpenAI Image] Failed: {e}")
-
-    # Hugging Face fallback
-    if HF_TOKEN:
-        try:
-            print("[Image] Falling back to Hugging Face Stable Diffusion...")
-            url = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-xl-base-1.0"
-            headers = {"Authorization": f"Bearer {HF_TOKEN}"}
-            payload = {"inputs": f"highly detailed scientific diagram about: {query}"}
-            r = requests.post(url, headers=headers, json=payload, timeout=60)
-            r.raise_for_status()
-            img_path = f"/tmp/{query.replace(' ', '_')}.png"
-            with open(img_path, "wb") as f:
-                f.write(r.content)
-            return img_path
-        except Exception as e:
-            print(f"[HF Image] Failed: {e}")
+# -------- OpenAI Image --------
+def run_openai_image(prompt: str) -> str:
+    """Generate image using OpenAI API."""
+    if not OPENAI_API_KEY:
+        return None
+    try:
+        url = "https://api.openai.com/v1/images/generations"
+        headers = {"Authorization": f"Bearer {OPENAI_API_KEY}"}
+        payload = {
+            "model": "gpt-image-1",
+            "prompt": f"Highly detailed scientific diagram: {prompt}",
+            "size": "1024x1024"
+        }
+        r = requests.post(url, headers=headers, json=payload, timeout=60)
+        r.raise_for_status()
+        data = r.json()
+        return data["data"][0]["url"]
+    except Exception as e:
+        print(f"[OpenAI Image] Failed: {e}")
+        return None
 
-    return None
+# -------- Hugging Face Image (Fallback) --------
+def run_hf_image(prompt: str) -> str:
+    """Generate image using Hugging Face Stable Diffusion."""
+    if not HF_TOKEN:
+        return None
+    try:
+        url = "https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-2"
+        headers = {"Authorization": f"Bearer {HF_TOKEN}"}
+        payload = {"inputs": prompt}
+        r = requests.post(url, headers=headers, json=payload, timeout=60)
+        r.raise_for_status()
+        # Hugging Face returns raw image bytes, so we save to file
+        image_path = f"generated_{hash(prompt)}.png"
+        with open(image_path, "wb") as f:
+            f.write(r.content)
+        return image_path
+    except Exception as e:
+        print(f"[HF Image] Failed: {e}")
+        return None
 
 # -------- PubMed Fallback --------
 def pubmed_fallback_search(query: str, api_key: str, email: str) -> List[Dict]:
-    """Search PubMed for relevant articles if no citations found."""
+    """Search PubMed if no citations found."""
     results = []
     if not api_key or not email:
         return results