modules/config/ai_tab.py

import gradio as gr
import os
import sys

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from dotenv import load_dotenv
load_dotenv(os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), ".env"))

PROVIDERS = {
    "groq": {
        "name": "Groq", "env_key": "GROQ_API_KEY",
        "base_url": "https://api.groq.com/openai/v1",
        "models": ["llama-3.3-70b-versatile", "llama-3.1-8b-instant", "deepseek-r1-distill-llama-70b", "mixtral-8x7b-32768", "gemma2-9b-it"],
    },
    "openrouter": {
        "name": "OpenRouter", "env_key": "OPENROUTER_API_KEY",
        "base_url": "https://openrouter.ai/api/v1",
        "models": ["meta-llama/llama-3.3-70b-instruct:free", "google/gemma-4-26b-a4b-it:free", "deepseek/deepseek-v4-flash:free", "anthropic/claude-sonnet-4.5", "openai/gpt-5.4", "mistralai/mistral-small-2603"],
    },
    "mistral": {
        "name": "Mistral AI", "env_key": "MISTRAL_API_KEY",
        "base_url": "https://api.mistral.ai/v1",
        "models": ["mistral-small-2506", "mistral-small-2603", "mistral-medium-2508", "mistral-medium-3-5", "mistral-large-2512", "magistral-medium-2509", "magistral-small-2509", "codestral-2508", "devstral-2512", "ministral-8b-2512", "ministral-14b-2512"],
    },
    "gemini": {
        "name": "Google Gemini", "env_key": "GEMINI_API_KEY",
        "base_url": "https://generativelanguage.googleapis.com/v1beta/openai",
        "models": ["gemini-2.5-flash", "gemini-2.5-pro", "gemini-2.0-flash", "gemini-3-flash-preview", "gemini-3-pro-preview"],
    },
    "deepseek": {
        "name": "DeepSeek", "env_key": "DEEPSEEK_API_KEY",
        "base_url": "https://api.deepseek.com/v1",
        "models": ["deepseek-chat", "deepseek-reasoner", "deepseek-v4-flash", "deepseek-v4-pro"],
    },
    "nebius": {
        "name": "Nebius", "env_key": "NEBIUS_API_KEY",
        "base_url": "https://api.tokenfactory.nebius.com/v1",
        "models": ["deepseek-ai/DeepSeek-V3.2", "deepseek-ai/DeepSeek-V4-Pro", "meta-llama/Llama-3.3-70B-Instruct", "Qwen/Qwen3-235B-A22B-Instruct-2507", "Qwen/Qwen3-32B"],
    },
    "azure": {
        "name": "Azure OpenAI", "env_key": "AZURE_API_KEY",
        "base_url": "https://letxinet.openai.azure.com/openai/deployments",
        "models": ["gpt-4o-mini", "gpt-4o", "o3-mini", "o4-mini", "gpt-4.1-mini"],
    },
    "huggingface": {
        "name": "HuggingFace", "env_key": "HF_TOKEN",
        "base_url": "https://api-inference.huggingface.co/v1",
        "models": ["deepseek-ai/DeepSeek-V3.2", "deepseek-ai/DeepSeek-R1", "meta-llama/Llama-3.3-70B-Instruct", "Qwen/Qwen3-235B-A22B-Instruct-2507"],
    },
}

DEFAULT_MODEL = "mistral-small-2506"

MODEL_ROLES = {
    "search": {"label": "🔍 Búsqueda", "description": "Modelo para planificar queries y estrategia de búsqueda", "default": DEFAULT_MODEL},
    "synthesis": {"label": "📝 Síntesis", "description": "Modelo para sintetizar y redactar el reporte final", "default": DEFAULT_MODEL},
    "translation": {"label": "🌐 Traducción", "description": "Modelo para traducción de contenido académico", "default": DEFAULT_MODEL},
}


