Update agent.py

agent.py

@@ -1,4 +1,7 @@
 import os
+import json
+import re
+from typing import Tuple, Dict, Any
 from transformers import pipeline
 from tools.asr_tool import transcribe_audio
 from tools.excel_tool import analyze_excel
@@ -9,43 +12,139 @@ class GaiaAgent:
         token = os.getenv("HUGGINGFACEHUB_API_TOKEN")
         if not token:
             raise ValueError("Missing HUGGINGFACEHUB_API_TOKEN environment variable.")
-
+        
+        # Använd en mer kapabel modell för bättre reasoning
         self.llm = pipeline(
             "text-generation",
             model="mistralai/Mistral-7B-Instruct-v0.2",
             use_auth_token=token,
             device="cpu",
-            max_new_tokens=512,
-            do_sample=False
+            max_new_tokens=1024,  # Öka för mer detaljerade svar
+            do_sample=False,
+            temperature=0.1,
+            return_full_text=False
         )
+        
+        # System prompt enligt GAIA:s instruktioner
+        self.system_prompt = """You are a general AI assistant. I will ask you a question. Report your thoughts, and finish your answer with the following template: FINAL ANSWER: [YOUR FINAL ANSWER]. YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings. If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise. If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise. If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string."""
+
+    def extract_final_answer(self, text: str) -> str:
+        """Extrahera det slutliga svaret från modellens output"""
+        # Leta efter FINAL ANSWER: mönster
+        final_answer_match = re.search(r'FINAL ANSWER:\s*(.+?)(?:\n|$)', text, re.IGNORECASE)
+        if final_answer_match:
+            return final_answer_match.group(1).strip()
+        
+        # Fallback: ta sista meningen om inget FINAL ANSWER hittas
+        sentences = text.strip().split('\n')
+        return sentences[-1].strip() if sentences else text.strip()
 
-    def __call__(self, question: str):
-        trace_log = ""
-        # Search tool
-        if "http" in question or "www." in question:
-            trace_log += "Detected URL or web reference. Performing search...\n"
-            result = search_duckduckgo(question)
-            trace_log += f"Search result: {result}\n"
-            return result, trace_log
+    def needs_tool(self, question: str) -> Tuple[str, bool]:
+        """Bestäm vilket verktyg som behövs baserat på frågan"""
+        question_lower = question.lower()
+        
+        # Kontrollera för audio-filer
+        if any(ext in question_lower for ext in ['.mp3', '.wav', '.m4a', '.flac']):
+            return 'audio', True
+        
+        # Kontrollera för Excel-filer
+        if any(ext in question_lower for ext in ['.xlsx', '.xls', '.csv']):
+            return 'excel', True
+        
+        # Kontrollera för web-sökning
+        if any(keyword in question_lower for keyword in ['search', 'find', 'lookup', 'http', 'www.', 'website']):
+            return 'search', True
+        
+        # Kontrollera för matematiska beräkningar
+        if any(keyword in question_lower for keyword in ['calculate', 'compute', 'sum', 'average', 'count']):
+            return 'math', True
+        
+        return 'llm', False
 
-        # Audio tool
-        if question.lower().endswith(".mp3") or question.lower().endswith(".wav"):
-            trace_log += "Detected audio file. Performing transcription...\n"
-            result = transcribe_audio(question)
-            trace_log += f"Transcription result: {result}\n"
-            return result, trace_log
+    def process_with_tools(self, question: str, tool_type: str) -> Tuple[str, str]:
+        """Bearbeta frågan med specifika verktyg"""
+        trace_log = f"Detected {tool_type} task. Processing...\n"
+        
+        try:
+            if tool_type == 'audio':
+                # Extrahera filnamn från frågan
+                audio_files = re.findall(r'\b[\w\-_]+\.(mp3|wav|m4a|flac)\b', question, re.IGNORECASE)
+                if audio_files:
+                    result = transcribe_audio(audio_files[0])
+                    trace_log += f"Audio transcription: {result}\n"
+                    return result, trace_log
+            
+            elif tool_type == 'excel':
+                # Extrahera filnamn från frågan
+                excel_files = re.findall(r'\b[\w\-_]+\.(xlsx|xls|csv)\b', question, re.IGNORECASE)
+                if excel_files:
+                    result = analyze_excel(excel_files[0])
+                    trace_log += f"Excel analysis: {result}\n"
+                    return result, trace_log
+            
+            elif tool_type == 'search':
+                # Extrahera sökfråga
+                search_query = question
+                result = search_duckduckgo(search_query)
+                trace_log += f"Search results: {result}\n"
+                return result, trace_log
+            
+        except Exception as e:
+            trace_log += f"Error using {tool_type} tool: {str(e)}\n"
+            return f"Error: {str(e)}", trace_log
+        
+        return "No valid input found for tool", trace_log
 
-        # Excel tool
-        if question.lower().endswith(".xlsx") or question.lower().endswith(".xls"):
-            trace_log += "Detected Excel file. Performing analysis...\n"
-            result = analyze_excel(question)
-            trace_log += f"Excel analysis result: {result}\n"
-            return result, trace_log
+    def reason_with_llm(self, question: str, context: str = "") -> Tuple[str, str]:
+        """Använd LLM för reasoning med kontext"""
+        trace_log = "Using LLM for reasoning...\n"
+        
+        # Bygg prompt med system instruktioner
+        if context:
+            prompt = f"{self.system_prompt}\n\nContext: {context}\n\nQuestion: {question}\n\nPlease analyze this step by step and provide your final answer."
+        else:
+            prompt = f"{self.system_prompt}\n\nQuestion: {question}\n\nPlease analyze this step by step and provide your final answer."
+        
+        try:
+            response = self.llm(prompt)[0]["generated_text"]
+            trace_log += f"LLM response: {response}\n"
+            return response, trace_log
+        except Exception as e:
+            trace_log += f"Error with LLM: {str(e)}\n"
+            return f"Error: {str(e)}", trace_log
 
-        # LLM fallback
-        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
+    def __call__(self, question: str) -> Tuple[str, str]:
+        """Huvudfunktion som bearbetar frågan"""
+        total_trace = f"Processing question: {question}\n"
+        
+        # Bestäm vilka verktyg som behövs
+        tool_type, needs_tool = self.needs_tool(question)
+        total_trace += f"Tool needed: {tool_type}\n"
+        
+        context = ""
+        if needs_tool and tool_type != 'llm':
+            # Använd verktyg för att samla kontext
+            tool_result, tool_trace = self.process_with_tools(question, tool_type)
+            total_trace += tool_trace
+            context = tool_result
+        
+        # Använd LLM för reasoning
+        llm_response, llm_trace = self.reason_with_llm(question, context)
+        total_trace += llm_trace
+        
+        # Extrahera slutligt svar
+        final_answer = self.extract_final_answer(llm_response)
+        total_trace += f"Final answer extracted: {final_answer}\n"
+        
+        return final_answer, total_trace
 
+    def format_for_submission(self, task_id: str, question: str) -> Dict[str, Any]:
+        """Formatera svar för GAIA-submission"""
+        answer, trace = self.__call__(question)
+        
+        return {
+            "task_id": task_id,
+            "model_answer": answer,
+            "reasoning_trace": trace
+        }