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)