add router to smalltalk chain and rang chain

app.py

@@ -3,11 +3,14 @@ from st_audiorec import st_audiorec
 from Modules.Speech2Text.transcribe import transcribe
 import base64
 from langchain_mistralai import ChatMistralAI
+from langchain_core.prompts import ChatPromptTemplate
 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.PoseEstimation.pose_agent import agent_executor
 
 mistral_api_key = os.getenv("MISTRAL_API_KEY")
 from Modules.PoseEstimation import pose_estimator
@@ -26,6 +29,15 @@ st.set_page_config(layout="wide", initial_sidebar_state="collapsed")
 col1, col2 = st.columns(2)
 video_uploaded = None
 llm = ChatMistralAI(model="mistral-large-latest", mistral_api_key=mistral_api_key, temperature=0)
+prompt = ChatPromptTemplate.from_template(
+    template =""" You are a personal AI sports coach with an expertise in nutrition and fitness. 
+    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 
 
 # First column containers
 with col1:
@@ -56,9 +68,19 @@ with col1:
 
         with st.chat_message("assistant"):
             # Build answer from LLM
-            response = rag_chain.invoke(
-                        instruction
-                        )
+            direction = router_chain.invoke({"question":instruction})
+            if direction=='fitness_advices':
+                response = rag_chain.invoke(
+                            instruction
+                            )
+            elif direction=='smalltalk':
+                response = base_chain.invoke(
+                    {"question":instruction}
+                ).content
+            # elif direction =='movement_analysis':
+            #     response = agent_executor.invoke(
+            #         {"input" : instruction}
+            #     )["output"]
             print(type(response))
             st.session_state.messages.append({"role": "assistant", "content": response})
             st.markdown(response)