Update app.py

app.py

@@ -1,5 +1,6 @@
 from fastapi import FastAPI, HTTPException
 from fastapi.responses import JSONResponse
+from pydantic import BaseModel
 from typing import Optional
 import os
 import tempfile
@@ -17,7 +18,8 @@ app = FastAPI(title="Gemini Chat API", description="API para interagir com Googl
 
 # Inicializar chatbot globalmente
 
-chatbot = None
+# Inicializar chatbots globalmente
+chatbots = {}
 upscale_chatbot = None
 
 async def update_cookie_if_needed(cookie_path: str, secure_1psid: str, secure_1psidts: str, additional_cookies: dict):
@@ -35,7 +37,7 @@ async def init_chatbot(retry_count=0, max_retries=2):
     Tenta atualizar cookies automaticamente se falhar.
     """
 
-    global chatbot, upscale_chatbot
+    global chatbots, upscale_chatbot
     cookie_path = os.getenv("COOKIE_PATH", "cookies.json")
     
     if not os.path.exists(cookie_path):
@@ -49,15 +51,29 @@ async def init_chatbot(retry_count=0, max_retries=2):
         if retry_count == 0:
             secure_1psidts = await update_cookie_if_needed(cookie_path, secure_1psid, secure_1psidts, additional_cookies)
         
-        # Criar AsyncChatbot diretamente
-        chatbot = await AsyncChatbot.create(
+        # Criar Chatbot Flash (Padrão/Rápido)
+        chatbots['flash'] = await AsyncChatbot.create(
             secure_1psid=secure_1psid,
             secure_1psidts=secure_1psidts,
-            model=Model.G_2_5_PRO,
+            model=Model.G_3_0_FLASH,
             additional_cookies=additional_cookies,
             cookie_path=cookie_path
         )
-        print(f"Chatbot inicializado com sucesso usando {cookie_path}")
+        print(f"Chatbot Flash (3.0) inicializado com sucesso")
+
+        # Criar Chatbot Thinking (Raciocínio) - Timeout maior
+        chatbots['thinking'] = await AsyncChatbot.create(
+            secure_1psid=secure_1psid,
+            secure_1psidts=secure_1psidts,
+            model=Model.G_3_0_THINKING,
+            additional_cookies=additional_cookies,
+            cookie_path=cookie_path,
+            timeout=120 # Timeout maior para thinking
+        )
+        print(f"Chatbot Thinking (3.0) inicializado com sucesso")
+
+        # Fallback/Default
+        chatbots['default'] = chatbots['flash']
 
         # Criar instância de Upscale separada
         upscale_chatbot = await AsyncChatbot.create(
@@ -304,12 +320,52 @@ def download_file_with_retry(url: str, max_retries: int = 3, timeout: int = 300)
         detail=f"Falha ao baixar arquivo após {max_retries} tentativas"
     )
 
+class ChatRequest(BaseModel):
+    message: str
+    context: Optional[str] = None
+    model: Optional[str] = "flash" # 'flash' or 'thinking'
+
+@app.post("/chat")
+async def chat_endpoint(request: ChatRequest):
+    """
+    Endpoint para conversas de texto simples.
+    """
+    if not chatbots:
+        raise HTTPException(status_code=500, detail="Chatbot não inicializado")
+    
+    try:
+        requested_model = request.model.lower() if request.model else "flash"
+        if "thinking" in requested_model:
+            selected_chatbot = chatbots.get('thinking', chatbots['default'])
+        else:
+            selected_chatbot = chatbots.get('flash', chatbots['default'])
+
+        prompt = request.message
+        if request.context:
+            prompt = f"Contexto: {request.context}\n\nMensagem: {request.message}"
+            
+        print(f"💬 Chat request ({requested_model}): {prompt[:50]}...")
+        response_gemini = await selected_chatbot.ask(prompt)
+        
+        if response_gemini.get("error"):
+            raise HTTPException(
+                status_code=500,
+                detail=f"Erro no Gemini: {response_gemini.get('content', 'Erro desconhecido')}"
+            )
+            
+        return {"response": response_gemini.get("content", "")}
+        
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
 @app.get("/subtitle")
 async def generate_subtitle(
     file: str,
     context: Optional[str] = None,
     start: Optional[float] = None,
-    end: Optional[float] = None
+    end: Optional[float] = None,
+    model: Optional[str] = "flash" # 'flash' or 'thinking'
 ):
     """
     Endpoint para gerar legendas SRT a partir de um arquivo (imagem, vídeo ou áudio).
@@ -323,9 +379,15 @@ async def generate_subtitle(
     Retorna:
     - Arquivo SRT formatado (cortado se start e end forem fornecidos)
     """
-    if chatbot is None:
+    if not chatbots:
         raise HTTPException(status_code=500, detail="Chatbot não inicializado")
     
+    requested_model = model.lower() if model else "flash"
+    if "thinking" in requested_model:
+        selected_chatbot = chatbots.get('thinking', chatbots['default'])
+    else:
+        selected_chatbot = chatbots.get('flash', chatbots['default'])
+    
     if not file:
         raise HTTPException(status_code=400, detail="Parâmetro 'file' é obrigatório")
     
@@ -497,11 +559,11 @@ TRADUZA TUDO DE IMPORTANTE NO {media_desc_final}, que tenha dialogo... Nunca dei
         
         # Determinar qual parâmetro usar baseado no tipo de mídia
         if media_type == 'image':
-            response_gemini = await chatbot.ask(full_prompt, image=temp_file.name)
+            response_gemini = await selected_chatbot.ask(full_prompt, image=temp_file.name)
         elif media_type == 'video':
-            response_gemini = await chatbot.ask(full_prompt, video=temp_file.name)
+            response_gemini = await selected_chatbot.ask(full_prompt, video=temp_file.name)
         else:  # audio
-            response_gemini = await chatbot.ask(full_prompt, audio=temp_file.name)
+            response_gemini = await selected_chatbot.ask(full_prompt, audio=temp_file.name)
         
         if response_gemini.get("error"):
             raise HTTPException(