genesis/providers.py

import os
import requests
from typing import List, Dict

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
HF_TOKEN = os.getenv("HF_TOKEN")
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
NCBI_API_KEY = os.getenv("NCBI_API_KEY")
NCBI_EMAIL = os.getenv("NCBI_EMAIL")

# -------- DeepSeek Summary --------
def run_deepseek_summary(prompt: str) -> str:
    """Run a dense scientific summary using DeepSeek API."""
    try:
        url = "https://api.deepseek.com/v1/chat/completions"
        headers = {"Authorization": f"Bearer {HF_TOKEN}", "Content-Type": "application/json"}
        payload = {
            "model": "deepseek-science",
            "messages": [{"role": "user", "content": prompt}],
            "temperature": 0.2
        }
        r = requests.post(url, headers=headers, json=payload, timeout=60)
        r.raise_for_status()
        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."""
    if not GEMINI_API_KEY:
        return text
    try:
        url = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-pro:generateContent?key={GEMINI_API_KEY}"
        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()
        data = r.json()
        return data["candidates"][0]["content"]["parts"][0]["text"]
    except Exception as e:
        print(f"[Gemini] Failed: {e}")
        return text

# -------- 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

# -------- 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 if no citations found."""
    results = []
    if not api_key or not email:
        return results
    try:
        base_url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi"
        params = {
            "db": "pubmed",
            "term": query,
            "retmax": 3,
            "api_key": api_key,
            "email": email,
            "retmode": "json"
        }
        r = requests.get(base_url, params=params, timeout=10)
        r.raise_for_status()
        ids = r.json().get("esearchresult", {}).get("idlist", [])
        for pmid in ids:
            results.append({
                "type": "PMID",
                "id": pmid,
                "url": f"https://pubmed.ncbi.nlm.nih.gov/{pmid}/"
            })
    except Exception as e:
        print(f"[PubMed] Failed: {e}")
    return results