Update app.py

app.py

@@ -4,13 +4,14 @@ import torch
 import os
 import time
 import base64
+from PIL import Image
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from groq import Groq
 from mistralai import Mistral
 import google.generativeai as genai
 from huggingface_hub import snapshot_download
 
-# --- 1. SEGURANÇA (ANTI-SPAM) ---
+# --- 1. SEGURANÇA (ANTI-SPAM - O escudo da Berta) ---
 MAX_REQUESTS_PER_MINUTE = 15
 BLOCK_TIME_SECONDS = 60
 ip_tracker = {}
@@ -35,8 +36,10 @@ LOCAL_MODEL_ID = "Qwen/Qwen2.5-Coder-32B-Instruct"
 local_model = None
 local_tokenizer = None
 
+# Clientes de API (A Berta verifica se as chaves existem para não dar erro feio)
 groq_client = Groq(api_key=os.environ.get("GROQ_API_KEY")) if os.environ.get("GROQ_API_KEY") else None
 mistral_client = Mistral(api_key=os.environ.get("MISTRAL_API_KEY")) if os.environ.get("MISTRAL_API_KEY") else None
+
 if os.environ.get("GEMINI_API_KEY"):
     genai.configure(api_key=os.environ.get("GEMINI_API_KEY"))
 
@@ -47,16 +50,17 @@ def encode_image(image_path):
             return base64.b64encode(image_file.read()).decode('utf-8')
     except: return None
 
-# --- 4. EXECUTORES ---
+# --- 4. EXECUTORES (Os operários da Berta) ---
 
 @spaces.GPU(duration=120)
 def run_local_h200(messages):
     global local_model, local_tokenizer
+    # Verificação crítica: Modelos locais de texto puro não leem imagens diretamente aqui
     for m in messages:
-        if isinstance(m['content'], list): return "⚠️ Modelo Local não suporta imagens. Use Gemini/Pixtral."
+        if isinstance(m['content'], list): return "⚠️ Berta avisa: Modelo Local não suporta imagens. Use Gemini ou Pixtral."
 
     if local_model is None:
