Update agent.py

agent.py

@@ -1,7 +1,12 @@
-# --- agent.py ---
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-from duckduckgo_search import DDGS
 import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from ddgs import DDGS
+import re
+import pandas as pd
+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.
@@ -9,52 +14,70 @@ Then return only the answer without any explanation or formatting.
 Do not say 'Final answer' or anything else. Just output the raw answer string.
 """
 
-def web_search(query: str, max_results: int = 3) -> list[str]:
-    results = []
-    try:
-        with DDGS() as ddgs:
-            for r in ddgs.text(query, max_results=max_results):
-                snippet = f"{r['title']}: {r['body']} (URL: {r['href']})"
-                results.append(snippet)
-    except Exception as e:
-        results.append(f"[Web search error: {e}]")
-    return results
-
-
 class GaiaAgent:
     def __init__(self, model_id="google/flan-t5-base"):
         self.tokenizer = AutoTokenizer.from_pretrained(model_id)
         self.model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.model.to(self.device)
+        self.transcriber = whisper.load_model("base")
+
+    def search(self, query: str) -> str:
+        try:
+            with DDGS() as ddgs:
+                results = list(ddgs.text(query, safesearch="off"))
+                if results:
+                    return results[0]['body']
+        except Exception as e:
+            return f"Search failed: {e}"
+        return ""
 
-    def __call__(self, question: str) -> tuple[str, str]:
+    def transcribe_audio(self, file_path: str) -> str:
         try:
-            # Heuristik: gör webbsök om frågan kräver externa fakta
-            search_required = any(keyword in question.lower() for keyword in [
-                "wikipedia", "who", "when", "where", "youtube", "mp3", "video", "article", "name", "code", "city", "award", "nasa"
-            ])
-
-            if search_required:
-                search_results = web_search(question)
-                context = "\n".join(search_results)
-                prompt = f"{SYSTEM_PROMPT}\n\nSearch context:\n{context}\n\nQuestion: {question}"
-                trace = f"Search used:\n{context}"
+            result = self.transcriber.transcribe(file_path)
+            return result['text']
+        except Exception as e:
+            return f"Audio transcription failed: {e}"
+
+    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}"
+
+    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}"
+                        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}"
-                trace = "Search not used."
 
             inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True).to(self.device)
             outputs = self.model.generate(
                 **inputs,
                 max_new_tokens=128,
                 do_sample=False,
+                temperature=0.0,
                 pad_token_id=self.tokenizer.pad_token_id
             )
             output_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
             final = output_text.strip()
-            return final, trace
-
+            return final, output_text
         except Exception as e:
             return "ERROR", f"Agent failed: {e}"