import os
import requests
API_URL = "https://router.huggingface.co/featherless-ai/v1/chat/completions"
headers = {
    "Authorization": f"Bearer {os.environ['HF_API_KEY']}",
}


def query_test(payload):
    payload = {
        "messages": [
            {
                "role": "user",
                "content": "What is the capital of France?"
            }
        ],
        "model": "mistralai/Mistral-7B-Instruct-v0.2"
    }
    response = requests.post(API_URL, headers=headers, json=payload)
    return response.json()["choices"][0]["message"]