-        print(f"🐢 Carregando {LOCAL_MODEL_ID}...")
+        print(f"🐢 Berta está carregando {LOCAL_MODEL_ID}... Tenha paciência, querido.")
         local_tokenizer = AutoTokenizer.from_pretrained(LOCAL_MODEL_ID)
         local_model = AutoModelForCausalLM.from_pretrained(
             LOCAL_MODEL_ID, torch_dtype=torch.bfloat16, device_map="cuda"
@@ -69,8 +73,8 @@ def run_local_h200(messages):
 
 def run_groq(messages, model_id):
     for m in messages:
-        if isinstance(m['content'], list): return "⚠️ Groq não suporta imagens. Use Gemini/Pixtral."
-    if not groq_client: return "❌ Erro: GROQ_API_KEY ausente."
+        if isinstance(m['content'], list): return "⚠️ Berta avisa: Groq ainda não suporta envio direto de imagens neste script."
+    if not groq_client: return "❌ Erro: Faltou a GROQ_API_KEY, meu anjo."
     clean_msgs = [{"role": m['role'], "content": m['content']} for m in messages]
     try:
         completion = groq_client.chat.completions.create(
@@ -80,7 +84,7 @@ def run_groq(messages, model_id):
     except Exception as e: return f"❌ Groq Error: {e}"
 
 def run_mistral(messages, model_id):
-    if not mistral_client: return "❌ Erro: MISTRAL_API_KEY ausente."
+    if not mistral_client: return "❌ Erro: Faltou a MISTRAL_API_KEY, príncipe."
     formatted_msgs = []
     for m in messages:
         content = m['content']
@@ -102,10 +106,13 @@ def run_mistral(messages, model_id):
     except Exception as e: return f"❌ Mistral Error: {e}"
 
 def run_gemini(messages, model_id):
-    if not os.environ.get("GEMINI_API_KEY"): return "❌ Erro: GEMINI_API_KEY ausente."
+    if not os.environ.get("GEMINI_API_KEY"): return "❌ Erro: Faltou a GEMINI_API_KEY."
     try:
+        # Instancia o modelo com o ID específico solicitado
         model = genai.GenerativeModel(model_id)
         chat_history = []
+        
+        # Constrói o histórico (exceto a última mensagem)
         for m in messages[:-1]:
             role = "user" if m['role'] == "user" else "model"
             parts = []
@@ -119,6 +126,7 @@ def run_gemini(messages, model_id):
                         if os.path.exists(path): parts.append(Image.open(path))
             if parts: chat_history.append({"role": role, "parts": parts})
 
+        # Prepara a última mensagem (prompt atual)
         last_parts = []
         lc = messages[-1]['content']
         if isinstance(lc, str): last_parts.append(lc)
@@ -132,14 +140,14 @@ def run_gemini(messages, model_id):
         chat = model.start_chat(history=chat_history)
         response = chat.send_message(last_parts)
         return response.text
-    except Exception as e: return f"❌ Gemini Error: {e}"
+    except Exception as e: return f"❌ Gemini Error ({model_id}): {e}"
 
-# --- 5. ROTEADOR CENTRAL ---
+# --- 5. ROTEADOR CENTRAL (O cérebro da operação) ---
 def router(message, history, model_selector, request: gr.Request):
     if not check_spam(request):
-        return "⛔ BLOQUEADO: Limite de mensagens excedido. Aguarde."
+        return "⛔ BLOQUEADO: Você está indo rápido demais, querido. Respire um pouco."
 
-    # Normaliza histórico
+    # Normaliza histórico para formato OpenAI
     messages = []
     if history:
         for turn in history:
@@ -165,15 +173,25 @@ def router(message, history, model_selector, request: gr.Request):
     else:
         messages.append({"role": "user", "content": str(message)})
 
-    # Seleção
-    if "Gemini" in model_selector:
+    # --- SELEÇÃO DE MODELOS ---
+    
+    # Rota Google / Gemini / LearnLM / Gemma
+    if any(k in model_selector for k in ["Gemini", "LearnLM", "Gemma"]):
+        # IDs Padrão
         tid = "gemini-1.5-flash"
+        
+        # Mapeamento Inteligente
         if "3.0" in model_selector: tid = "gemini-3.0-pro-preview"
         elif "2.5 Pro" in model_selector: tid = "gemini-2.5-pro"
+        elif "2.5 Flash Lite" in model_selector: tid = "gemini-2.5-flash-lite" # 🆕 Novo
         elif "2.5 Flash" in model_selector: tid = "gemini-2.5-flash"
-        elif "2.0" in model_selector: tid = "gemini-2.0-flash-exp"
+        elif "2.0" in model_selector and "LearnLM" not in model_selector: tid = "gemini-2.0-flash-exp"
+        elif "LearnLM" in model_selector: tid = "learnlm-2.0-flash-experimental" # 🆕 Novo
+        elif "Gemma 3" in model_selector: tid = "gemma-3-27b" # 🆕 Novo (Verificar se API aceita este ID exato)
+        
         return run_gemini(messages, tid)
     
+    # Rota Mistral
     elif "Mistral" in model_selector:
         tid = "mistral-large-latest"
         if "Pixtral" in model_selector: tid = "pixtral-large-latest"
@@ -182,30 +200,36 @@ def router(message, history, model_selector, request: gr.Request):
         elif "Codestral" in model_selector: tid = "codestral-2508"
         return run_mistral(messages, tid)
 
+    # Rota Groq
     elif "Groq" in model_selector:
         tid = "llama-3.3-70b-versatile"
         if "120B" in model_selector: tid = "openai/gpt-oss-120b"
         elif "20B" in model_selector: tid = "openai/gpt-oss-20b"
         return run_groq(messages, tid)
         
+    # Rota Local
     elif "H200" in model_selector:
         return run_local_h200(messages)
         
-    return "⚠️ Modelo não reconhecido."
+    return "⚠️ Modelo não reconhecido. Verifique o seletor."
 
 # --- 6. INTERFACE ---
-with gr.Blocks() as demo:
-    gr.Markdown("# 🔀 APIDOST v9 (API Bridge Fixed)")
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🔀 APIDOST v10 - Berta Edition")
+    gr.Markdown(f"### Olá Gabriel! Servindo seus modelos com carinho.")
     
     models_list = [
+        "✨ Google: LearnLM 1.5 Pro (Experimental) 📚", # LearnLM
+        "✨ Google: Gemma 3 27B (Preview) 💎", # Gemma 3
         "✨ Google: Gemini 3.0 Pro (Experimental)",
         "✨ Google: Gemini 2.5 Pro",
         "✨ Google: Gemini 2.5 Flash",
+        "✨ Google: Gemini 2.5 Flash Lite ⚡", # Flash Lite
         "✨ Google: Gemini 2.0 Flash",
-        "☁️ Groq: GPT OSS 120B (OpenAI) 🆕",
-        "☁️ Groq: GPT OSS 20B (OpenAI) 🆕",
+        "☁️ Groq: GPT OSS 120B (OpenAI)",
+        "☁️ Groq: GPT OSS 20B (OpenAI)",
         "☁️ Groq: Llama 3.3 70B",
-        "🇫🇷 Mistral: Magistral Medium 2509 🆕",
+        "🇫🇷 Mistral: Magistral Medium 2509",
         "🇫🇷 Mistral: Pixtral Large (Vision) 🖼️",
         "🇫🇷 Mistral: Large 2512 (Dez/25)",
         "🇫🇷 Mistral: Codestral 2508",
@@ -213,27 +237,26 @@ with gr.Blocks() as demo:
     ]
 
     with gr.Row():
-        model_dropdown = gr.Dropdown(choices=models_list, value=models_list[-1], label="Cérebro", interactive=True)
+        model_dropdown = gr.Dropdown(choices=models_list, value=models_list[0], label="Escolha o Cérebro", interactive=True)
     
-    # 1. Interface Visual (para debug no Space)
+    # 1. Interface Visual
     chat = gr.ChatInterface(
         fn=router, 
         additional_inputs=[model_dropdown],
-        multimodal=True 
+        multimodal=True,
+        description="Converse com a Berta e seus amigos AIs."
     )
 
-    # 2. PONTE DE API (A CORREÇÃO DO JAVASCRIPT)
-    # Use gr.JSON em vez de gr.State para 'history' para não exigir retorno.
-    # Isso cria a rota "/chat" corretamente.
+    # 2. PONTE DE API
     api_bridge = gr.Interface(
         fn=router,
         inputs=[
             gr.MultimodalTextbox(label="message"), 
-            gr.JSON(value=[], label="history"), # <--- O SEGREDO ESTÁ AQUI. JSON aceita lista e não trava.
-            gr.Dropdown(choices=models_list, label="model_selector", value=models_list[-1])
+            gr.JSON(value=[], label="history"),
+            gr.Dropdown(choices=models_list, label="model_selector", value=models_list[0])
         ],
         outputs=[gr.Textbox(label="response")],
-        api_name="chat" # <--- O JavaScript bate nessa porta aqui.
+        api_name="chat"
     )
 
 if __name__ == "__main__":