Update mistral_hf_wrapper.py

mistral_hf_wrapper.py

@@ -1,37 +1,32 @@
-# mistral_hf_wrapper.py
-
 import os
 import requests
 
 class MistralInference:
-    def __init__(self):
-        self.api_url = os.getenv("HF_MISTRAL_URL")
-        self.api_token = os.getenv("HF_TOKEN")
-        if not self.api_url or not self.api_token:
-            raise ValueError("Missing HF_MISTRAL_URL or HF_TOKEN environment variables")
+    def __init__(self, api_url=None, api_token=None):
+        self.api_url = api_url or os.getenv("HF_MISTRAL_ENDPOINT")
+        self.api_token = api_token or os.getenv("HF_TOKEN")
 
     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,
-                "temperature": 0.7,
-                "return_full_text": False
-            }
+            "parameters": {"max_new_tokens": 512}
         }
-
-        response = requests.post(self.api_url, headers=headers, json=payload)
-        response.raise_for_status()
-
-        output = response.json()
-        if isinstance(output, list) and "generated_text" in output[0]:
-            return output[0]["generated_text"]
-        elif isinstance(output, dict) and "generated_text" in output:
-            return output["generated_text"]
-        else:
-            raise ValueError("Unexpected response format from Mistral endpoint")
+        try:
+            response = requests.post(self.api_url, headers=headers, json=payload)
+            response.raise_for_status()
+            output = response.json()
+            # Check different possible keys depending on model
+            if isinstance(output, list) and "generated_text" in output[0]:
+                return output[0]["generated_text"]
+            elif "generated_text" in output:
+                return output["generated_text"]
+            elif "text" in output:
+                return output["text"]
+            else:
+                return str(output)
+        except Exception as e:
+            return f"Error querying Mistral: {str(e)}"