Robustify agent: better context passing, error handling, logging, prompt engineering, and dependencies

app.py

@@ -1,4 +1,4 @@
-import os
+chess - screenshote - screenshote - screenshote - screenshote - screenshotimport os
 import gradio as gr
 import requests
 import inspect
@@ -25,6 +25,7 @@ import cv2
 import torch
 from bs4 import BeautifulSoup
 import openai
+import magic  # for robust file type detection
 
 logging.basicConfig(filename='gaia_agent.log', level=logging.INFO, format='%(asctime)s %(levelname)s:%(message)s')
 logger = logging.getLogger(__name__)
@@ -38,7 +39,7 @@ def llama3_chat(prompt):
             messages=[{"role": "user", "content": prompt}],
         )
         return completion.choices[0].message.content
-    except Exception as e:
+        except Exception as e:
         logging.error(f"llama3_chat error: {e}")
         return f"LLM error: {e}"
 
@@ -243,116 +244,179 @@ TOOL_REGISTRY = {
     "gpt4_chat": gpt4_chat,
 }
 
+# --- Utility: Robust file type detection ---
+def detect_file_type_magic(file_name):
+    try:
+        mime = magic.Magic(mime=True)
+        filetype = mime.from_file(file_name)
+        if 'audio' in filetype:
+            return 'audio'
+        elif 'image' in filetype:
+            return 'image'
+        elif 'python' in filetype or file_name.endswith('.py'):
+            return 'code'
+        elif 'spreadsheet' in filetype or file_name.endswith('.xlsx'):
+            return 'excel'
+        elif 'csv' in filetype or file_name.endswith('.csv'):
+            return 'csv'
+        elif 'json' in filetype or file_name.endswith('.json'):
+            return 'json'
+        elif 'text' in filetype or file_name.endswith(('.txt', '.md')):
+            return 'text'
+        else:
+            return 'unknown'
+    except Exception as e:
+        logger.error(f"magic file type detection error: {e}")
+        return 'unknown'
+
+# --- Improved prompt template for LLMs ---
+def build_prompt(context, question):
+    return f"""
+Context:
+{context}
+
+Question:
+{question}
+
+Answer:
+"""
+
+# --- Refactored ModularGAIAAgent ---
 class ModularGAIAAgent:
-    def __init__(self, api_url=DEFAULT_API_URL, tool_registry=TOOL_REGISTRY):
+    def __init__(self, api_url=DEFAULT_API_URL, tool_registry=None):
         self.api_url = api_url
-        self.tools = tool_registry
+        self.tools = tool_registry or TOOL_REGISTRY
         self.reasoning_trace = []
         self.file_cache = set(os.listdir('.'))
 
     def fetch_questions(self, from_api=True, questions_path="Hugging Face Questions"):
-        if from_api:
-            r = requests.get(f"{self.api_url}/questions")
-            r.raise_for_status()
-            return r.json()
-        else:
-            with open(questions_path) as f:
-                data = f.read()
-            start = data.find("[")
-            end = data.rfind("]") + 1
-            questions = json.loads(data[start:end])
-            return questions
+        """Fetch questions from API or local file."""
+        try:
+            if from_api:
+                r = requests.get(f"{self.api_url}/questions")
+                r.raise_for_status()
+                return r.json()
+            else:
+                with open(questions_path) as f:
+                    data = f.read()
+                start = data.find("[")
+                end = data.rfind("]") + 1
+                questions = json.loads(data[start:end])
+                return questions
+        except Exception as e:
+            logger.error(f"fetch_questions error: {e}")
+            return []
 
     def download_file(self, file_id, file_name=None):
-        if not file_name:
-            file_name = file_id
-        if file_name in self.file_cache:
-            return file_name
-        url = f"{self.api_url}/files/{file_id}"
-        r = requests.get(url)
-        if r.status_code == 200:
-            with open(file_name, "wb") as f:
-                f.write(r.content)
-            self.file_cache.add(file_name)
-            return file_name
-        else:
-            self.reasoning_trace.append(f"Failed to download file {file_id} (status {r.status_code})")
+        """Download file if not present locally."""
+        try:
+            if not file_name:
+                file_name = file_id
+            if file_name in self.file_cache:
+                return file_name
+            url = f"{self.api_url}/files/{file_id}"
+            r = requests.get(url)
+            if r.status_code == 200:
+                with open(file_name, "wb") as f:
+                    f.write(r.content)
+                self.file_cache.add(file_name)
+                return file_name
+            else:
+                self.reasoning_trace.append(f"Failed to download file {file_id} (status {r.status_code})")
+                logger.error(f"Failed to download file {file_id} (status {r.status_code})")
+                return None
+        except Exception as e:
+            logger.error(f"download_file error: {e}")
+            self.reasoning_trace.append(f"Download error: {e}")
             return None
 
     def detect_file_type(self, file_name):
