app.py

import os
import inspect
import chainlit as cl
import PyPDF2

from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.chains import RetrievalQAWithSourcesChain, LLMChain
from langchain.chat_models import ChatOpenAI
from langchain.prompts.chat import (
    ChatPromptTemplate,
    SystemMessagePromptTemplate,
    HumanMessagePromptTemplate,
)

# Clase personalizada que cumple con la nueva interfaz de EmbeddingFunction de Chroma
class CustomOpenAIEmbeddings(OpenAIEmbeddings):
    def __call__(self, input):
        # Llama al método embed_documents para generar las embeddings a partir de una lista de textos
        return self.embed_documents(input)

# Forzamos la firma de __call__ para que tenga exactamente ("self", "input")
CustomOpenAIEmbeddings.__call__.__signature__ = inspect.Signature(
    parameters=[
        inspect.Parameter("self", inspect.Parameter.POSITIONAL_OR_KEYWORD),
        inspect.Parameter("input", inspect.Parameter.POSITIONAL_OR_KEYWORD)
    ]
)

# --- CONFIGURACIÓN ---
# Obtenemos la API key de OpenAI desde las variables de entorno
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
if not OPENAI_API_KEY:
    raise ValueError(
        "No se encontró la variable de entorno 'OPENAI_API_KEY'. Defínela en tu entorno o en los secrets."
    )

# Configuración del text splitter (puedes ajustar chunk_size y chunk_overlap según tus necesidades)
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)

# --- PROMPTS Y PLANTILLAS ---

# Plantilla del sistema para consultas basadas en PDF + Conocimiento General
system_template = """\
Eres un asistente en español basado en ChatGPT-4 con grandes capacidades de razonamiento y análisis.
Tienes acceso a los siguientes documentos, y también cuentas con conocimientos generales para responder
toda clase de preguntas, tanto del contexto provisto como de tu conocimiento general.

- Si la pregunta está claramente respondida por el contenido de los textos, proporciona la información relevante y cita tus fuentes.
- Si no está respondida por los textos, utiliza tu conocimiento general y responde de forma analítica, extensa y detallada.
- Siempre que utilices información proveniente de los PDFs, al final de tu respuesta indica las fuentes de la forma:
  FUENTES: nombre_del_pdf

----------------
{summaries}
"""

messages_pdf = [
    SystemMessagePromptTemplate.from_template(system_template),
    HumanMessagePromptTemplate.from_template("{question}")
]
pdf_prompt = ChatPromptTemplate.from_messages(messages_pdf)

# Cadena/prompt para conocimiento general (fallback), en caso de que no haya nada relevante en los PDF
fallback_system_template = """\
Eres ChatGPT-4, un modelo de lenguaje altamente analítico y con amplio conocimiento.
Responde en español de manera extensa, detallada y muy analítica.
"""

messages_fallback = [
    SystemMessagePromptTemplate.from_template(fallback_system_template),
    HumanMessagePromptTemplate.from_template("{question}")
]
fallback_prompt = ChatPromptTemplate.from_messages(messages_fallback)


