Update app.py

app.py

@@ -4,15 +4,23 @@ import torch
 import os
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from groq import Groq
+from mistralai import Mistral # SDK Oficial Atualizado
 
 # --- CONFIGURAÇÕES ---
+
+# 1. LOCAL (H200)
 LOCAL_MODEL_ID = "Qwen/Qwen2.5-Coder-32B-Instruct"
 local_model = None
 local_tokenizer = None
 
-# Cliente Groq
-api_key = os.environ.get("GROQ_API_KEY")
-groq_client = Groq(api_key=api_key) if api_key else None
+# 2. CLIENTE GROQ
+groq_api_key = os.environ.get("GROQ_API_KEY")
+groq_client = Groq(api_key=groq_api_key) if groq_api_key else None
+
+# 3. CLIENTE MISTRAL (NOVO!)
+mistral_api_key = os.environ.get("MISTRAL_API_KEY")
+mistral_client = Mistral(api_key=mistral_api_key) if mistral_api_key else None
+
 
 # --- FUNÇÃO 1: H200 (ZeroGPU - Cota Limitada) ---
 @spaces.GPU(duration=60)
@@ -39,70 +47,84 @@ def run_local_h200(messages):
     )
     return local_tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
 
-# --- FUNÇÃO 2: GROQ (Nuvem - Rápido e Grátis) ---
+# --- FUNÇÃO 2: GROQ (Nuvem) ---
 def run_groq(messages, model_id):
     if not groq_client:
-        return "❌ Erro: Chave GROQ_API_KEY não configurada nos Secrets."
-    
+        return "❌ Erro: Chave GROQ_API_KEY não configurada."
     try:
         completion = groq_client.chat.completions.create(
             model=model_id,
             messages=messages,
             temperature=0.7,
             max_tokens=4096,
-            top_p=1,
-            stream=False,
-            stop=None,
+            top_p=1, stream=False, stop=None,
         )
         return completion.choices[0].message.content
     except Exception as e:
-        return f"❌ Erro na Groq: {str(e)}"
+        return f"❌ Erro no Groq: {str(e)}"
+
+# --- FUNÇÃO 3: MISTRAL (Nuvem - NOVO!) ---
+def run_mistral(messages, model_id):
+    if not mistral_client:
+        return "❌ Erro: Chave MISTRAL_API_KEY não configurada."
+    
+    print(f"🇫🇷 Chamando Mistral: {model_id}")
+    try:
+        # A nova SDK da Mistral usa essa sintaxe
+        chat_response = mistral_client.chat.complete(
+            model=model_id,
+            messages=messages,
+        )
+        return chat_response.choices[0].message.content
+    except Exception as e:
+        return f"❌ Erro na Mistral AI: {str(e)}"
 
-# --- ROTEADOR ---
+# --- ROTEADOR CENTRAL ---
 def router(message, history, model_selector):
-    # Converte histórico
     messages = []
     for user_msg, bot_msg in history:
         if user_msg: messages.append({"role": "user", "content": user_msg})
         if bot_msg: messages.append({"role": "assistant", "content": bot_msg})
     messages.append({"role": "user", "content": message})
 
-    # Lógica baseada no SEU PRINT da Groq
+    # --- SELEÇÃO ---
     if "H200" in model_selector:
         return run_local_h200(messages)
     
-    # ATUALIZADO PARA LLAMA 3.3 (O novo do seu print)
-    elif "Llama 3.3 70B" in model_selector:
+    # Rota Groq
+    elif "Llama 3.3" in model_selector:
         return run_groq(messages, "llama-3.3-70b-versatile")
-        
     elif "Llama 3.1 8B" in model_selector:
         return run_groq(messages, "llama-3.1-8b-instant")
         
+    # Rota Mistral (NOVO)
+    elif "Mistral Large" in model_selector:
+        return run_mistral(messages, "mistral-large-latest") # O monstro do seu print
+    elif "Codestral" in model_selector:
+        return run_mistral(messages, "codestral-latest") # Especialista em código
+        
     else:
-        return "⚠️ Modelo não reconhecido. Verifique o dropdown."
+        return "⚠️ Modelo não reconhecido."
 
 # --- INTERFACE ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🔀 APIDOST Router V2")
+    gr.Markdown("# 🔀 APIDOST Router V3: Trinity")
     
     with gr.Row():
         model_dropdown = gr.Dropdown(
             choices=[
-                "☁️ Groq: Llama 3.3 70B (Novo & Versátil)", 
-                "☁️ Groq: Llama 3.1 8B (Flash - Instantâneo)",
-                "🔥 Local H200: Qwen 2.5 Coder 32B (Gasta Cota!)"
+                "☁️ Groq: Llama 3.3 70B (Versátil)", 
+                "☁️ Groq: Llama 3.1 8B (Flash)",
+                "🇫🇷 Mistral: Large 3 (SOTA - Inteligente)",
+                "🇫🇷 Mistral: Codestral (Especialista em Code)",
+                "🔥 Local H200: Qwen 2.5 Coder 32B (Cota ZeroGPU)"
             ],
-            value="☁️ Groq: Llama 3.3 70B (Novo & Versátil)", 
+            value="☁️ Groq: Llama 3.3 70B (Versátil)", 
             label="Escolha o Cérebro",
             interactive=True
         )
     
-    # Isso cria automaticamente o endpoint /chat
-    chat = gr.ChatInterface(
-        fn=router,
-        additional_inputs=[model_dropdown]
-    )
+    chat = gr.ChatInterface(fn=router, additional_inputs=[model_dropdown])
 
 if __name__ == "__main__":
-    # SEM parâmetros extras. Isso corrige o erro de inicialização.
     demo.launch()