-        ext = os.path.splitext(file_name)[-1].lower()
-        if ext in ['.mp3', '.wav', '.flac']:
-            return 'audio'
-        elif ext in ['.png', '.jpg', '.jpeg', '.bmp']:
-            return 'image'
-        elif ext in ['.py']:
-            return 'code'
-        elif ext in ['.xlsx']:
-            return 'excel'
-        elif ext in ['.csv']:
-            return 'csv'
-        elif ext in ['.json']:
-            return 'json'
-        elif ext in ['.txt', '.md']:
-            return 'text'
-        else:
-            return 'unknown'
+        """Detect file type using magic and extension as fallback."""
+        file_type = detect_file_type_magic(file_name)
+        if file_type == 'unknown':
+            ext = os.path.splitext(file_name)[-1].lower()
+            if ext in ['.mp3', '.wav', '.flac']:
+                return 'audio'
+            elif ext in ['.png', '.jpg', '.jpeg', '.bmp']:
+                return 'image'
+            elif ext in ['.py']:
+                return 'code'
+            elif ext in ['.xlsx']:
+                return 'excel'
+            elif ext in ['.csv']:
+                return 'csv'
+            elif ext in ['.json']:
+                return 'json'
+            elif ext in ['.txt', '.md']:
+                return 'text'
+            else:
+                return 'unknown'
+        return file_type
 
     def analyze_file(self, file_name, file_type):
-        if file_type == 'audio':
-            transcript = self.tools['asr_transcribe'](file_name)
-            self.reasoning_trace.append(f"Transcribed audio: {transcript[:100]}...")
-            return transcript
-        elif file_type == 'image':
-            caption = self.tools['image_caption'](file_name)
-            self.reasoning_trace.append(f"Image caption: {caption}")
-            return caption
-        elif file_type == 'code':
-            result = self.tools['code_analysis'](file_name)
-            self.reasoning_trace.append(f"Code analysis result: {result}")
-            return result
-        elif file_type == 'excel':
-            wb = openpyxl.load_workbook(file_name)
-            ws = wb.active
-            data = list(ws.values)
-            headers = data[0]
-            table = [dict(zip(headers, row)) for row in data[1:]]
-            self.reasoning_trace.append(f"Excel table loaded: {table[:2]}...")
-            return table
-        elif file_type == 'csv':
-            df = pd.read_csv(file_name)
-            table = df.to_dict(orient='records')
-            self.reasoning_trace.append(f"CSV table loaded: {table[:2]}...")
-            return table
-        elif file_type == 'json':
-            with open(file_name) as f:
-                data = json.load(f)
-            self.reasoning_trace.append(f"JSON loaded: {str(data)[:100]}...")
-            return data
-        elif file_type == 'text':
-            with open(file_name) as f:
-                text = f.read()
-            self.reasoning_trace.append(f"Text loaded: {text[:100]}...")
-            return text
-        else:
-            self.reasoning_trace.append(f"Unknown file type: {file_name}")
+        """Analyze file and return context for the question."""
+        try:
+            if file_type == 'audio':
+                transcript = self.tools['asr_transcribe'](file_name)
+                self.reasoning_trace.append(f"Transcribed audio: {transcript[:100]}...")
+                return transcript
+            elif file_type == 'image':
+                caption = self.tools['image_caption'](file_name)
+                self.reasoning_trace.append(f"Image caption: {caption}")
+                return caption
+            elif file_type == 'code':
+                result = self.tools['code_analysis'](file_name)
+                self.reasoning_trace.append(f"Code analysis result: {result}")
+                return result
+            elif file_type == 'excel':
+                wb = openpyxl.load_workbook(file_name)
+                ws = wb.active
+                data = list(ws.values)
+                headers = data[0]
+                table = [dict(zip(headers, row)) for row in data[1:]]
+                self.reasoning_trace.append(f"Excel table loaded: {table[:2]}...")
+                return table
+            elif file_type == 'csv':
+                df = pd.read_csv(file_name)
+                table = df.to_dict(orient='records')
+                self.reasoning_trace.append(f"CSV table loaded: {table[:2]}...")
+                return table
+            elif file_type == 'json':
+                with open(file_name) as f:
+                    data = json.load(f)
+                self.reasoning_trace.append(f"JSON loaded: {str(data)[:100]}...")
+                return data
+            elif file_type == 'text':
+                with open(file_name) as f:
+                    text = f.read()
+                self.reasoning_trace.append(f"Text loaded: {text[:100]}...")
+                return text
+            else:
+                self.reasoning_trace.append(f"Unknown file type: {file_name}")
+                logger.warning(f"Unknown file type: {file_name}")
+                return None
+        except Exception as e:
+            logger.error(f"analyze_file error: {e}")
+            self.reasoning_trace.append(f"Analyze file error: {e}")
             return None
 
     def smart_tool_select(self, question, file_type=None):
         """Select the best tool(s) for the question, optionally using GPT-4.1 for planning."""