def create_ai_tab():
    with gr.Tab("🤖 Modelos IA", id="ai"):
        gr.Markdown("## 🤖 Configuración de Modelos IA")
        gr.Markdown("Asigna un modelo a cada rol del pipeline de investigación. **Por defecto: mistral-small-2506**")

        with gr.Row():
            # ─── Selector de Proveedor ───
            with gr.Column(scale=1):
                gr.Markdown("### Proveedor")
                provider = gr.Dropdown(
                    choices=list(PROVIDERS.keys()), value="mistral",
                    label="Proveedor IA"
                )
                api_key = gr.Textbox(
                    label="API Key", type="password",
                    placeholder="sk-...",
                    value=os.getenv("MISTRAL_API_KEY", "")
                )
                key_status = gr.Markdown(
                    f"**Estado:** {'✅ Configurada' if os.getenv('MISTRAL_API_KEY') else '❌ No configurada'}"
                )

                def update_key_status(prov_name):
                    cfg = PROVIDERS.get(prov_name, {})
                    key = os.getenv(cfg.get("env_key", ""), "")
                    status = "✅ Configurada" if key else "❌ No configurada"
                    return gr.update(value=key), f"**Estado:** {status}"

                provider.change(fn=update_key_status, inputs=[provider], outputs=[api_key, key_status])

            # ─── Asignación de Modelos por Rol ───
            with gr.Column(scale=2):
                gr.Markdown("### Asignación de Modelos por Rol")
                gr.Markdown("Selecciona qué modelo usar para cada fase del pipeline:")

                model_dropdowns = {}
                role_info = {}

                for role_key, role_info_data in MODEL_ROLES.items():
                    with gr.Row():
                        model_dropdowns[role_key] = gr.Dropdown(
                            choices=PROVIDERS["mistral"]["models"],
                            value=DEFAULT_MODEL,
                            label=f"{role_info_data['label']} — {role_info_data['description']}",
                            allow_custom_value=True,
                        )
                        role_info[role_key] = gr.Markdown(f"**Rol:** `{role_key}` | **Modelo:** `{DEFAULT_MODEL}`")

                def update_models_on_provider_change(prov_name):
                    cfg = PROVIDERS.get(prov_name, PROVIDERS["mistral"])
                    models = cfg["models"]
                    updates = []
                    for role_key in MODEL_ROLES:
                        updates.append(gr.update(choices=models, value=models[0]))
                    return updates

                provider.change(
                    fn=update_models_on_provider_change,
                    inputs=[provider],
                    outputs=[model_dropdowns["search"], model_dropdowns["synthesis"], model_dropdowns["translation"]]
                )

                for role_key, dd in model_dropdowns.items():
                    def make_update(rk):
                        def updater(val):
                            return f"**Rol:** `{rk}` | **Modelo:** `{val}`"
                        return updater
                    dd.change(fn=make_update(role_key), inputs=[dd], outputs=[role_info[role_key]])

        # ─── Información de Roles ───
        with gr.Accordion("ℹ️ Guía de Roles del Pipeline", open=False):
            gr.Markdown("""
### Roles del Pipeline de Investigación

| Rol | Función | Cuándo se usa |
|-----|---------|---------------|
| **🔍 Búsqueda (search)** | Planificar queries, generar variaciones, detectar vacíos | Fase de búsqueda iterativa |
| **📝 Síntesis (synthesis)** | Redactar secciones del reporte, generar plan maestro | Fase de redacción |
| **🌐 Traducción (translation)** | Traducir contenido entre idiomas | Cuando se necesita traducción |

### Pipeline de Ejecución

```
Query → [search model] → Plan de Búsqueda → Búsqueda en fuentes
     → [synthesis model] → Plan Maestro → Redacción por secciones
     → [translation model] → Traducción (si aplica)
```

### Recomendaciones por Proveedor

| Proveedor | Mejor para | Velocidad |
|-----------|------------|-----------|
| **Mistral** | Síntesis académica de alta calidad | Media |
| **Groq** | Búsqueda rápida y planificación | Muy rápida |
| **OpenRouter** | Acceso a múltiples modelos (Claude, GPT, Llama) | Variable |
| **Gemini** | Ventanas de contexto grandes (1M tokens) | Rápida |
| **DeepSeek** | Razonamiento profundo (deepseek-reasoner) | Lenta |
| **Nebius** | Modelos open-source de alta calidad | Media |
            """)

        # ─── Exportar configuración ───
        def get_model_config():
            return {
                "provider": provider.value,
                "api_key": api_key.value,
                "search_model": model_dropdowns["search"].value,
                "synthesis_model": model_dropdowns["synthesis"].value,
                "translation_model": model_dropdowns["translation"].value,
            }

        config_output = gr.JSON(label="Configuración Actual", visible=False)

        return {
            "provider": provider,
            "api_key": api_key,
            "search_model": model_dropdowns["search"],
            "synthesis_model": model_dropdowns["synthesis"],
            "translation_model": model_dropdowns["translation"],
        }