@cl.on_chat_start
async def on_chat_start():
    await cl.Message(
        content="¡Bienvenido! Estoy listo para ayudarte con gestión de conflictos y cualquier otra pregunta que tengas."
    ).send()

    # Rutas de los PDFs
    pdf_paths = [
        "gestios de conflictos.pdf",
        "Managing Conflict with Your Boss .pdf"
    ]
    
    all_texts = []
    all_metadatas = []
    
    # Procesar cada PDF: extraer texto, dividirlo en fragmentos y asignar metadata
    for path in pdf_paths:
        base_name = os.path.basename(path)
        with open(path, "rb") as f:
            reader = PyPDF2.PdfReader(f)
            pdf_text = ""
            for page in reader.pages:
                text = page.extract_text()
                if text:
                    pdf_text += text
        
        chunks = text_splitter.split_text(pdf_text)
        all_texts.extend(chunks)
        all_metadatas.extend([{"source": base_name} for _ in chunks])
    
    # Crear la base vectorial usando nuestra clase personalizada de embeddings
    embeddings = CustomOpenAIEmbeddings(openai_api_key=OPENAI_API_KEY)
    docsearch = await cl.make_async(Chroma.from_texts)(
        all_texts,
        embeddings,
        metadatas=all_metadatas,
        persist_directory="./chroma_db"  # Directorio de persistencia (ajústalo si necesitas)
    )
    
    # Cadena para preguntas que sí tengan match en los PDFs
    pdf_chain = RetrievalQAWithSourcesChain.from_chain_type(
        llm=ChatOpenAI(
            temperature=0.7,
            model_name="gpt-4",          # Asegúrate de que tu cuenta tenga acceso a GPT-4
            openai_api_key=OPENAI_API_KEY,
            max_tokens=2000
        ),
        chain_type="stuff",
        retriever=docsearch.as_retriever(),
        chain_type_kwargs={"prompt": pdf_prompt}
    )
    
    # Cadena de fallback para preguntas fuera de contexto PDF (conocimiento general)
    fallback_chain = LLMChain(
        llm=ChatOpenAI(
            temperature=0.7,
            model_name="gpt-4",          # Asegúrate de que tu cuenta tenga acceso a GPT-4
            openai_api_key=OPENAI_API_KEY,
            max_tokens=2000
        ),
        prompt=fallback_prompt
    )

    # Guardar en la sesión de usuario
    cl.user_session.set("pdf_chain", pdf_chain)
    cl.user_session.set("fallback_chain", fallback_chain)
    cl.user_session.set("metadatas", all_metadatas)
    cl.user_session.set("texts", all_texts)
    
    await cl.Message(content="¡Listo! Puedes comenzar a hacer tus preguntas.").send()


@cl.on_message
async def main(message: cl.Message):
    query = message.content
    pdf_chain = cl.user_session.get("pdf_chain")
    fallback_chain = cl.user_session.get("fallback_chain")
    metadatas = cl.user_session.get("metadatas")
    texts = cl.user_session.get("texts")
    
    # Callback para hacer streaming de la respuesta
    cb = cl.AsyncLangchainCallbackHandler(
        stream_final_answer=True,
        answer_prefix_tokens=["FINAL", "ANSWER"]
    )
    cb.answer_reached = True
    
    # 1) Intentar obtener respuesta del PDF chain
    res = await pdf_chain.acall(query, callbacks=[cb])
    answer = res["answer"]
    sources = res["sources"].strip()
    
    # Verificamos si la respuesta indica que no se encontró nada relevante
    # o si la cadena devolvió algo muy corto que parezca "No lo sé".
    # Ajusta la condición según tu preferencia.
    if ("no lo sé" in answer.lower()) or ("no sé" in answer.lower()) or (len(answer) < 30):
        # 2) Fallback a la cadena de conocimiento general
        res_fallback = await fallback_chain.acall({"question": query})
        answer = res_fallback["text"]
        # En fallback no tenemos "FUENTES", pues responde con conocimiento general
        sources = ""
    else:
        # Agregar fuentes si las hay
        if sources:
            # Buscamos los fragmentos correspondientes
            found_sources = []
            source_elements = []
            
            all_sources = [m["source"] for m in metadatas]
            for src in sources.split(","):
                src_name = src.strip().replace(".", "")
                try:
                    index = all_sources.index(src_name)
                except ValueError:
                    continue
                found_sources.append(src_name)
                source_elements.append(cl.Text(content=texts[index], name=src_name))
            
            if found_sources:
                answer += f"\n\nFUENTES: {', '.join(found_sources)}"
                
            # Si estamos haciendo streaming, actualizamos el mensaje con los elementos
            if cb.has_streamed_final_answer:
                cb.final_stream.elements = source_elements
                await cb.final_stream.update()
                return
            else:
                # Si no hubo streaming, mandamos el mensaje completo al final
                await cl.Message(content=answer, elements=source_elements).send()
                return

    # Si llegamos aquí, simplemente enviamos la respuesta (sea PDF o fallback)
    # y no hay fuentes que mostrar (o ya se procesaron).
    if cb.has_streamed_final_answer:
        # Si fue streaming, actualizamos el mensaje final sin fuentes
        await cb.final_stream.update(content=answer)
    else:
        await cl.Message(content=answer).send()


if __name__ == "__main__":
    from chainlit.cli import run_chainlit
    file_name = __file__ if '__file__' in globals() else "app.py"
    run_chainlit(file_name)