Update mistral_hf_wrapper.py

mistral_hf_wrapper.py

@@ -1,21 +1,20 @@
 import os
 import requests
 
-API_URL = os.getenv("HF_MISTRAL_ENDPOINT")
-API_TOKEN = os.getenv("HF_TOKEN")
+class MistralInference:
+    def __init__(self):
+        self.api_url = os.getenv("HF_MISTRAL_URL")
+        self.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()
+    def run(self, prompt: str) -> str:
+        headers = {
+            "Authorization": f"Bearer {self.api_token}",
+            "Content-Type": "application/json"
+        }
+        payload = {
+            "inputs": prompt,
+            "parameters": {"max_new_tokens": 512}
+        }
+        response = requests.post(self.api_url, headers=headers, json=payload)
+        response.raise_for_status()
+        return response.json()[0]["generated_text"]