Update app.py

app.py

@@ -26,20 +26,31 @@ load_dotenv()
 import io
 import contextlib
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from smolagents import Tool
 
-class ZephyrToolCallingAgentTool:
+class ZephyrToolCallingAgentTool(Tool):
     name = "zephyr_tool_agent"
     description = "Uses Zephyr-7B to answer questions using code or reasoning"
 
+    inputs = {
+        "question": {
+            "type": "string",
+            "description": "The user's question involving reasoning or code execution."
+        }
+    }
+    output_type = "string"
+
     def __init__(self):
         self.model_id = "HuggingFaceH4/zephyr-7b-beta"
-        token = os.getenv("HF_TOKEN")  # Optional unless private
+        token = os.getenv("HF_TOKEN")
 
         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
+            self.model_id,
+            device_map="auto",
+            torch_dtype="auto",
+            token=token
         )
-
         self.pipeline = pipeline(
             "text-generation",
             model=self.model,
@@ -48,30 +59,21 @@ class ZephyrToolCallingAgentTool:
             temperature=0.2
         )
 
-    def _run_code(self, code: str) -> str:
-        buffer = io.StringIO()
-        try:
-            with contextlib.redirect_stdout(buffer):
-                exec(code, {})
-            return buffer.getvalue().strip()
-        except Exception as e:
-            return f"Error during code execution: {e}"
+    def _run(self, question: str) -> str:
+        return self.run(question)
 
     def run(self, question: str) -> str:
         prompt = f"""You are a helpful assistant. Use code to solve questions that involve calculations.
 If code is needed, return a block like <tool>code</tool>. End your answer with <final>answer</final>.
 
-
 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"FINAL ANSWER (code output): {output}"
+            return self._run_code(code)
 
         elif "<final>" in result and "</final>" in result:
             final = result.split("<final>")[1].split("</final>")[0].strip()
@@ -79,6 +81,15 @@ Answer:"""
 
         return "Could not determine how to respond. No <tool> or <final> block detected."
 
+    def _run_code(self, code: str) -> str:
+        buffer = io.StringIO()
+        try:
+            with contextlib.redirect_stdout(buffer):
+                exec(code, {})
+            return f"FINAL ANSWER (code output): {buffer.getvalue().strip()}"
+        except Exception as e:
+            return f"Error during code execution: {e}"
+
 
 #from smolagents import Tool
 #from langchain_community.document_loaders import WikipediaLoader
@@ -253,7 +264,7 @@ class BasicAgent:
         video_transcription_tool = VideoTranscriptionTool()
 
         # ✅ New Mistral-based Tool
-        mistral_tool = ZephyrToolCallingAgentTool()
+        zephyr_tool = ZephyrToolCallingAgentTool()
 
         system_prompt = f"""
 You are my general AI assistant. Your task is to answer the question I asked.
@@ -272,7 +283,7 @@ If the answer is a comma-separated list, apply the above rules for each element
                 keywords_extract_tool, speech_to_text_tool,
                 visit_webpage_tool, final_answer_tool,
                 parse_excel_to_json, video_transcription_tool,
-                mistral_tool  # 🔧 Add here
+                zephyr_tool  # 🔧 Add here
             ],
             add_base_tools=True
         )