MultiDocument RAG con LangChain

Uno de los desafios al implementar RAG es que puede ser que al realizar una consulta obtenga la informacion del documento equivocado.

Para solucionar este problema LangChain tiene disponible create_vectorstore_router_agent, que basado en la pregunta decide que documento debe usar y utiliza ese para obtener la informacion, sino encuentra nada en el documento seleccionado va al siguiente y busca informacion.

from langchain.agents import AgentType
from langchain.agents.agent_toolkits import (
    create_vectorstore_router_agent,
    VectorStoreRouterToolkit,
    VectorStoreInfo,
)
from langchain.document_loaders import UnstructuredPDFLoader
from langchain.vectorstores import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_huggingface.embeddings import HuggingFaceEmbeddings
from langchain_openai import ChatOpenAI
import os
import re

dbs = []
query_rag_agent_agent = None

def query_rag_agent():
    global dbs 
    global query_rag_agent_agent
    rag_llm = ChatOpenAI(model="gpt-3.5-turbo")
    embeddings = HuggingFaceEmbeddings(
        model_name="sentence-transformers/all-mpnet-base-v2", model_kwargs={"device": "cpu"}, encode_kwargs={"normalize_embeddings": False}
    )

    # Solo crea vectores si no existen
    if len(dbs) == 0:
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000, chunk_overlap=0
        )

        for file in os.listdir("source-folder"):
            if re.search("pdf", file):
                loader = UnstructuredPDFLoader(
                    os.path.join("source-folder", file)
                )
                documents = loader.load()
                doc_texts = text_splitter.split_documents(documents)
                vStore = Chroma.from_documents(doc_texts, embeddings)
                ruff_vectorstore_info = VectorStoreInfo(
                    name=f"{file}", description=f"{file}", vectorstore=vStore
                )
                dbs.append(ruff_vectorstore_info)

    # Solo creamos agente sino existe
    if query_rag_agent_agent is None:
        router_toolkit = VectorStoreRouterToolkit(
            vectorstores=dbs, llm=rag_llm
        )
        query_rag_agent_agent = create_vectorstore_router_agent(
            llm=rag_llm,
            toolkit=router_toolkit,
            verbose=True,
            agent=AgentType.SELF_ASK_WITH_SEARCH,
            prefix="Responde en español las siguientes preguntas lo mejor que puedas. Tienes acceso a las siguientes herramientas:"
        )

    question= input("question: ")
    output = query_rag_agent_agent.invoke({"input": question})
    print(output)
  
query_rag_agent()

Este codigo toma todos los documentos PDF de la carpeta llamada “source-folder” y los convierte en DB de vectores, luego crea un agente que tiene acceso a estas bases de datos. Tambien guarda el nombre del archivo como dato que puede ser usado al decidir si el documento es relevante para responder la pregunta.

No solo puedes cargar PDFs a la base de datos de vectores sino otros tipos de archivo como TXT, docx y WebScraping. Pero en este ejemplo solo cargaremos PDFs para mantenerlo sencillo.

Para poder ejecutar el codigo debes asegurarse de instalar las dependencias con:

pip install langchain unstructured[pdf] langchain-huggingface langchain-openai langchain_chroma langchain-community

Para ejecutar el ejemplo, obtuvimos dos archivos sobre preguntas frecuentes, uno sobre el censo y otro sobre beneficios para estudiantes.

El agente nos mostrara el proceso de pensamiento porque indicamos Verbose=True. Al preguntar cuando termina el censo podemos ver las acciones del agente:

Respuesta: “El censo termina el viernes 28 de junion 2024”

Si preguntamos sobre los beneficios estudiantiles vemos que utiliza el otro documento para responder la pregunta:

Respuesta: “Para completar el FUAS, debes ingresar al sitio web www.fuas.cl, registrarte, completar el formulario con tus datos personales y los de tu grupo familiar, revisar la información ingresada y finalizar el proceso.“