Hacemos refactor de ingest.py

app.py

@@ -7,8 +7,7 @@ from langchain.globals import set_verbose, set_debug
 
 
 def respond(message, history):
-    prompt = llm.getMainTemplate()
-    chain = prompt | llm.llm
+    chain = llm.main_prompt | llm.llm
     response = chain.invoke({"message": message, "history": history})
     return response.content
 

chatbot/llm.py

@@ -8,26 +8,22 @@ class GeminiAI:
     Google Gemini AI class.
     """
 
+    system_prompt = """
+        Eres un asesor experto en la Resolucion 1650/24 de la DGCyE de la Provincia de Buenos Aires.\n
+        Tu tarea es utiliza la información de la conversación y el contexto disponible para responder las consultas del usuario.
+        """
+
     def __init__(self, llm_model_name: str) -> None:
         self.llm = ChatGoogleGenerativeAI(model=llm_model_name)
 
-    def getMainTemplate(self) -> ChatPromptTemplate:
-        """
-        Devuelve el system prompt principal.
-        """
-
-        prompt = ChatPromptTemplate.from_messages(
+        self.main_prompt = ChatPromptTemplate.from_messages(
             [
-                ("system",
-                "Eres un asesor experto en la Resolucion 1650/24 de la DGCyE de la Provincia de Buenos Aires.\n"
-                "Tu tarea es utiliza la información de la conversación y el contexto disponible para responder las consultas del usuario.\n"),
+                ("system",self.system_prompt),
                 ("placeholder", "{history}"),
                 ("human", "{message}"),
             ]
         )
 
-        return prompt
-
 
 
 if __name__ == "__main__":

ingest.py

@@ -3,23 +3,55 @@ Modulo para procesar el PDF de la resolucion e indexar su contenido en la DB, pa
 
 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 langchain_community.document_loaders import PyPDFLoader
 from chatbot.embeddings import init_embeddings
 from chatbot.vectorstore import ChromaDB
+from typing import List
+import sys, argparse
 
-if __name__ == "__main__":
-    #loader = PyPDFLoader("2024_DP_134.pdf")
+
+
+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])
+    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)
 
-    results = vector_store.db.similarity_search(
-        "Cuantos anexos contiene la resolucion?",
-        k=2,
-    )
+    sys.exit(0)
 
-    print(results)