app.py

import gradio as gr
from huggingface_hub import InferenceClient
import chromadb
from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings
import os
from openai import OpenAI

# Configurar la API Key de OpenAI
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

# Inicializar el cliente de OpenAI
client = OpenAI(api_key=OPENAI_API_KEY)

# Inicializar el cliente de ChromaDB
chroma_client = chromadb.PersistentClient(path="chroma_db_crminbox/content/chroma_db")  # Ajusta la ruta según tu entorno

# Cargar la base de datos de Chroma como un vector store
vectorstore = Chroma(
    client=chroma_client,
    collection_name="docs",  # Nombre de la colección en Chroma
    embedding_function=OpenAIEmbeddings(model="text-embedding-3-small", openai_api_key=OPENAI_API_KEY)
)

# Crear un retriever
retriever = vectorstore.as_retriever()

# Función para obtener extractos relevantes
def obtener_extractos(pregunta):
    docs_relevantes = retriever.invoke(pregunta)
    return [(doc.page_content, doc.metadata.get("url", "URL no disponible")) for doc in docs_relevantes]

def extract_unique_citations_paragraph(response):
    """Extrae todas las URL y títulos únicos de un objeto Response, sin duplicados."""
    citations = []
    seen = set()  # Para evitar duplicados (usaremos la URL como clave única)

    for item in getattr(response, "output", []):
        # Solo buscamos en mensajes del asistente
        if getattr(item, "type", None) == "message":
            for block in getattr(item, "content", []) or []:
                if getattr(block, "type", None) == "output_text":
                    for ann in getattr(block, "annotations", []) or []:
                        if getattr(ann, "type", None) == "url_citation":
                            url = getattr(ann, "url", None)
                            title = getattr(ann, "title", None)

                            # Evita duplicados basados en URL
                            if url and url not in seen:
                                seen.add(url)
                                citations.append({
                                    "title": title,
                                    "url": url,
                                    "start": getattr(ann, "start_index", None),
                                    "end": getattr(ann, "end_index", None)
                                })

    return citations




