Update agent.py

agent.py

@@ -1,38 +1,50 @@
+import os
+import mimetypes
 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
 
 class GaiaAgent:
     def __init__(self):
-        self.llm = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.2", device="cpu")
+        token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
+        if not token:
+            raise ValueError("Missing HUGGINGFACEHUB_API_TOKEN environment variable.")
+
+        self.llm = pipeline(
+            "text-generation",
+            model="mistralai/Mistral-7B-Instruct-v0.2",
+            token=token,
+            device="cpu",
+            max_new_tokens=512,
+            do_sample=False
+        )
 
     def __call__(self, question: str):
-        trace = ""
-        question_lower = question.lower()
+        trace_log = ""
 
-        if "http" in question_lower or "wikipedia" in question_lower:
-            trace += "Detected URL or web reference. Performing search...\n"
+        if "http" in question or "www." in question:
+            trace_log += "Detected URL or web reference. Performing search...\n"
             result = search_duckduckgo(question)
-            trace += result
-            return result.strip(), trace
+            trace_log += f"Search result: {result}\n"
+            return result, trace_log
 
-        elif question_lower.endswith(".mp3"):
-            trace += "Detected audio reference. Transcribing audio...\n"
+        if question.lower().endswith(".mp3") or question.lower().endswith(".wav"):
+            trace_log += "Detected audio file. Performing transcription...\n"
             result = transcribe_audio(question)
-            trace += result
-            return result.strip(), trace
+            trace_log += f"Transcription result: {result}\n"
+            return result, trace_log
 
-        elif question_lower.endswith(".xlsx") or question_lower.endswith(".csv"):
-            trace += "Detected Excel file. Analyzing...\n"
+        if question.lower().endswith(".xlsx") or question.lower().endswith(".xls"):
+            trace_log += "Detected Excel file. Performing analysis...\n"
             result = analyze_excel(question)
-            trace += result
-            return result.strip(), trace
-
-        else:
-            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
+            trace_log += f"Excel analysis result: {result}\n"
+            return result, trace_log
+
+        trace_log += "General question. Using local model...\n"
+        response = self.llm(question)[0]["generated_text"]
+        trace_log += f"LLM response: {response}\n"
+
+        return response.strip(), trace_log