add workout plan generator

Modules/rag.py

@@ -63,6 +63,12 @@ prompt = ChatPromptTemplate.from_template(
     Use the following pieces of retrieved context to answer the question.
     If you don't know the answer, use your common knowledge.
     Use three sentences maximum and keep the answer concise.
+    If the user asks you a full program workout, structure your response in this way (this is an example):
+    - First workout : Lower body (1 hour)
+    1. Barbelle squat / 4 sets of 8 reps / 2'30 recovery
+    2. Lunges / 4 sets of 10 reps / 2'recovery
+    3. etc
+    - Second workout  .... and so on.
 
     Question: {question} 
 
@@ -86,6 +92,6 @@ rag_chain = (
 
 
 
-print(rag_chain.invoke("WHi I'm Susan. Can you make a fitness program for me please?"))
+# print(rag_chain.invoke("WHi I'm Susan. Can you make a fitness program for me please?"))
 
 # 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"))

Modules/router.py

@@ -10,9 +10,10 @@ from langchain_core.output_parsers import StrOutputParser
 router_chain = (
     ChatPromptTemplate.from_template(
         """Given the user question below, classify it as either being about :
-        - `fitness_advices` if the user query is about nutrition or fitness program, exercices
-        - `movement_analysis` if the user asks to analyse or give advice on his exercice execution?
-        - `smalltalk` if other.
+        - 'fitness_advices`' if the user query is about nutrition or fitness strategies, exercices
+        - 'workout_plan' if the user asks for a detailed workout plan or a full fitness program
+        - 'movement_analysis' if the user asks to analyse or give advice on his exercice execution?
+        - 'smalltalk' if other.
 
 Do not respond with more than one word.
 

Modules/workout_plan.py

@@ -0,0 +1,139 @@
+import os
+os.environ['TOKENIZERS_PARALLELISM'] = 'true'
+from dotenv import load_dotenv
+load_dotenv() # load .env api keys 
+
+mistral_api_key = os.getenv("MISTRAL_API_KEY")
+print("mistral_api_key", mistral_api_key)
+import pandas as pd
+from langchain.output_parsers import PandasDataFrameOutputParser
+from langchain_community.document_loaders import PyPDFLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.vectorstores import Chroma
+from langchain_mistralai import MistralAIEmbeddings
+from langchain import hub
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.runnables import RunnablePassthrough
+from typing import Literal
+from langchain_core.prompts import PromptTemplate
+from langchain_mistralai import ChatMistralAI
+from pathlib import Path
+from langchain.retrievers import (
+    MergerRetriever,
+)
+import pprint
+from typing import Any, Dict
+from huggingface_hub import login
+login(token=os.getenv("HUGGING_FACE_TOKEN"))
+
+def load_chunk_persist_pdf(task) -> Chroma:
+    
+    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'):
+            pdf_path = os.path.join(pdf_folder_path, file)
+            loader = PyPDFLoader(pdf_path)
+            documents.extend(loader.load())
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=10)
+    chunked_documents = text_splitter.split_documents(documents)
+    os.makedirs("data/chroma_store/", exist_ok=True)
+    vectorstore = Chroma.from_documents(
+        documents=chunked_documents,
+        embedding=MistralAIEmbeddings(),
+        persist_directory= os.path.join(os.getcwd(),Path("data/chroma_store/"))
+    )
+    vectorstore.persist()
+    return vectorstore
+
+df = pd.DataFrame(
+    {
+        "exercise": ["Squat","Bench Press","Lunges","Pull ups"],
+        "sets": [4, 4, 3, 3],
+        "repetitions": [10, 8, 8, 8],
+        "rest":["2:30","2:00","1:30","2:00"]
+    }
+)
+
+# parser = PandasDataFrameOutputParser(dataframe=df)
+
+# personal_info_vectorstore = load_chunk_persist_pdf("personal_info")
+# 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()
+# personal_info_retriever = personal_info_vectorstore.as_retriever()
+
+llm = ChatMistralAI(model="mistral-large-latest", mistral_api_key=mistral_api_key, temperature=0)
+
+# prompt = PromptTemplate(
+#     template="""
+#     You are a professional AI coach specialized in building fitness plans, full workout programs.
+#     You must adapt to the user according to personal informations in the context. A You are gentle and motivative.
+#     Use the following pieces of retrieved context to answer the user's query.
+
+#     Context: {context} 
+    
+#     \n{format_instructions}\n{question}\n
+#     """,
+#     input_variables=["question","context"],
+#     partial_variables={"format_instructions": parser.get_format_instructions()},
+# )
+
+# def format_docs(docs):
+#         return "\n\n".join(doc.page_content for doc in docs)
+
+# def format_parser_output(parser_output: Dict[str, Any]) -> None:
+#     for key in parser_output.keys():
+#         parser_output[key] = parser_output[key].to_dict()
+#     return pprint.PrettyPrinter(width=4, compact=True).pprint(parser_output)
+    
+# retriever = MergerRetriever(retrievers=[zero2hero_retriever, bodyweight_retriever, nutrition_retriever, workout_retriever, personal_info_retriever])
+
+# chain = (
+#         {"context": zero2hero_retriever | format_docs, "question": RunnablePassthrough()}
+#         | prompt
+#         | llm
+#         | parser
+#     )
+
+# # chain = prompt | llm | parser
+# format_parser_output(chain.invoke("Build me a full body workout plan for summer body."))
+
+
+from pydantic import BaseModel, Field
+from typing import List
+from langchain_core.output_parsers import JsonOutputParser
+
+class Exercise(BaseModel):
+    exercice: str = Field(description="Name of the exercise")
+    nombre_series: int = Field(description="Number of sets for the exercise")
+    nombre_repetitions: int = Field(description="Number of repetitions for the exercise")
+    temps_repos: str = Field(description="Rest time between sets")
+
+class MusculationProgram(BaseModel):
+    exercises: List[Exercise]
+
+
+from langchain.prompts import PromptTemplate
+
+# Define your query to get a musculation program.
+musculation_query = "Provide a musculation program with exercises, number of sets, number of repetitions, and rest time between sets."
+
+# Set up a parser + inject instructions into the prompt template.
+parser = JsonOutputParser(pydantic_object=MusculationProgram)
+
+prompt = PromptTemplate(
+    template="Answer the user query.\n{format_instructions}\n{query}\n",
+    input_variables=["query"],
+    partial_variables={"format_instructions": parser.get_format_instructions()},
+)
+
+# Set up a chain to invoke the language model with the prompt and parser.
+workout_chain = prompt | llm | parser
+
+

app.py

@@ -4,12 +4,15 @@ from Modules.Speech2Text.transcribe import transcribe
 import base64
 from langchain_mistralai import ChatMistralAI
 from langchain_core.prompts import ChatPromptTemplate
+import pandas as pd 
+import json
 from dotenv import load_dotenv
 load_dotenv() # load .env api keys 
 import os
 
 from Modules.rag import rag_chain
 from Modules.router import router_chain
+from Modules.workout_plan import workout_chain
 # from Modules.PoseEstimation.pose_agent import agent_executor
 
 mistral_api_key = os.getenv("MISTRAL_API_KEY")
@@ -27,11 +30,13 @@ prompt = ChatPromptTemplate.from_template(
     You are having a conversation with your client, which is either a beginner or an advanced athlete. 
     You must be gentle, kind, and motivative.
     Always try to answer concisely to the queries.
+    
     User: {question}
     AI Coach:"""
 )
 base_chain = prompt | llm 
 
+display_workout = False
 # First column containers
 with col1:
     
@@ -52,6 +57,8 @@ with col1:
         with st.chat_message("assistant"):
             # Build answer from LLM
             direction = router_chain.invoke({"question":prompt})
+            print(type(direction))
+            print(direction)
             if direction=='fitness_advices':
                 response = rag_chain.invoke(
                             prompt
@@ -60,15 +67,26 @@ with col1:
                 response = base_chain.invoke(
                     {"question":prompt}
                 ).content
-            # elif direction =='movement_analysis':
+            elif direction =='movement_analysis':
+                response = "I can't do that for the moment"
             #     response = agent_executor.invoke(
             #         {"input" : instruction}
             #     )["output"]
+            # elif direction == 'workout_plan':
+            else:
+                response = "Sure! I just made a workout for you. Check on the table I just provided you."
+                json_output = workout_chain.invoke({"query":prompt})
+                exercises_list = json_output['exercises']
+                workout_df = pd.DataFrame(exercises_list)
+                workout_df.columns = ["exercice", "nombre_series", "nombre_repetitions", "temps_repos"]
+                display_workout=True
             print(type(response))
             st.session_state.messages.append({"role": "assistant", "content": response})
             st.markdown(response)
 
-    st.subheader("Movement Analysis")
+    if display_workout:
+        st.subheader("Workout")
+        st.data_editor(workout_df)
         # TO DO 
 # Second column containers
 with col2: