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 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.

<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?"}))