Update mistral_hf_wrapper.py

mistral_hf_wrapper.py

@@ -1,32 +1,21 @@
-import requests
 import os
+import requests
 
-class MistralInference:
-    def __init__(self, api_url: str, api_token: str):
-        self.api_url = api_url.rstrip("/")
-        self.headers = {
-            "Authorization": f"Bearer {api_token}",
-            "Content-Type": "application/json"
-        }
-
-    def generate(self, prompt: str, temperature: float = 0.7, max_tokens: int = 512) -> str:
-        payload = {
-            "inputs": prompt,
-            "parameters": {
-                "temperature": temperature,
-                "max_new_tokens": max_tokens,
-                "return_full_text": False
-            }
-        }
-
-        response = requests.post(
-            f"{self.api_url}/generate",
-            headers=self.headers,
-            json=payload
-        )
+API_URL = os.getenv("HF_MISTRAL_ENDPOINT")
+API_TOKEN = os.getenv("HF_TOKEN")
 
-        if response.status_code != 200:
-            raise RuntimeError(f"Request failed: {response.status_code} - {response.text}")
+headers = {
+    "Authorization": f"Bearer {API_TOKEN}",
+    "Content-Type": "application/json"
+}
 
-        data = response.json()
-        return data["generated_text"] if "generated_text" in data else data[0]["generated_text"]
+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()