Primera version de RAG

.gitignore

@@ -1,4 +1,5 @@
 Pipfile.lock
 *.pdf
+*.zip
 .env
 chroma_db/

app.py

@@ -3,18 +3,30 @@ Chatbot Nuevo Régimen Académico
 """
 from chatbot.ui import ChatbotInterface
 from chatbot.llm import GeminiAI
+from chatbot.embeddings import init_embeddings
+from chatbot.vectorstore import ChromaDB
 from langchain.globals import set_verbose, set_debug
 
 
+def format_docs(docs):
+    return "\n\n".join(doc.page_content for doc in docs)
+
+
 def respond(message, history):
+    docs = retriever.invoke(message)
     chain = llm.main_prompt | llm.llm
-    response = chain.invoke({"message": message, "history": history})
+
+    response = chain.invoke({"context": format_docs(docs), "message": message, "history": history})
     return response.content
 
 
 if __name__ == "__main__":
     set_verbose(True)
     set_debug(True)
+
+    embeddings = init_embeddings()
+    retriever = ChromaDB(embeddings).db.as_retriever(search_type="similarity", search_kwargs={"k": 3})
+
     llm = GeminiAI("gemini-1.5-flash")
     ui = ChatbotInterface(respond)
 

chatbot/llm.py

@@ -9,8 +9,9 @@ class GeminiAI:
     """
 
     system_prompt = """
-        Eres un asesor experto en la Resolucion 1650/24 de la DGCyE de la Provincia de Buenos Aires.\n
+        Eres un asesor experto en la Resolucion 1650/24 de la DGCyE de la Provincia de Buenos Aires.
         Tu tarea es utiliza la información de la conversación y el contexto disponible para responder las consultas del usuario.
+        Contexto: {context}
         """
 
     def __init__(self, llm_model_name: str) -> None: