from huggingface_hub import InferenceClient
import os

HF_TOKEN = os.getenv("HF_TOKEN")

client = InferenceClient(
    model="mistralai/Mistral-7B-Instruct-v0.2",
    token=HF_TOKEN
)


def query_model(prompt):

    response = client.chat.completions.create(
        messages=[
            {"role": "system", "content": "You are a professional fitness trainer"},
            {"role": "user", "content": prompt}
        ],
        max_tokens=700
    )

    return response.choices[0].message.content