import requests
import os
from typing import Optional

HUGGING_FACE_TOKEN = os.environ["HUGGING_FACE_TOKEN"]
MODEL_API_URL = (
    "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.1"
)

headers = {"Authorization": f"Bearer {HUGGING_FACE_TOKEN}"}


def query(payload: Optional[dict] = None, inputs: Optional[str] = None):
    if payload is None:
        payload = {
            "inputs": inputs,
            "parameters": {
                "max_new_tokens": 240,
                "return_full_text": False,
            },
            "options": {"use_cache": True, "wait_for_model": True},
        }
    response = requests.post(MODEL_API_URL, headers=headers, json=payload)
    return response.json()[0]["generated_text"]