Update app.py

app.py

@@ -5,24 +5,26 @@ import os
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from groq import Groq
 
-# --- 1. Configuração do Peso Pesado (H200 Local) ---
-# Esse roda na placa da Nvidia de graça, mas gasta sua cota.
+# --- Configurações ---
 LOCAL_MODEL_ID = "Qwen/Qwen2.5-Coder-32B-Instruct"
 local_model = None
 local_tokenizer = None
 
-# --- 2. Configuração da Groq (Nuvem Rápida) ---
-# Esse roda na API deles, não gasta sua cota do ZeroGPU.
-# Pegue a chave em: https://console.groq.com/keys
-groq_client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
+# Tenta pegar a chave, se não tiver, avisa (evita crash)
+api_key = os.environ.get("GROQ_API_KEY")
+if api_key:
+    groq_client = Groq(api_key=api_key)
+else:
+    groq_client = None
+    print("⚠️ AVISO: GROQ_API_KEY não encontrada nos Secrets!")
 
-# --- Função ZeroGPU (Blindada com Cota Pequena) ---
-@spaces.GPU(duration=60) # 60s pra não queimar tudo de uma vez
+# --- Função H200 (ZeroGPU) ---
+@spaces.GPU(duration=60)
 def run_local_h200(messages):
     global local_model, local_tokenizer
     
     if local_model is None:
-        print(f"🐢 Cold Start: Carregando {LOCAL_MODEL_ID} na H200...")
+        print(f"🐢 Cold Start: Carregando {LOCAL_MODEL_ID}...")
         local_tokenizer = AutoTokenizer.from_pretrained(LOCAL_MODEL_ID)
         local_model = AutoModelForCausalLM.from_pretrained(
             LOCAL_MODEL_ID,
@@ -30,7 +32,6 @@ def run_local_h200(messages):
             device_map="cuda"
         )
     
-    # Prepara o prompt
     text = local_tokenizer.apply_chat_template(
         messages,
         tokenize=False,
@@ -38,20 +39,20 @@ def run_local_h200(messages):
     )
     inputs = local_tokenizer([text], return_tensors="pt").to(local_model.device)
 
-    # Gera a resposta
     outputs = local_model.generate(
         **inputs, 
-        max_new_tokens=2048, # Aumentei pra caber código grande
+        max_new_tokens=2048,
         temperature=0.6,
         do_sample=True
     )
     
-    response = local_tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
-    return response
+    return local_tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
 
-# --- Função Groq (Rápida e Grátis) ---
+# --- Função Groq ---
 def run_groq(messages, model_id):
-    print(f"⚡ Roteando para Groq: {model_id}")
+    if not groq_client:
+        return "❌ Erro: Configure a GROQ_API_KEY nos Settings do Space."
+        
     try:
         completion = groq_client.chat.completions.create(
             model=model_id,
@@ -64,40 +65,33 @@ def run_groq(messages, model_id):
         )
         return completion.choices[0].message.content
     except Exception as e:
-        return f"❌ Erro no Groq (Verifique a API Key): {str(e)}"
+        return f"❌ Erro no Groq: {str(e)}"
 
-# --- O CÉREBRO (Roteador) ---
+# --- Roteador ---
 def router(message, history, model_selector):
-    # Formata 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})
 
-    # Mapa de Roteamento
     if "H200" in model_selector:
         return run_local_h200(messages)
-    
-    elif "llama-3.1-8b-instant" in model_selector:
+    elif "Llama 3.1 70B" in model_selector:
+        return run_groq(messages, "llama-3.1-70b-versatile")
+    elif "Llama 3.1 8B" in model_selector:
         return run_groq(messages, "llama-3.1-8b-instant")
-        
-    elif "llama-3.3-70b-versatile" in model_selector:
-        return run_groq(messages, "llama-3.3-70b-versatile")
-        
-    elif "openai/gpt-oss-120b" in model_selector:
-        return run_groq(messages, "openai/gpt-oss-120b")
-        
-    elif "openai/gpt-oss-20b" in model_selector:
-        return run_groq(messages, "openai/gpt-oss-20b")
-    
+    elif "Gemma 2 9B" in model_selector:
+        return run_groq(messages, "gemma2-9b-it")
+    elif "Mixtral" in model_selector:
+        return run_groq(messages, "mixtral-8x7b-32768")
     else:
-        return "⚠️ Modelo não configurado no roteador."
+        return "⚠️ Modelo não encontrado."
 
-# --- Interface Gráfica ---
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🔀 APIDOST: O Hub do Gabriel")
-    gr.Markdown("Selecione o cérebro que você quer usar.")
+# --- Interface (SIMPLIFICADA) ---
+# Removi o theme=gr.themes.Soft() que estava quebrando
+with gr.Blocks() as demo:
+    gr.Markdown("# 🔀 APIDOST Router")
     
     with gr.Row():
         model_dropdown = gr.Dropdown(
@@ -109,11 +103,10 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                 "🔥 Local H200: Qwen 2.5 Coder 32B (Gasta Cota!)"
             ],
             value="☁️ Groq: Llama 3.1 70B (Inteligente & Rápido)", 
-            label="🤖 Escolha o Modelo",
+            label="Escolha o Modelo",
             interactive=True
         )
     
-    # A interface de chat conecta no roteador
     chat = gr.ChatInterface(
         fn=router,
         additional_inputs=[model_dropdown]