Update agent.py

agent.py

@@ -7,11 +7,8 @@ import tempfile
 import os
 import whisper
 
-
 SYSTEM_PROMPT = """
-You are a general AI assistant. I will ask you a question. Think step by step to find the best possible answer.
-Then return only the answer without any explanation or formatting.
-Do not say 'Final answer' or anything else. Just output the raw answer string.
+You are a helpful AI assistant. Think step by step to solve the problem. If the question requires reasoning, perform it. If it refers to a search or file, use the result provided. At the end, return ONLY the final answer string. No explanations.
 """
 
 class GaiaAgent:
@@ -29,44 +26,42 @@ class GaiaAgent:
                 if results:
                     return results[0]['body']
         except Exception as e:
-            return f"Search failed: {e}"
+            return ""
         return ""
 
     def transcribe_audio(self, file_path: str) -> str:
         try:
             result = self.transcriber.transcribe(file_path)
             return result['text']
-        except Exception as e:
-            return f"Audio transcription failed: {e}"
+        except Exception:
+            return ""
 
     def handle_excel(self, file_path: str) -> str:
         try:
             df = pd.read_excel(file_path)
-            food_sales = df[df['Category'].str.lower() != 'drink']['Sales'].sum()
-            return f"{food_sales:.2f}"
-        except Exception as e:
-            return f"Excel parsing failed: {e}"
+            df.columns = [col.lower() for col in df.columns]
+            if 'category' in df.columns and 'sales' in df.columns:
+                food_sales = df[df['category'].str.lower() != 'drink']['sales'].sum()
+                return f"{food_sales:.2f}"
+        except Exception:
+            return ""
+        return ""
 
     def __call__(self, question: str, files: dict = None) -> tuple[str, str]:
         try:
-            if "http" in question or "Wikipedia" in question:
-                web_context = self.search(question)
-                prompt = f"{SYSTEM_PROMPT}\n\n{web_context}\n\nQuestion: {question}"
-            elif files:
-                file_keys = list(files.keys())
-                for key in file_keys:
-                    if key.endswith(".mp3"):
-                        audio_txt = self.transcribe_audio(files[key])
-                        prompt = f"{SYSTEM_PROMPT}\n\n{audio_txt}\n\n{question}"
+            context = ""
+            if files:
+                for filename, filepath in files.items():
+                    if filename.endswith(".mp3"):
+                        context = self.transcribe_audio(filepath)
                         break
-                    elif key.endswith(".xlsx"):
-                        excel_result = self.handle_excel(files[key])
-                        return excel_result, excel_result
-                    else:
-                        prompt = f"{SYSTEM_PROMPT}\n\n{question}"
-            else:
-                prompt = f"{SYSTEM_PROMPT}\n\n{question}"
+                    elif filename.endswith(".xlsx"):
+                        excel_result = self.handle_excel(filepath)
+                        return excel_result.strip(), excel_result.strip()
+            elif "http" in question.lower() or "wikipedia" in question.lower():
+                context = self.search(question)
 
+            prompt = f"{SYSTEM_PROMPT}\n\n{context}\n\nQuestion: {question.strip()}"
             inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True).to(self.device)
             outputs = self.model.generate(
                 **inputs,
@@ -77,7 +72,7 @@ class GaiaAgent:
             )
             output_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
             final = output_text.strip()
-            return final, output_text
+            return final, final
         except Exception as e:
             return "ERROR", f"Agent failed: {e}"