Update app.py

app.py

@@ -35,11 +35,25 @@ class MistralToolCallingAgentTool:
 
     def __init__(self):
         self.model_id = "mistralai/Mistral-7B-Instruct-v0.3"
-        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)
-        self.model = AutoModelForCausalLM.from_pretrained(self.model_id, device_map="auto", torch_dtype="auto")
+        token = os.getenv("HF_TOKEN")
+
+        if token is None:
+            raise EnvironmentError("HF_TOKEN is not set in environment variables.")
+
+        try:
+            self.tokenizer = AutoTokenizer.from_pretrained(self.model_id, token=token)
+            self.model = AutoModelForCausalLM.from_pretrained(
+                self.model_id, device_map="auto", torch_dtype="auto", token=token
+            )
+        except Exception as e:
+            raise RuntimeError(f"Error loading Mistral model: {e}")
+
         self.pipeline = pipeline(
-            "text-generation", model=self.model, tokenizer=self.tokenizer,
-            max_new_tokens=512, temperature=0.2
+            "text-generation",
+            model=self.model,
+            tokenizer=self.tokenizer,
+            max_new_tokens=512,
+            temperature=0.2
         )
 
     def _run_code(self, code: str) -> str:
@@ -49,7 +63,7 @@ class MistralToolCallingAgentTool:
                 exec(code, {})
             return buffer.getvalue().strip()
         except Exception as e:
-            return f"Error during execution: {e}"
+            return f"Error during code execution: {e}"
 
     def run(self, question: str) -> str:
         prompt = f"""You are a helpful assistant. Use code to solve questions that involve calculations.
@@ -57,19 +71,21 @@ If code is needed, return a block like <tool>code</tool>. End your answer with <
 
 Question: {question}
 Answer:"""
+
         result = self.pipeline(prompt)[0]["generated_text"]
 
         # Process result
         if "<tool>" in result and "</tool>" in result:
             code = result.split("<tool>")[1].split("</tool>")[0].strip()
             output = self._run_code(code)
-            return f"Code result: {output}"
+            return f"FINAL ANSWER (code output): {output}"
 
         elif "<final>" in result and "</final>" in result:
             final = result.split("<final>")[1].split("</final>")[0].strip()
             return f"FINAL ANSWER: {final}"
 
-        return "Mistral agent could not determine how to respond."
+        return "Could not determine how to respond. No <tool> or <final> block detected."
+
 
 #from smolagents import Tool
 #from langchain_community.document_loaders import WikipediaLoader