Spaces:

CVision-assitant
/

asistente-virtual-CVision

Sleeping

asistente-virtual-CVision / modules /utils /chat_utils.py

Update modules/utils/chat_utils.py

e3025f2 verified 6 months ago

9 kB

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

	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