Blame for modules/utils/chat_utils.py

import threading
import re
from concurrent.futures import ThreadPoolExecutor, TimeoutError

from modules.talker.talk import speak

from modules.langchain_assistant.langchain_productivy_assistant import (
    LangChainProductivityAssistant
)
from modules.langchain_assistant.langchain_browser_assistant import (
    LangChainBrowserAssistant
)
from modules.langchain_assistant.langchain_scientist_assistant import (
    LangChainScientisAssistant
)
from modules.langchain_assistant.langchain_chat_assistant import (
    LangChainChatAssistant
)
from modules.langchain_assistant.langchain_document_assistant import (
    LangChainDocumentAssistant
)

from modules.menus.role import create_notion_page


# ----------------------------------------
# ESTADO GLOBAL DE API KEYS
# ----------------------------------------
api_keys = {
    "gemini": None,
    "nvidia": None,
    "notion": None
}

# ----------------------------------------
# INSTANCIAS GLOBALES
# ----------------------------------------
productivity_role = None
browser_role      = None
scientist_role    = None
regular_chat_role = None


# ----------------------------------------
# RECONSTRUIR ROLES
# ----------------------------------------
def initialize_roles():
    global productivity_role, browser_role, scientist_role, regular_chat_role
    productivity_role = LangChainProductivityAssistant()
    browser_role      = LangChainBrowserAssistant()
    scientist_role    = LangChainScientisAssistant()
    regular_chat_role = LangChainChatAssistant()

initialize_roles()


# ----------------------------------------
# ACTUALIZAR API KEYS DESDE GRADIO
# ----------------------------------------
def update_api_keys(gemini_key, nvidia_key, notion_key):
    global api_keys
    api_keys = {
        "gemini": gemini_key.strip() if gemini_key else None,
        "nvidia": nvidia_key.strip() if nvidia_key else None,
        "notion": notion_key.strip() if notion_key else None
    }
    initialize_roles()
    return "✅ API Keys actualizadas correctamente."


# ----------------------------------------
# CONVERTIR HTML IMG → MARKDOWN
# ----------------------------------------
def convert_images_to_markdown(text):
    if not text:
        return ""
    img_pattern = r'<img[^>]*src="([^"]+)"[^>]*>'
    def replace(match):
        url = match.group(1)
        return f"\n\n![Imagen científica]({url})\n\n"
    text = re.sub(img_pattern, replace, text)
    text = re.sub(r"</?div[^>]*>", "", text)
    return text


# ----------------------------------------
# BUSCAR IMAGEN RELACIONADA (chat web)
# ----------------------------------------
def fetch_related_image(query):
    try:
        from ddgs import DDGS
        with DDGS() as ddgs:
            results = list(ddgs.images(query, max_results=5, safesearch="moderate"))
        for r in results:
            url = r.get("image") or r.get("url") or r.get("thumbnail")
            if url and url.startswith("http"):
                return url
    except Exception as e:
        print(f"⚠️ Error buscando imagen: {e}")
    return None


# ----------------------------------------
# GUARDAR EN NOTION ASÍNCRONO
# ----------------------------------------
def save_to_notion_async(content):
    try:
        threading.Thread(
            target=create_notion_page,
            args=(content,),
            daemon=True
        ).start()
    except Exception as e:
        print("⚠️ Error iniciando guardado en Notion:", e)


# ----------------------------------------
# TIMEOUT CONTROLADO
# Solo para roles que NO son BrowserAssistant
# ----------------------------------------
def run_with_timeout(role, user_input, timeout=120):
    def task():
        return role.chat(input=user_input)

    with ThreadPoolExecutor(max_workers=1) as executor:
        future = executor.submit(task)
        try:
            return future.result(timeout=timeout)
        except TimeoutError:
            print("⚠️ El modelo tardó demasiado.")
            raise Exception("Model timeout")


# ----------------------------------------
# EXTRAER CONTENIDO
# ----------------------------------------
def extract_content(role, assistant_response):
    if assistant_response is None:
        return "No se recibió respuesta."
    if isinstance(role, LangChainChatAssistant):
        return str(assistant_response)
    if isinstance(assistant_response, dict):
        return assistant_response.get("content", "Sin contenido.")
    return str(assistant_response)


# ----------------------------------------
# STREAMING RESPUESTA
# HF usa tuplas [user, bot] — NO dicts
# ----------------------------------------
def stream_response(chat_history, user_input, content):
    if chat_history is None:
        chat_history = []
    response = ""
    for letter in content:
        response += letter
        yield chat_history + [[user_input, response]], None


# ----------------------------------------
# FUNCIÓN BASE (todos los roles excepto browser)
# ----------------------------------------
def chat(chat_history, user_input, role):
    if chat_history is None:
        chat_history = []

    try:
        assistant_response = run_with_timeout(role, user_input)

    except Exception:
        yield chat_history + [[user_input, "PENSANDO 🧠🎯🖥️..."]], None

        try:
            assistant_response = role.chat(input=user_input)
        except Exception as e:
            yield chat_history + [[user_input, f"❌ Error del asistente: {str(e)}"]], None
            return

    content = extract_content(role, assistant_response)

    if isinstance(role, LangChainProductivityAssistant):
        if isinstance(assistant_response, dict):
            save_to_notion_async(assistant_response)

    if isinstance(role, LangChainScientisAssistant):
        content = convert_images_to_markdown(content)

    # Generar audio
    audio_path = None
    try:
        audio_path = speak(content)
        print("🔊 Audio generado:", audio_path)
    except Exception as e:
        print("⚠️ Error en voz:", e)

    for history in stream_response(chat_history, user_input, content):
        yield history[0], audio_path


# ----------------------------------------
# BROWSER CHAT — flujo especial para HF
# No usa run_with_timeout (NVIDIA tarda ~90s)
# Agrega imagen relacionada al final
# ----------------------------------------
def browser_chat(chat_history, user_input):
    if chat_history is None:
        chat_history = []

    # Mensaje de espera mientras busca
    yield chat_history + [[user_input, "🔍 Buscando en la web en tiempo real..."]], None

    try:
        # Llamada directa sin timeout externo
        assistant_response = browser_role.chat(input=user_input)
        content = extract_content(browser_role, assistant_response)
    except Exception as e:
        yield chat_history + [[user_input, f"❌ Error en búsqueda web: {str(e)}"]], None
        return

    # Agregar imagen relacionada
    try:
        image_url = fetch_related_image(user_input)
        if image_url:
            content += f"\n\n![Imagen relacionada]({image_url})"
    except Exception:
        pass

    # Audio
    audio_path = None
    try:
        audio_path = speak(content)
    except Exception as e:
        print("⚠️ Error en voz browser:", e)

    # Stream letra a letra
    for history in stream_response(chat_history, user_input, content):
        yield history[0], audio_path


# ----------------------------------------
# CHATS POR MÓDULO
# ----------------------------------------
def productivity_chat(chat_history, user_input):
    yield from chat(
        chat_history=chat_history,
        user_input=user_input,
        role=productivity_role
    )


def scientist_chat(chat_history, user_input):
    yield from chat(
        chat_history=chat_history,
        user_input=user_input,
        role=scientist_role
    )


def regular_chat(chat_history, user_input):
    yield from chat(
        chat_history=chat_history,
        user_input=user_input,
        role=regular_chat_role
    )


# ----------------------------------------
# CHAT PDF
# ----------------------------------------
def document_chat(chat_history, user_input, document_file):
    if chat_history is None:
        chat_history = []

    if document_file is None:
        yield chat_history + [[user_input, "⚠️ Debes subir un archivo PDF primero."]], None
        return

    try:
        file_path = document_file.name
        document_role = LangChainDocumentAssistant(document_name=file_path)
        assistant_response = document_role.chat(query=user_input)
    except Exception as e:
        yield chat_history + [[user_input, f"❌ Error procesando el PDF: {str(e)}"]], None
        return

    content = str(assistant_response)

    audio_path = None
    try:
        audio_path = speak(content)
    except Exception as e:
        print("⚠️ Error en voz PDF:", e)

    for history in stream_response(chat_history, user_input, content):
        yield history[0], audio_path