Update app.py

app.py

@@ -2,12 +2,14 @@ import gradio as gr
 import spaces
 import torch
 import os
+import base64
+from PIL import Image
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from groq import Groq
 from mistralai import Mistral
 import google.generativeai as genai
 
-# --- CONFIGURAÇÕES ---
+# --- CONFIGURAÇÕES GLOBAIS ---
 
 # 1. LOCAL (H200)
 LOCAL_MODEL_ID = "Qwen/Qwen2.5-Coder-32B-Instruct"
@@ -25,9 +27,18 @@ gemini_key = os.environ.get("GEMINI_API_KEY")
 if gemini_key:
     genai.configure(api_key=gemini_key)
 
+# --- HELPER: Converte imagem para Base64 (Para Mistral) ---
+def encode_image(image_path):
+    with open(image_path, "rb") as image_file:
+        return base64.b64encode(image_file.read()).decode('utf-8')
+
 # --- FUNÇÃO 1: LOCAL H200 (ZeroGPU) ---
 @spaces.GPU(duration=60)
 def run_local_h200(messages):
+    # Verifica se tem imagem (não suportado)
+    if isinstance(messages[-1]['content'], list):
+         return "⚠️ Erro: O modelo local Qwen não suporta imagens."
+
     global local_model, local_tokenizer
     if local_model is None:
         print(f"🐢 Carregando {LOCAL_MODEL_ID}...")
@@ -43,6 +54,10 @@ def run_local_h200(messages):
 
 # --- FUNÇÃO 2: GROQ ---
 def run_groq(messages, model_id):
