augment rag docs

Modules/rag.py

@@ -7,26 +7,22 @@ mistral_api_key = os.getenv("MISTRAL_API_KEY")
 
 from langchain_community.document_loaders import PyPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_community.document_loaders import WebBaseLoader
-from langchain_community.vectorstores import Chroma, FAISS
-from langchain.chains.combine_documents import create_stuff_documents_chain
+from langchain_community.vectorstores import Chroma
 from langchain_mistralai import MistralAIEmbeddings
 from langchain import hub
-from langchain.chains import (
-    create_history_aware_retriever,
-    create_retrieval_chain,
-)
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.runnables import RunnablePassthrough
 from typing import Literal
-from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
-from langchain_core.pydantic_v1 import BaseModel, Field
+from langchain_core.prompts import ChatPromptTemplate
 from langchain_mistralai import ChatMistralAI
-from langchain.embeddings.sentence_transformer import SentenceTransformerEmbeddings
-from langchain_community.tools import DuckDuckGoSearchRun
 from pathlib import Path
 
-def load_chunk_persist_pdf() -> Chroma:
+from langchain.retrievers import (
+    MergerRetriever,
+)
+def load_chunk_persist_pdf(task) -> Chroma:
     
-    pdf_folder_path = os.path.join(os.getcwd(),Path("data/pdf/"))
+    pdf_folder_path = os.path.join(os.getcwd(),Path(f"data/pdf/{task}"))
     documents = []
     for file in os.listdir(pdf_folder_path):
         if file.endswith('.pdf'):
@@ -44,30 +40,23 @@ def load_chunk_persist_pdf() -> Chroma:
     vectorstore.persist()
     return vectorstore
 
-vectorstore = load_chunk_persist_pdf()
-retriever = vectorstore.as_retriever()
-prompt = hub.pull("rlm/rag-prompt")
-
-
-# Data model
-class RouteQuery(BaseModel):
-    """Route a user query to the most relevant datasource."""
+zero2hero_vectorstore = load_chunk_persist_pdf("zero2hero")
+bodyweight_vectorstore = load_chunk_persist_pdf("bodyweight")
+nutrition_vectorstore = load_chunk_persist_pdf("nutrition")
+workout_vectorstore = load_chunk_persist_pdf("workout")
+zero2hero_retriever = zero2hero_vectorstore.as_retriever()
+nutrition_retriever = nutrition_vectorstore.as_retriever()
+bodyweight_retriever = bodyweight_vectorstore.as_retriever()
+workout_retriever = workout_vectorstore.as_retriever()
 
-    datasource: Literal["vectorstore", "websearch"] = Field(
-        ...,
-        description="Given a user question choose to route it to web search or a vectorstore.",
-    )
-
-# LLM with function call 
 llm = ChatMistralAI(model="mistral-large-latest", mistral_api_key=mistral_api_key, temperature=0)
 
-
 prompt = ChatPromptTemplate.from_template(
     """
     You are a professional AI coach specialized in fitness, bodybuilding and nutrition.
     You must adapt to the user : if he is a beginner, use simple words. You are gentle and motivative.
     Use the following pieces of retrieved context to answer the question.
-    If you don't know the answer, just say that you don't know, and to refer to a nutritionist or a doctor.
+    If you don't know the answer, use your common knowledge.
     Use three sentences maximum and keep the answer concise.
 
     Question: {question} 
@@ -77,12 +66,11 @@ prompt = ChatPromptTemplate.from_template(
     Answer:
     """, 
 )
-from langchain_core.output_parsers import StrOutputParser
-from langchain_core.runnables import RunnablePassthrough
 
 def format_docs(docs):
         return "\n\n".join(doc.page_content for doc in docs)
     
+retriever = MergerRetriever(retrievers=[zero2hero_retriever, bodyweight_retriever, nutrition_retriever, workout_retriever])
 
 rag_chain = (
         {"context": retriever | format_docs, "question": RunnablePassthrough()}
@@ -92,6 +80,7 @@ rag_chain = (
     )
 
 
-# print(rag_chain.invoke("Build a fitness program for me. Be precise in terms of exercises"))
 
-# print(rag_chain.invoke("I am a 45 years old woman and I have to loose weight for the summer. Provide me with a fitness program"))
+print(rag_chain.invoke("What supplement could i buy to improve my sleep?"))
+
+# print(rag_chain.invoke("I am a 45 years old woman and I have to loose weight for the summer. Provide me with a fitness program, and a nutrition program"))