CVision-assitant's picture
Update modules/utils/chat_utils.py
e3025f2 verified
Raw History Blame Contribute Delete
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