Modules/router.py

import os
from dotenv import load_dotenv
load_dotenv() 
mistral_api_key = os.getenv("MISTRAL_API_KEY")
from langchain_core.prompts import ChatPromptTemplate
from langchain_mistralai import ChatMistralAI
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.

Do not respond with more than one word.

<question>
{question}
</question>

Classification:"""
    )
    | ChatMistralAI(model="mistral-large-latest", mistral_api_key=mistral_api_key, temperature=0)
    | StrOutputParser()
)

# print(chain.invoke({"question": "Hi AI Bro coach! how are you? I was thinking about going to the gym today. What can i train?"}))