Update app.py

app.py

@@ -14,8 +14,10 @@ from typing import List, Dict, TypedDict, Annotated
 from tenacity import retry, stop_after_attempt, wait_exponential
 import io
 import sys
+import speech_recognition as sr
+import tempfile
 
-# Set recursion limit for Python and LangGraph
+# Set recursion limit for Python
 sys.setrecursionlimit(1500)
 
 # --- Setup Logging ---
@@ -46,6 +48,23 @@ class RateLimitedDuckDuckGoSearch(DuckDuckGoSearchRun):
             logging.warning(f"All retries failed for query '{query}'. Falling back to reasoning.")
             return "Web search failed; using reasoning."
 
+# --- Audio Transcription ---
+def transcribe_audio(file_content: bytes) -> str:
+    try:
+        recognizer = sr.Recognizer()
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as temp_file:
+            temp_file.write(file_content)
+            temp_file_path = temp_file.name
+        with sr.AudioFile(temp_file_path) as source:
+            audio = recognizer.record(source)
+        text = recognizer.recognize_google(audio)
+        os.unlink(temp_file_path)
+        logging.info(f"Transcribed audio: {text[:200]}...")
+        return text
+    except Exception as e:
+        logging.error(f"Audio transcription failed: {str(e)}")
+        return "ERROR: Audio transcription failed"
+
 # --- Answer Formatting ---
 def format_answer(answer: str, question: str) -> str:
     """Post-process answers to match GAIA's exact-match requirements."""
@@ -95,7 +114,7 @@ class AgentState(TypedDict):
     messages: Annotated[List, add_messages]
     question: str
     task_id: str
-    file_data: pd.DataFrame | None
+    file_data: pd.DataFrame | bytes | None
     search_result: str
     answer: str
     attempts: int
@@ -117,7 +136,7 @@ class GaiaAgent:
         Analyze the question to determine the required action:
 
         - List tasks (keywords: "list," "comma-separated," "categorize"): Action: reason, Input: none
-        - File tasks (keywords: "attached," "Excel," "spreadsheet", ".mp3"): Action: fetch_file, Input: none
+        - File tasks (keywords: "attached," "Excel," "spreadsheet", ".mp3", "image"): Action: fetch_file, Input: none
         - Search tasks (keywords: "who," "when," "historical," "competition", "wikipedia"): Action: search, Input: specific query
         - Number tasks (keywords: "how many," "count"): Action: search, Input: specific query
         - Fallback: Action: reason, Input: none
@@ -146,11 +165,11 @@ class GaiaAgent:
         Steps:
         1. Analyze question format (list, USD, name, number).
         2. Use data:
-           - File: Analyze DataFrame (e.g., sum sales where 'Category' != 'Drink').
+           - File: Analyze DataFrame, audio transcript, or image description.
            - Search: Extract names, numbers, codes.
         3. Reason:
            - Lists: Filter, alphabetize.
-           - Files: Sum columns, handle missing data.
+           - Files: Sum columns, extract text, or describe visuals.
            - Search: Parse for key info.
         4. Format:
            - Lists: 'item1, item2' (alphabetized).
@@ -162,7 +181,8 @@ class GaiaAgent:
         Specific instructions:
         - Mercedes Sosa studio albums (2000-2009): Misa Criolla (2000), Acústico (2002), Corazón Libre (2005), Cantora 1 (2009), Cantora 2 (2009). Answer: 5
         - Vegetable lists: Only non-botanical fruits (broccoli, celery, fresh basil, green beans, lettuce, sweet potatoes).
-        - .mp3 files: Return 'ERROR: Audio transcription not supported'.
+        - .mp3 files: Transcribe audio and extract relevant info.
+        - Image-based (e.g., chess): Describe image or return error if no image provided.
 
         Examples:
         - Q: List of vegetables...
@@ -183,17 +203,17 @@ class GaiaAgent:
         Validate the answer:
 
         - Check format (list, USD, name, number).
-        - Verify relevance.
+        - Verify relevance to question.
         - Return 'valid' or 'invalid: [reason]'.
         - Numbers: Ensure non-negative for counts.
