Update app.py

app.py

@@ -25,17 +25,19 @@ load_dotenv()
 
 import io
 import contextlib
+import traceback
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-from smolagents import Tool, CodeAgent, DuckDuckGoSearchTool, HfApiModel, FinalAnswerTool
+from smolagents import Tool, CodeAgent, DuckDuckGoSearchTool, FinalAnswerTool, HfApiModel
+
 
 class CodeLlamaTool(Tool):
     name = "code_llama_tool"
-    description = "Uses Code Llama 7B Instruct to answer questions with code or reasoning"
+    description = "Solves reasoning/code questions using Meta Code Llama 7B Instruct"
 
     inputs = {
         "question": {
             "type": "string",
-            "description": "The user's question requiring reasoning or code execution."
+            "description": "The question requiring code-based or reasoning-based solution"
         }
     }
     output_type = "string"
@@ -46,43 +48,69 @@ class CodeLlamaTool(Tool):
 
         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,
             tokenizer=self.tokenizer,
             max_new_tokens=512,
-            truncation=True,
-            temperature=0.1            
+            temperature=0.2,
+            truncation=True
         )
 
     def forward(self, question: str) -> str:
-        prompt = f"""
-    "You are an assistant. Solve step-by-step. "
-    "End your answer with a concise answer ina few words: FINAL ANSWER: [value]"
-"""
+        prompt = f"""You are an AI that uses Python code to answer questions.
 
-        result = self.pipeline(prompt)[0]["generated_text"]
+Question: {question}
 
-        if "<tool>" in result and "</tool>" in result:
-            code = result.split("<tool>")[1].split("</tool>")[0].strip()
-            return self._run_code(code)
+Instructions:
+- If solving requires code, use a block like <tool>code</tool>.
+- Always end with <final>FINAL ANSWER</final> containing the final number or string.
 
-        elif "<final>" in result and "</final>" in result:
-            final = result.split("<final>")[1].split("</final>")[0].strip()
-            return f"FINAL ANSWER: {final}"
+Example:
+Question: What is 5 * sqrt(36)?
+Answer:
+<tool>
+import math
+print(5 * math.sqrt(36))
+</tool>
+<final>30.0</final>
 
-        return "Could not determine how to respond. No <tool> or <final> block detected."
+Answer:"""
+
+        response = self.pipeline(prompt)[0]["generated_text"]
+        return self.parse_and_execute(response)
+
+    def parse_and_execute(self, response: str) -> str:
+        try:
+            # Extract and run code if exists
+            if "<tool>" in response and "</tool>" in response:
+                code = response.split("<tool>")[1].split("</tool>")[0].strip()
+                result = self._run_code(code)
+                return f"FINAL ANSWER (code output): {result}"
+
+            # Extract final result directly
+            elif "<final>" in response and "</final>" in response:
+                final = response.split("<final>")[1].split("</final>")[0].strip()
+                return f"FINAL ANSWER: {final}"
+
+            return f"Could not extract final answer.\n\n{response}"
+
+        except Exception as e:
+            return f"Error in parse_and_execute: {str(e)}\n\nFull response:\n{response}"
 
     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}"
+            return buffer.getvalue().strip()
+        except Exception:
+            return f"Error executing code:\n{traceback.format_exc()}"
 
 
 
@@ -259,7 +287,7 @@ class BasicAgent:
         video_transcription_tool = VideoTranscriptionTool()
 
         # ✅ New Llama Tool
-        my_tool = CodeLlamaTool()
+        code_llama_tool = CodeLlamaTool()
 
         system_prompt = f"""
 You are my general AI assistant. Your task is to answer the question I asked.
@@ -278,7 +306,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,
-                my_tool  # 🔧 Add here
+                code_llama_tool  # 🔧 Add here
             ],
             add_base_tools=True
         )