def respond(message, history: list[tuple[str, str]], domain_table):
    """Genera una respuesta basada en el historial y documentos relevantes."""

    # --- Limpiar el DataFrame o lista de dominios ---
    if domain_table is None:
        allowed_domains = []
    elif hasattr(domain_table, "values"):  # DataFrame
        allowed_domains = [d for d in domain_table.iloc[:, 0].dropna().tolist() if d]
    else:  # lista de listas
        allowed_domains = [d[0] for d in domain_table if d and d[0]]

    system_message = '''Eres un agente especializado con alta competencia en investigación médica. Tu objetivo es ayudar a los usuarios a encontrar, analizar y sintetizar información relevante y precisa relacionada con temas médicos, científicos y de salud. Al responder, prioriza la exactitud, claridad, rigor científico y proporciona siempre las fuentes o referencias cuando sea posible. Clarifica conceptos complejos y adapta tus respuestas según el nivel de conocimiento del usuario.
                        Para tu investigación usarás siempre tu herrammienta de búsqueda en línea.
                        
                        # Detalles adicionales
                        
                        - Mantente actualizado en cuanto a literatura y evidencia médica reciente (hasta tu fecha de conocimiento).
                        - No proporciones diagnósticos médicos personalizados ni recomendaciones clínicas específicas, pero sí puedes guiar sobre dónde encontrar información fiable.
                        - Estructura las respuestas inicialmente explicando el razonamiento y la evidencia antes de presentar conclusiones o resúmenes.
                        
                        # Pasos sugeridos
                        
                        1. Analiza la consulta médica o científica recibida.
                        2. Identifica las fuentes de información y evidencia relevantes.
                        3. Sintetiza la información resaltando los hallazgos clave y referencias.
                        4. Presenta la respuesta, primero explicando el razonamiento y las evidencias encontradas, y luego dando la conclusión o resumen.
                        
                        # Formato de salida
                        
                        La respuesta debe estar bien estructurada, comenzando con la explicación y razonamiento seguido por las conclusiones o respuestas directas. Si es pertinente, incluye referencias en formato [autor, año] o enlaces.
                        
                        # Ejemplo
                        
                        Consulta del usuario: "¿Cuáles son los tratamientos más eficaces para la diabetes tipo 2 según la evidencia actual?"
                        
                        Respuesta:
                        
                        Primero, revisé guías internacionales y revisiones sistemáticas recientes para identificar las opciones terapéuticas recomendadas para diabetes tipo 2. Diversas fuentes coinciden en que la metformina sigue siendo el fármaco de primera elección debido a su eficacia, seguridad y bajo coste ([ADA, 2023]). Además, existen nuevas familias farmacológicas como los inhibidores de SGLT2 y los agonistas del GLP-1, que han demostrado beneficios cardiovasculares adicionales en estudios recientes ([Zinman et al., 2015]; [Marso et al., 2016]). La selección del tratamiento debe individualizarse según perfil del paciente y comorbilidades.
                        
                        En conclusión, el abordaje terapéutico debe basarse en guías actuales, iniciando generalmente con metformina y considerando otras opciones según características individuales y evidencia más reciente. 
                        
                        # Notas
                        
                        - Si la consulta es ambigua o muy general, clarifica la pregunta antes de responder.
                        - No inventes evidencia ni referencias; si la información es incierta, indícalo claramente.
                        - Si el usuario requiere información técnica (por ejemplo, sobre mecanismos moleculares), adapta el nivel de detalle según lo solicitado.
                        
                        Recuerda: Tu enfoque principal es brindar investigación y síntesis experta en el ámbito médico.'''

    # Obtener documentos relevantes desde ChromaDB
    contexto = ""#obtener_extractos(message)

    # Construir el mensaje del sistema con el contexto directamente incluido
    system_message_final = f"""{system_message}

    Información relevante extraída de los documentos, en caso de que estos documentos tenga la informacion que necesitas, no olvides tomar
    el historial de conversacion con el usuario:
    {contexto}
    """

    messages = [{"role": "system", "content": system_message}]

    # Agregar historial del chat
    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    # Agregar la nueva pregunta del usuario
    messages.append({"role": "user", "content": message})

    # --- Construir dinámicamente la tool ---
    web_search_tool = {
        "type": "web_search",
        "search_context_size": "high",
    }
    
    if allowed_domains:  # 👈 Solo agregar si hay dominios
        web_search_tool["filters"] = {"allowed_domains": allowed_domains}
    
    # --- Configurar include dinámico ---
    include_params = ["web_search_call.action.sources"] if allowed_domains else []
    
    # --- Llamada a la API con streaming ---
    with client.responses.stream(
        model="gpt-5-mini",
        input=messages,
        tools=[web_search_tool],
        include=include_params,  # 👈 también condicional
    ) as stream:
        response = ""
        for event in stream:
            if event.type == "response.output_text.delta":
                response += event.delta
                yield response
            elif event.type == "response.completed":
                response_stream = stream.get_final_response()
                citations = extract_unique_citations_paragraph(response_stream)
                response += "\n\n🔗 Fuentes citadas:"
                for c in citations:
                    response += f"\n• {c['title'] or 'Sin título'} → {c['url']}"
                    yield response
                break


def add_domain(domains, new_domain):
    # Si domains está vacío o None, inicializarlo como lista vacía
    if domains is None:
        domains = []
    # Si llega como DataFrame (Gradio lo hace internamente)
    elif hasattr(domains, "values"):
        domains = domains.values.tolist()
    # Si llega como lista con una fila vacía, limpiarla
    domains = [d for d in domains if d and d[0]]

    # Agregar el nuevo dominio si no está vacío
    if new_domain and new_domain.strip():
        domains.append([new_domain.strip()])

    return domains, ""  # limpia el textbox

with gr.Blocks() as demo:
    gr.Markdown("### 🧠 Chat Médico AI")
    gr.Markdown("Este es un chatbot especializdo en la búsqueda de información médica con fuentes verídicas. Puedes colocar que busque sólo en dominios específicos.")

    with gr.Row():
        new_domain = gr.Textbox(label="Búscar en sitios web específicos (opcional)")
        add_button = gr.Button("➕ Añadir")
    
    domain_table = gr.Dataframe(
        headers=["Dominios"],
        row_count=(1, "dynamic"),
        col_count=(1, "fixed"),
        label="Sitios web permitidos"
    )

    custom_chatbot = gr.Chatbot(height=500, label="Chat Médico AI")

    chat = gr.ChatInterface(
        fn=respond,
        chatbot=custom_chatbot,   
        additional_inputs=[domain_table],
    )

    add_button.click(add_domain, [domain_table, new_domain], [domain_table, new_domain])

if __name__ == "__main__":
    demo.launch()