-        - Allow reasonable answers even if not perfect to avoid excessive retries.
+        - Allow reasonable answers to avoid excessive retries.
 
         Question: {question}
         Answer: {answer}
         """)
         self.graph = self._build_graph()
 
-    async def fetch_file(self, task_id: str) -> pd.DataFrame | None:
+    async def fetch_file(self, task_id: str) -> pd.DataFrame | bytes | None:
         @retry(stop=stop_after_attempt(6), wait=wait_exponential(multiplier=1, min=2, max=5))
         def _fetch():
             file_url = f"{self.api_url}/files/{task_id}"
@@ -203,8 +223,14 @@ class GaiaAgent:
             content = io.BytesIO(response.content)
             if file_url.endswith('.xlsx') or file_url.endswith('.xls'):
                 df = pd.read_excel(content)
-                logging.info(f"File fetched. Shape: {df.shape}, Columns: {list(df.columns)}")
+                logging.info(f"Excel file fetched. Shape: {df.shape}, Columns: {list(df.columns)}")
                 return df
+            elif file_url.endswith('.mp3'):
+                logging.info("Audio file fetched.")
+                return response.content
+            elif file_url.endswith('.png') or file_url.endswith('.jpg'):
+                logging.info("Image file fetched.")
+                return response.content
             logging.warning(f"Unsupported file type for task {task_id}")
             return None
 
@@ -216,7 +242,7 @@ class GaiaAgent:
 
     def _build_graph(self):
         workflow = StateGraph(AgentState)
-        # Set higher recursion limit for LangGraph
+        # Explicitly set recursion limit
         workflow.config = {"recursion_limit": 50}
 
         def classify(state: AgentState) -> AgentState:
@@ -233,7 +259,7 @@ class GaiaAgent:
                     task_id=state["task_id"],
                     attempts=state["attempts"]
                 ))
-                response_content = str(response.content)
+                response_content = str(response.content).strip()
                 action = response_content.split("Action: ")[1].split("\n")[0].strip() if "Action: " in response_content else "reason"
                 action_input = response_content.split("Input: ")[1].strip() if "Input: " in response_content else "none"
             except Exception as e:
@@ -247,7 +273,7 @@ class GaiaAgent:
 
         async def fetch_file_node(state: AgentState) -> AgentState:
             logging.info(f"Fetching file for task {state['task_id']}")
-            if state["file_data"] is None and state["attempts"] < 3:
+            if state["file_data"] is None and state["attempts"] < 2:
                 try:
                     state["file_data"] = await asyncio.wait_for(self.fetch_file(state["task_id"]), timeout=30)
                     state["messages"].append(f"Fetched file: {state['file_data'] is not None}")
@@ -260,7 +286,7 @@ class GaiaAgent:
 
         async def search_node(state: AgentState) -> AgentState:
             logging.info(f"Searching for task {state['task_id']}")
-            if state["attempts"] < 3:
+            if state["attempts"] < 2:
                 action_input = state["action_input"]
                 if action_input != "none":
                     try:
@@ -284,7 +310,14 @@ class GaiaAgent:
 
         def generate_answer(state: AgentState) -> AgentState:
             logging.info(f"Generating answer for task {state['task_id']}")
-            file_data_str = state["file_data"].to_dict() if state["file_data"] is not None else "none"
+            file_data_str = "none"
+            if isinstance(state["file_data"], pd.DataFrame):
+                file_data_str = state["file_data"].to_dict()
+            elif isinstance(state["file_data"], bytes):
+                if "mp3" in state["question"].lower():
+                    file_data_str = transcribe_audio(state["file_data"])
+                else:
+                    file_data_str = "Image data; cannot process directly."
             try:
                 response = self.llm.invoke(self.answer_prompt.format(
                     question=state["question"],
@@ -311,9 +344,8 @@ class GaiaAgent:
                 logging.error(f"Critique error: {str(e)}")
                 critique = "valid"  # Avoid looping on critique failure
             state["messages"].append(f"Critique: {critique}")
-            # Stricter exit condition to prevent infinite loops
             state["attempts"] += 1
-            if "invalid" in critique.lower() and state["attempts"] < 2:
+            if "invalid" in critique.lower() and state["attempts"] < 1:  # Reduced retries to 1
                 state["action"] = "reason"
                 state["action_input"] = "none"
             else:
@@ -327,7 +359,7 @@ class GaiaAgent:
         workflow.add_node("critique", critique_answer)
 
         def route_action(state: AgentState) -> str:
-            if state["attempts"] >= 3 or state["action"] is None:
+            if state["attempts"] >= 2 or state["action"] is None:
                 return END
             return state["action"] if state["action"] in ["fetch_file", "search", "reason"] else "generate_answer"