Update agent.py

agent.py

@@ -1,40 +1,46 @@
+import mimetypes
 from transformers import pipeline
 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.qa_pipeline = pipeline("text2text-generation", model="google/flan-t5-base")
+        self.llm = pipeline("text-generation", model="mistralai/Mistral-7B-Instruct-v0.2", max_new_tokens=512, device="cpu")
+
+    def __call__(self, question: str, files: list = None):
+        trace = []
+        context = ""
 
-    def __call__(self, question: str):
-        trace = ""
+        if files:
+            for file in files:
+                mime, _ = mimetypes.guess_type(file.name)
+                if mime and mime.startswith("audio"):
+                    transcription = transcribe_audio(file.name)
+                    trace.append(f"Transcribed audio: {transcription}")
+                    context += f"\nTranscription: {transcription}"
+                elif mime and ("spreadsheet" in mime or file.name.endswith(".xlsx")):
+                    result = analyze_excel(file.name)
+                    trace.append(f"Excel analysis: {result}")
+                    context += f"\nSpreadsheet data: {result}"
 
-        # Handle audio
-        if question.lower().strip().endswith(('.mp3', '.wav')):
-            trace += "Audio detected. Running transcription...\n"
-            text = transcribe_audio(question.strip())
-            trace += f"Transcribed text: {text}\n"
-            answer = self.qa_pipeline(text, max_new_tokens=64)[0]['generated_text']
-            return answer.strip(), trace
+        if "http" in question or "Wikipedia" in question or "YouTube" in question or "search" in question.lower():
+            trace.append("Performing DuckDuckGo search...")
+            search_result = search_duckduckgo(question)
+            trace.append(f"Summary from search: {search_result}")
+            context += f"\nSearch Result: {search_result}"
 
-        # Handle Excel
-        if question.lower().strip().endswith(('.xls', '.xlsx')):
-            trace += "Excel detected. Running analysis...\n"
-            answer = analyze_excel(question.strip())
-            trace += f"Extracted value: {answer}\n"
-            return answer.strip(), trace
+        # Include the original question
+        prompt = f"""
+Answer the question based on the context below.
+Context: {context}
+Question: {question}
+Answer:
+"""
+        response = self.llm(prompt)[0]['generated_text'].split("Answer:")[-1].strip()
+        trace.append(response)
 
-        # Handle web search
-        if any(keyword in question.lower() for keyword in ["wikipedia", "video", "youtube", "article"]):
-            trace += "Performing DuckDuckGo search...\n"
-            summary = search_duckduckgo(question)
-            trace += f"Summary from search: {summary}\n"
-            answer = self.qa_pipeline(summary + "\n" + question, max_new_tokens=64)[0]['generated_text']
-            return answer.strip(), trace
+        return response, "\n".join(trace)
 
-        trace += "General question. Using local model...\n"
-        answer = self.qa_pipeline(question, max_new_tokens=64)[0]['generated_text']
-        return answer.strip(), trace
+