Update agent.py

agent.py

@@ -1,43 +1,38 @@
 from transformers import pipeline
+from tools.search_tool import search_duckduckgo
 from tools.asr_tool import transcribe_audio
 from tools.excel_tool import analyze_excel
-from tools.search_tool import search_duckduckgo
-import mimetypes
-
 
 class GaiaAgent:
     def __init__(self):
-        print("Loading model...")
-        self.model = pipeline(
-            "text2text-generation",
-            model="MBZUAI/LaMini-Flan-T5-783M",
-            tokenizer="MBZUAI/LaMini-Flan-T5-783M"
-        )
-        print("Model loaded.")
+        self.llm = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.2", device="cpu")
 
-    def __call__(self, query):
+    def __call__(self, question: str):
         trace = ""
-        final_answer = ""
-
-        # Försök identifiera om det är en filreferens
-        if isinstance(query, str) and (query.endswith(".mp3") or query.endswith(".wav")):
-            trace = "Detected audio file. Transcribing..."
-            final_answer = transcribe_audio(query)
-
-        elif isinstance(query, str) and (query.endswith(".xls") or query.endswith(".xlsx")):
-            trace = "Detected Excel file. Analyzing..."
-            final_answer = analyze_excel(query)
-
-        elif "http" in query:
-            trace = "Detected URL or web reference. Performing search..."
-            final_answer = search_duckduckgo(query)
+        question_lower = question.lower()
+
+        if "http" in question_lower or "wikipedia" in question_lower:
+            trace += "Detected URL or web reference. Performing search...\n"
+            result = search_duckduckgo(question)
+            trace += result
+            return result.strip(), trace
+
+        elif question_lower.endswith(".mp3"):
+            trace += "Detected audio reference. Transcribing audio...\n"
+            result = transcribe_audio(question)
+            trace += result
+            return result.strip(), trace
+
+        elif question_lower.endswith(".xlsx") or question_lower.endswith(".csv"):
+            trace += "Detected Excel file. Analyzing...\n"
+            result = analyze_excel(question)
+            trace += result
+            return result.strip(), trace
 
         else:
-            trace = "General question. Using local model..."
-            output = self.model(query, max_new_tokens=128)
-            final_answer = output[0]["generated_text"].strip()
-
-        return final_answer, trace
-
+            trace += "General question. Using local model...\n"
+            output = self.llm(question, max_new_tokens=256, do_sample=False)[0]["generated_text"]
+            trace += output
+            return output.strip(), trace