mistral_hf_wrapper.py

import os
import requests

API_URL = os.getenv("HF_MISTRAL_ENDPOINT")
API_TOKEN = os.getenv("HF_TOKEN")

headers = {
    "Authorization": f"Bearer {API_TOKEN}",
    "Content-Type": "application/json"
}

def query_mistral(system_prompt: str, user_prompt: str) -> str:
    """Query the Mistral model hosted on Hugging Face."""
    prompt = f"<s>[INST] {system_prompt.strip()}\n\n{user_prompt.strip()} [/INST]"
    response = requests.post(
        API_URL,
        headers=headers,
        json={"inputs": prompt}
    )
    response.raise_for_status()
    return response.json()["generated_text"].strip()