ingest.py

"""
Modulo para procesar el PDF de la resolucion e indexar su contenido en la DB, para su posterior utilización por parte del chatbot.

Por simplicidad, se indexo un documento por cada página completa del documento. TODO: Implementar estrategia ParentDocumentRetriever.
"""
from langchain_community.document_loaders import PyPDFLoader
from chatbot.embeddings import init_embeddings
from chatbot.vectorstore import ChromaDB
from typing import List
import sys, argparse



def query(consulta: str) -> List:
    """ Realiza una consulta a la vectorstore. """
    embedding_model = init_embeddings()
    vector_store = ChromaDB(embedding_model)
    docs = vector_store.db.similarity_search(
            consulta,
            k=2,
        )
    
    return docs



def ingest(filename: str) -> None:
    """ Ingiere un archivo PDF y lo indexa en la vectorstore. """
    loader = PyPDFLoader(filename)
    embedding_model = init_embeddings()
    vector_store = ChromaDB(embedding_model)

    for page in loader.lazy_load():
        print(f"Procesando pagina {page.metadata['page']} - len: {len(page.page_content)}")
        vector_store.add_documents([page])



if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Este script permite interactuar con la vectorstore.')
    parser.add_argument('operacion', 
                        type=str,
                        choices=["ingest", "query"],
                        help="La operación que se desea realizar sobre la DB. 'ingest' para agregar un PDF a la DB. 'query' para hacer una consulta.")
    parser.add_argument("argumento", help="Nombre del archivo a ingerir o consulta a realizar.")
    args = parser.parse_args()

    if args.operacion == 'ingest':
        ingest(args.argumento)
        print("Listo.")

    elif args.operacion == 'query':
        docs = query(args.argumento)
        print(docs)

    sys.exit(0)