| """ | |
| 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) | |