+    # Verifica se tem imagem (não suportado)
+    if isinstance(messages[-1]['content'], list):
+         return "⚠️ Erro: Os modelos da Groq não suportam imagens."
+         
     if not groq_client: return "❌ Erro: Configure a GROQ_API_KEY."
     try:
         completion = groq_client.chat.completions.create(
@@ -51,52 +66,133 @@ def run_groq(messages, model_id):
         return completion.choices[0].message.content
     except Exception as e: return f"❌ Groq Error: {e}"
 
-# --- FUNÇÃO 3: MISTRAL (IDs Novos da sua lista!) ---
+# --- FUNÇÃO 3: MISTRAL (Com suporte a Pixtral/Imagem) ---
 def run_mistral(messages, model_id):
     if not mistral_client: return "❌ Erro: Configure a MISTRAL_API_KEY."
+
+    # Processa mensagens para lidar com imagens para o Pixtral
+    processed_messages = []
+    for msg in messages:
+        if isinstance(msg['content'], list):
+            # É uma mensagem multimodal
+            content_list = []
+            for item in msg['content']:
+                if item['type'] == 'text':
+                    content_list.append({"type": "text", "text": item['text']})
+                elif item['type'] == 'image_url':
+                    # Gradio passa o caminho do arquivo, convertemos para base64 data URL
+                    image_path = item['image_url']['url']
+                    base64_image = encode_image(image_path)
+                    # Assumindo JPEG por padrão, o ideal seria detectar
+                    content_list.append({
+                        "type": "image_url",
+                        "image_url": f"data:image/jpeg;base64,{base64_image}" 
+                    })
+            processed_messages.append({"role": msg['role'], "content": content_list})
+        else:
+            # Mensagem só de texto
+            processed_messages.append(msg)
+
     try:
-        print(f"🇫🇷 Mistral Target: {model_id}")
-        res = mistral_client.chat.complete(model=model_id, messages=messages)
+        print(f"🇫🇷 Mistral Target: {model_id} (Multimodal: {'pixtral' in model_id})")
+        res = mistral_client.chat.complete(model=model_id, messages=processed_messages)
         return res.choices[0].message.content
-    except Exception as e: return f"❌ Mistral Error: {e}"
+    except Exception as e:
+        if "image" in str(e).lower() and "pixtral" not in model_id:
+             return f"⚠️ Erro: Você enviou uma imagem, mas o modelo '{model_id}' não suporta. Use o Pixtral."
+        return f"❌ Mistral Error: {e}"
 
-# --- FUNÇÃO 4: GEMINI (Modo "Hacker" - Aceita o que vc mandar) ---
+# --- FUNÇÃO 4: GEMINI (Com suporte a Imagem) ---
 def run_gemini(messages, model_id):
     if not gemini_key: return "❌ Erro: Configure a GEMINI_API_KEY."
     try:
         model = genai.GenerativeModel(model_id)
-        full_prompt = "\n".join([f"{m['role'].title()}: {m['content']}" for m in messages])
-        response = model.generate_content(full_prompt)
+        
+        gemini_history = []
+        for m in messages:
+            role = "user" if m['role'] == 'user' else "model"
+            
+            if isinstance(m['content'], list):
+                # Mensagem multimodal
+                parts = []
+                for item in m['content']:
+                    if item['type'] == 'text':
+                        parts.append(item['text'])
+                    elif item['type'] == 'image_url':
+                        # Gemini aceita PIL Image ou caminho
+                        image_path = item['image_url']['url']
+                        img = Image.open(image_path)
+                        parts.append(img)
+                gemini_history.append({"role": role, "parts": parts})
+            else:
+                # Mensagem só de texto
+                gemini_history.append({"role": role, "parts": [m['content']]})
+
+        # A última mensagem é o prompt atual
+        prompt_parts = gemini_history.pop()
+        
+        # Inicia o chat com o histórico
+        chat = model.start_chat(history=gemini_history)
+        response = chat.send_message(prompt_parts['parts'])
+        
         return response.text
     except Exception as e: return f"❌ Gemini Error ({model_id}): {e}"
 
-# --- ROTEADOR ---
+# --- ROTEADOR CENTRAL (Adaptado para Multimodal) ---
 def router(message, history, model_selector):
+    # 'message' agora é um dicionário: {"text": "...", "files": ["path/to/image.jpg"]}
+    
+    # 1. Formata o 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})
+        # Mensagem do usuário (pode ter imagem)
+        if isinstance(user_msg, dict):
+            content = [{"type": "text", "text": user_msg.get("text", "")}]
+            for file_path in user_msg.get("files", []):
+                content.append({"type": "image_url", "image_url": {"url": file_path}})
+            messages.append({"role": "user", "content": content})
+        elif user_msg:
+            messages.append({"role": "user", "content": user_msg})
+            
+        # Mensagem do bot (sempre texto por enquanto)
+        if bot_msg:
+            messages.append({"role": "assistant", "content": bot_msg})
 
-    # --- MAPEAMENTO DE MODELOS ---
+    # 2. Formata a Mensagem Atual
+    current_content = [{"type": "text", "text": message.get("text", "")}]
+    for file_path in message.get("files", []):
+        current_content.append({"type": "image_url", "image_url": {"url": file_path}})
     
-    # Rota Google (Seus modelos "Future")
+    # Se for só texto, simplifica para string (para compatibilidade com funções antigas)
+    if len(current_content) == 1 and current_content[0]['type'] == 'text' and not message.get("files"):
+        messages.append({"role": "user", "content": current_content[0]['text']})
+    else:
+        # Se tiver imagem ou for misto, manda a lista estruturada
+        messages.append({"role": "user", "content": current_content})
+
+    # --- MAPEAMENTO DE MODELOS (SEUS IDs COMPROVADOS!) ---
+    
+    # Rota Google
     if "Gemini 3" in model_selector:
-        return run_gemini(messages, "gemini-3.0-pro-exp") # Tentativa de ID
+        return run_gemini(messages, "gemini-3.0-pro-preview") #
+    elif "Gemini 2.5 Pro" in model_selector:
+        return run_gemini(messages, "gemini-2.5-pro") #
     elif "Gemini 2.5 Flash" in model_selector:
-        return run_gemini(messages, "gemini-2.0-flash-exp") # O Experimental atual
+        return run_gemini(messages, "gemini-2.5-flash") #
+    elif "Gemini 2.0 Flash" in model_selector:
+        return run_gemini(messages, "gemini-2.0-flash") #
         
     # Rota Groq
     elif "Llama 3.3" in model_selector:
         return run_groq(messages, "llama-3.3-70b-versatile")
         
-    # Rota Mistral (Sua Lista VIP)
+    # Rota Mistral
+    elif "Pixtral Large" in model_selector:
+        return run_mistral(messages, "pixtral-large-latest") # VISÃO!
     elif "Mistral Large 2512" in model_selector:
         return run_mistral(messages, "mistral-large-2512")
-    elif "Pixtral Large" in model_selector:
-        return run_mistral(messages, "pixtral-large-latest") # Multimodal!
     elif "Magistral Medium" in model_selector:
-        return run_mistral(messages, "magistral-medium-latest") # Exclusivo
+        return run_mistral(messages, "magistral-medium-latest")
     elif "Codestral 2508" in model_selector:
         return run_mistral(messages, "codestral-2508")
         
@@ -108,26 +204,34 @@ def router(message, history, model_selector):
 
 # --- UI ---
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🔀 APIDOST Router V5: Unleashed")
+    gr.Markdown("# 🔀 APIDOST Router V6: Vision Unleashed")
     
     with gr.Row():
         model_dropdown = gr.Dropdown(
             choices=[
-                "✨ Google: Gemini 3 (Experimental)",
+                "✨ Google: Gemini 3.0 Pro (Preview)",
+                "✨ Google: Gemini 2.5 Pro",
                 "✨ Google: Gemini 2.5 Flash",
+                "✨ Google: Gemini 2.0 Flash",
                 "☁️ Groq: Llama 3.3 70B",
+                "🇫🇷 Mistral: Pixtral Large (Vision) 🖼️",
                 "🇫🇷 Mistral: Large 2512 (Dez/25)",
-                "🇫🇷 Mistral: Pixtral Large (Vision)",
                 "🇫🇷 Mistral: Magistral Medium (VIP)",
                 "🇫🇷 Mistral: Codestral 2508 (Code)",
                 "🔥 Local H200: Qwen 2.5 Coder"
             ],
-            value="🇫🇷 Mistral: Large 2512 (Dez/25)", 
+            value="🇫🇷 Mistral: Pixtral Large (Vision) 🖼️", 
             label="Escolha o Cérebro",
             interactive=True
         )
     
-    chat = gr.ChatInterface(fn=router, additional_inputs=[model_dropdown])
+    # ATENÇÃO: multimodal=True HABILITA O UPLOAD DE IMAGENS
+    chat = gr.ChatInterface(
+        fn=router, 
+        additional_inputs=[model_dropdown],
+        multimodal=True, # O SEGREDO ESTÁ AQUI!
+        textbox=gr.MultimodalTextbox(placeholder="Digite e/ou cole uma imagem...", file_types=["image"], interactive=True)
+    )
 
 if __name__ == "__main__":
     demo.launch()