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 from tools.search_tool import search_duckduckgo class GaiaAgent: def __init__(self): 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=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 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 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 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 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 }