nvidia_nemotron.py

"""
╔══════════════════════════════════════════════════════════════════════════════╗
║              NVIDIA Nemotron Nano VL via OpenRouter - Doc2GL v2.0            ║
║                                                                               ║
║  Description : Génération de graphes Mermaid avec Nemotron Nano 12B VL      ║
║  Auteur      : YOUMBI CHATUE DANIELE                                         ║
╚══════════════════════════════════════════════════════════════════════════════╝
"""

import os
import base64
import logging
import requests

# Configuration du logging
logging.basicConfig(level=logging.INFO)


# ══════════════════════════════════════════════════════════════════════════════
#                          FONCTION PRINCIPALE
# ══════════════════════════════════════════════════════════════════════════════

def generate_mermaid_from_nvidia(base64_image):
    """
    Génère un diagramme Mermaid à partir d'une image en utilisant NVIDIA Nemotron via OpenRouter.

    Args:
        base64_image (str): Image encodée en Base64

    Returns:
        str: Code Mermaid représentant le graphe de connaissances

    Raises:
        EnvironmentError: Si la clé API n'est pas définie
        Exception: Si l'appel au modèle échoue

    Note:
        - Utilise NVIDIA Nemotron Nano 12B 2 VL via OpenRouter
        - Pas de téléchargement, pas de GPU requis
        - Clé API requise : NVIDIA_API_KEY (clé OpenRouter)
    """
    # ═════════════════════════════════════════════════════════════════════
    # ÉTAPE 1 : Récupération de la clé API OpenRouter
    # ═════════════════════════════════════════════════════════════════════
    # Essayer d'abord NVIDIA_API_KEY, sinon utiliser QWEN_API_KEY (même clé OpenRouter)
    api_key = os.environ.get("NVIDIA_API_KEY") or os.environ.get("QWEN_API_KEY")

    if not api_key:
        raise EnvironmentError(
            "❌ Aucune clé API OpenRouter trouvée.\n"
            "Définis NVIDIA_API_KEY ou QWEN_API_KEY (même clé OpenRouter)\n"
            "Format : sk-or-v1-xxxxxxxxxxxxxxxxxxxxx"
        )

    try:
        # ═════════════════════════════════════════════════════════════════
        # ÉTAPE 2 : Configuration de l'API OpenRouter
        # ═════════════════════════════════════════════════════════════════
        OPENROUTER_API_URL = "https://openrouter.ai/api/v1/chat/completions"

        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
            "HTTP-Referer": "https://doc2gl.app",
            "X-Title": "Doc2GL v2.0"
        }

        # ═════════════════════════════════════════════════════════════════
        # ÉTAPE 3 : Construction du message
        # ═════════════════════════════════════════════════════════════════
        payload = {
            # ✅ Modèle NVIDIA Nemotron via OpenRouter
            "model": "nvidia/nemotron-nano-12b-v2-vl:free",
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{base64_image}"
                            }
                        },
                        {
                            "type": "text",
                            "text": (
                                "Generate a Mermaid diagram of this mindmap image in the following "
                                "structured format: A main concept at the top connects to "
                                "several sub-concepts, which further branch into related "
                                "elements. Use a graph TD structure, and label each node "
                                "with meaningful names.\n\n"
                                "Example format:\n"
                                "graph TD\n"
                                "A[Main Topic] --> B[Subtopic 1]\n"
                                "A --> C[Subtopic 2]\n"
                                "B --> D[Detail 1]\n\n"
                                "Generate ONLY the Mermaid code, without any explanation or markdown formatting."
                            )
                        }
                    ]
                }
            ],
            "temperature": 0.7,
            "max_tokens": 1024,
            "top_p": 0.9
        }

        # ═════════════════════════════════════════════════════════════════
        # ÉTAPE 4 : Appel à l'API OpenRouter
        # ═════════════════════════════════════════════════════════════════
        logging.info("🤖 Génération en cours avec NVIDIA Nemotron (via OpenRouter)...")

        response = requests.post(
            OPENROUTER_API_URL,
            headers=headers,
            json=payload,
            timeout=60
        )

        # ═════════════════════════════════════════════════════════════════
        # ÉTAPE 5 : Validation de la réponse
        # ═════════════════════════════════════════════════════════════════
        if response.status_code != 200:
            error_msg = f"Erreur API OpenRouter : {response.status_code} - {response.text}"
            logging.error(f"❌ {error_msg}")
            raise Exception(error_msg)

        result = response.json()
        output_text = result["choices"][0]["message"]["content"]

        logging.info("✅ Génération terminée avec NVIDIA Nemotron (OpenRouter)")

        # ═════════════════════════════════════════════════════════════════
        # ÉTAPE 6 : Nettoyage du code Mermaid
        # ═════════════════════════════════════════════════════════════════
        clean_lines = []
        for line in output_text.split("\n"):
            clean_line = line.replace("```mermaid", "").replace("```", "").strip()
            clean_line = clean_line.split("'''")[0].strip()

            if clean_line and not clean_line.startswith("style"):
                clean_lines.append(clean_line)

        mermaid_code = "\n".join(clean_lines)

        if "graph" not in mermaid_code.lower() and "-->" not in mermaid_code:
            logging.warning("⚠️ Le code généré ne semble pas être un diagramme Mermaid valide")

        return mermaid_code

    except Exception as e:
        logging.error(f"❌ Erreur lors de l'appel à l'API NVIDIA (OpenRouter) : {e}")
        raise


# ══════════════════════════════════════════════════════════════════════════════
#                          FONCTION DE TEST
# ══════════════════════════════════════════════════════════════════════════════

if __name__ == "__main__":
    import sys

    test_image_path = "test_mindmap.png"

    if not os.path.exists(test_image_path):
        print(f"❌ Fichier {test_image_path} introuvable")
        sys.exit(1)

    with open(test_image_path, "rb") as f:
        base64_image = base64.b64encode(f.read()).decode('utf-8')

    print("🚀 Test de génération avec NVIDIA Nemotron (OpenRouter)...")

    try:
        mermaid = generate_mermaid_from_nvidia(base64_image)
        print("\n" + "="*80)
        print("📊 DIAGRAMME MERMAID GÉNÉRÉ :")
        print("="*80)
        print(mermaid)
        print("="*80)
        print("\n✅ SUCCESS !")
    except Exception as e:
        print(f"❌ Erreur : {e}")