Update agent.py

agent.py

@@ -1,92 +1,40 @@
-import torch
-from transformers import (
-    AutoTokenizer,
-    AutoModelForSeq2SeqLM,
-    pipeline,
-    AutoProcessor,
-    AutoModelForSpeechSeq2Seq
-)
-from duckduckgo_search import DDGS
-import pandas as pd
-import os
-
-SYSTEM_PROMPT = """
-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.
-"""
+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, 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)
-
-        # Whisper via HF
-        self.asr_model_id = "openai/whisper-small"
-        self.asr_processor = AutoProcessor.from_pretrained(self.asr_model_id)
-        self.asr_model = AutoModelForSpeechSeq2Seq.from_pretrained(self.asr_model_id).to(self.device)
-        self.pipe = pipeline(
-            "automatic-speech-recognition",
-            model=self.asr_model,
-            tokenizer=self.asr_processor.tokenizer,
-            feature_extractor=self.asr_processor.feature_extractor,
-            return_timestamps=False,
-            device=0 if torch.cuda.is_available() else -1
-        )
-
-    def search(self, query: str) -> str:
-        try:
-            with DDGS() as ddgs:
-                results = ddgs.text("your query", max_results=1)
-                if results:
-                    return results[0]['body']
-        except Exception:
-            return ""
-        return ""
-
-    def transcribe_audio(self, file_path: str) -> str:
-        try:
-            result = self.pipe(file_path)
-            return result['text']
-        except Exception:
-            return ""
-
-    def handle_excel(self, file_path: str) -> str:
-        try:
-            df = pd.read_excel(file_path)
-            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:
-            context = ""
-            if files:
-                for filename, filepath in files.items():
-                    if filename.endswith(".mp3") or filename.endswith(".wav"):
-                        context = self.transcribe_audio(filepath)
-                        break
-                    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,
-                max_new_tokens=128,
-                do_sample=False,
-                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, final
-        except Exception as e:
-            return "ERROR", f"Agent failed: {e}"
-            
+    def __init__(self):
+        print("Loading model...")
+        self.qa_pipeline = pipeline("text2text-generation", model="google/flan-t5-base")
+
+    def __call__(self, question: str):
+        trace = ""
+
+        # 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
+
+        # 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
+
+        # 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
+
+        trace += "General question. Using local model...\n"
+        answer = self.qa_pipeline(question, max_new_tokens=64)[0]['generated_text']
+        return answer.strip(), trace