Update agent.py

agent.py

@@ -1,6 +1,30 @@
-class BasicAgent:
-    def __init__(self):
-        pass
+import requests
+import os
+
+class MistralAgent:
+    def __init__(self, hf_token: str, model_id: str = "mistralai/Mistral-7B-Instruct-v0.2"):
+        self.api_url = f"https://api-inference.huggingface.co/models/{model_id}"
+        self.headers = {
+            "Authorization": f"Bearer {hf_token}",
+            "Content-Type": "application/json"
+        }
 
     def run(self, question: str) -> str:
-        return "This is a default answer."
+        payload = {
+            "inputs": f"[INST] {question} [/INST]",
+            "parameters": {
+                "temperature": 0.2,
+                "max_new_tokens": 512,
+                "return_full_text": False
+            }
+        }
+        response = requests.post(self.api_url, headers=self.headers, json=payload)
+        if response.status_code != 200:
+            return f"Error: {response.text}"
+
+        output = response.json()
+        if isinstance(output, list) and "generated_text" in output[0]:
+            return output[0]["generated_text"].strip()
+        else:
+            return "Error: Unexpected API output."
+