-        # Use GPT-4.1 to suggest a tool if available
         api_key = os.environ.get("OPENAI_API_KEY", "")
-        if api_key:
-            plan_prompt = f"""
+        try:
+            if api_key:
+                plan_prompt = f"""
 You are an expert AI agent. Given the following question and file type, suggest the best tool(s) to use from this list: {list(self.tools.keys())}.
 Question: {question}
 File type: {file_type}
 Respond with a comma-separated list of tool names only, in order of use. If unsure, start with web_search_duckduckgo.
 """
-            plan = gpt4_chat(plan_prompt, api_key=api_key)
-            tool_names = [t.strip() for t in plan.split(',') if t.strip() in self.tools]
-            if tool_names:
-                return tool_names
+                plan = gpt4_chat(plan_prompt, api_key=api_key)
+                tool_names = [t.strip() for t in plan.split(',') if t.strip() in self.tools]
+                if tool_names:
+                    return tool_names
+        except Exception as e:
+            logger.error(f"smart_tool_select planning error: {e}")
         # Fallback: heuristic
         if file_type == 'audio':
             return ['asr_transcribe']
@@ -370,6 +434,7 @@ Respond with a comma-separated list of tool names only, in order of use. If unsu
             return ['llama3_chat']
 
     def answer_question(self, question_obj):
+        """Answer a question using the best tool(s) and context."""
         self.reasoning_trace = []
         q = question_obj["question"]
         file_name = question_obj.get("file_name", "")
@@ -384,31 +449,40 @@ Respond with a comma-separated list of tool names only, in order of use. If unsu
         # Smart tool selection
         tool_names = self.smart_tool_select(q, file_type)
         answer = None
-        context = None
+        context = file_content
         for tool_name in tool_names:
             tool = self.tools[tool_name]
-            if tool_name == 'web_search_duckduckgo':
-                context = tool(q)
-                # Use LLM to synthesize answer from snippets
-                answer = llama3_chat(f"Answer the following question using ONLY the information below.\nQuestion: {q}\nSnippets:\n{context}\nAnswer:")
-            elif tool_name == 'gpt4_chat':
-                answer = tool(q)
-            elif tool_name == 'table_qa' and file_content:
-                answer = tool(q, file_content)
-            elif tool_name in ['asr_transcribe', 'image_caption', 'code_analysis'] and file_content:
-                answer = tool(file_name)
-            elif tool_name == 'youtube_video_qa':
-                answer = tool(q, q)
-            else:
-                answer = tool(q)
-            if answer:
-                break
+            try:
+                logger.info(f"Using tool: {tool_name} | Question: {q} | Context: {str(context)[:200]}")
+                if tool_name == 'web_search_duckduckgo':
+                    context = tool(q)
+                    answer = llama3_chat(build_prompt(context, q))
+                elif tool_name == 'gpt4_chat':
+                    answer = tool(build_prompt(context, q))
+                elif tool_name == 'table_qa' and file_content:
+                    answer = tool(q, file_content)
+                elif tool_name in ['asr_transcribe', 'image_caption', 'code_analysis'] and file_content:
+                    answer = tool(file_name)
+                elif tool_name == 'youtube_video_qa':
+                    answer = tool(q, q)
+                else:
+                    # Always pass context if available
+                    if context:
+                        answer = llama3_chat(build_prompt(context, q))
+                    else:
+                        answer = tool(q)
+                if answer:
+                    break
+            except Exception as e:
+                logger.error(f"Tool {tool_name} error: {e}")
+                self.reasoning_trace.append(f"Tool {tool_name} error: {e}")
+                continue
         self.reasoning_trace.append(f"Tools used: {tool_names}")
         self.reasoning_trace.append(f"Final answer: {answer}")
         return self.format_answer(answer), self.reasoning_trace
 
     def format_answer(self, answer):
-        # Strict GAIA: only the answer, no extra text, no prefix
+        """Strict GAIA: only the answer, no extra text, no prefix."""
         if isinstance(answer, str):
             return answer.strip().split('\n')[0]
         return str(answer)

requirements.txt

@@ -11,4 +11,7 @@ opencv-python
 beautifulsoup4
 yt-dlp
 ultralytics
-openai
+openai
+torchaudio
+ffmpeg-python
+python-magic