Update agent.py

agent.py

@@ -1,276 +1,173 @@
 import os
-import io
+import asyncio
 import re
-import requests
-import tempfile
-import subprocess
-from openai import OpenAI
-from duckduckgo_search import DDGS
-from PIL import Image
-import pytesseract
-import openpyxl
-
-try:
-    import whisper
-except ImportError:
-    whisper = None
-
-try:
-    import pdfplumber
-except ImportError:
-    pdfplumber = None
-
-AGENT_API_URL = "https://agents-course-unit4-scoring.hf.space"
-
-def safe_strip(text):
-    if not text:
-        return ""
-    if isinstance(text, bytes):
-        text = text.decode(errors="ignore")
-    return str(text).replace("\r", "").strip()
-
-def run_web_search(query, max_results=3):
+from typing import Any
+
+from llama_index.llms.openai import OpenAI
+from llama_index.core.agent.react import ReActAgent
+from llama_index.core.agent.workflow import AgentWorkflow
+from llama_index.core.tools import FunctionTool, ToolMetadata
+
+# Tool: DuckDuckGo Web Search
+from llama_index.tools.duckduckgo import DuckDuckGoSearchTool
+
+# Tool: Python code eval (for simple code/number/output questions)
+def eval_python_code(code: str) -> str:
+    """
+    Evaluate simple Python code and return result as string.
+    Use for 'What is the output of this code?' or math.
+    """
     try:
-        ddgs = DDGS()
-        results = ddgs.text(query)
-        bodies = []
-        for i, r in enumerate(results):
-            if i >= max_results:
-                break
-            if r.get('body'):
-                bodies.append(r['body'])
-            elif r.get('title'):
-                bodies.append(r['title'])
-        return "\n".join(bodies)
-    except Exception:
-        return ""
-
-def fetch_file(task_id):
-    url = f"{AGENT_API_URL}/files/{task_id}"
-    try:
-        resp = requests.get(url, timeout=30)
-        resp.raise_for_status()
-        content_type = resp.headers.get("Content-Type", "")
-        return resp.content, content_type
-    except Exception:
-        return None, None
-
-def ocr_image(img_bytes):
-    try:
-        img = Image.open(io.BytesIO(img_bytes))
-        return safe_strip(pytesseract.image_to_string(img))
-    except Exception:
+        # Only eval expressions (NOT exec for safety!)
+        return str(eval(code, {"__builtins__": {}}))
+    except Exception as e:
+        return f"ERROR: {e}"
+
+# Tool: Strict output formatting
+def format_gaia_answer(answer: str, question: str = "") -> str:
+    """Postprocess: GAIA strict answer format enforcement."""
+    if not answer:
         return ""
+    # Remove quotes/brackets/periods, apologies, "Final Answer:"
+    answer = re.sub(r'(?i)final answer:?\s*', '', answer).strip()
+    answer = re.sub(r'(?i)i(\'?m| cannot| can\'t| unable to| apologize| not available|process the file).*', '', answer).strip()
+    if answer.startswith('"') and answer.endswith('"'): answer = answer[1:-1]
+    if answer.startswith('[') and answer.endswith(']'): answer = answer[1:-1]
+    if not re.match(r'^[A-Za-z]+\.$', answer): answer = re.sub(r'\.$', '', answer)
+    # Numeric
+    if re.search(r'how many|number of|at bats|total sales|albums|output.*python|highest number', question, re.I):
+        num = re.search(r'(\$?\d[\d,\.]*)', answer)
+        if num: return num.group(1).replace(',', '')
+    # Surname/first name/code/city
+    if 'first name' in question: return answer.split()[0]
+    if 'surname' in question: return answer.split()[-1]
+    if 'city' in question: return answer.split()[0]
+    if re.search(r'IOC country code|award number|NASA', question, re.I):
+        code = re.search(r'[A-Z0-9]{3,}', answer)
+        if code: return code.group(0)
+    if re.search(r'list|comma.*separated|page numbers', question, re.I):
+        items = [x.strip('",.').lower() for x in re.split(r'[,\n]', answer) if x.strip()]
+        if 'page numbers' in question:
+            nums = [int(x) for x in re.findall(r'\d+', answer)]
+            return ', '.join(str(n) for n in sorted(nums))
+        if 'ingredient' in question or 'vegetable' in question:
+            merged = []
+            skip = False
+            for i, item in enumerate(items):
+                if skip: skip = False; continue
+                if i+1 < len(items) and item in ['sweet', 'green', 'lemon', 'ripe', 'whole', 'fresh']:
+                    merged.append(f"{item} {items[i+1]}")
+                    skip = True
+                else: merged.append(item)
+            merged = sorted(set(merged))
+            return ', '.join(merged)
+        return ', '.join(items)
+    return answer.strip().rstrip('.').strip()
 
-def read_excel(file_bytes):
+# Tool: OCR for images (incl. chessboards/screenshots)
+def ocr_image(file_path: str) -> str:
+    """Extract text from image file."""
+    from PIL import Image
+    import pytesseract
     try:
-        wb = openpyxl.load_workbook(io.BytesIO(file_bytes), data_only=True)
-        sheet = wb.active
-        rows = list(sheet.iter_rows(values_only=True))
-        text = "\n".join(["\t".join(str(cell) if cell is not None else "" for cell in row) for row in rows])
-        return safe_strip(text)
-    except Exception:
-        return ""
-
-def read_pdf(file_bytes):
-    if not pdfplumber:
-        return ""
+        img = Image.open(file_path)
+        return pytesseract.image_to_string(img)
+    except Exception as e:
+        return f"ERROR: {e}"
+
+# Tool: Audio transcription (Whisper)
+def transcribe_audio(file_path: str) -> str:
+    """Transcribe audio file with Whisper."""
     try:
-        with pdfplumber.open(io.BytesIO(file_bytes)) as pdf:
-            return safe_strip("\n".join(page.extract_text() or "" for page in pdf.pages))
-    except Exception:
-        return ""
-
-def transcribe_audio(audio_bytes):
-    if not whisper:
-        return ""
-    try:
-        with tempfile.NamedTemporaryFile(suffix=".mp3", delete=True) as tmpfile:
-            tmpfile.write(audio_bytes)
-            tmpfile.flush()
-            model = whisper.load_model("base")
-            result = model.transcribe(tmpfile.name)
-            return safe_strip(result.get("text", ""))
-    except Exception:
-        return ""
-
-def transcribe_youtube_audio(youtube_url):
-    if not whisper:
-        return ""
+        import whisper
+        model = whisper.load_model("base")
+        result = model.transcribe(file_path)
+        return result.get("text", "")
+    except Exception as e:
+        return f"ERROR: {e}"
+
+# Tool: YouTube video transcription
+def transcribe_youtube(url: str) -> str:
+    """Download and transcribe a YouTube video (audio only)."""
+    import tempfile, os
     try:
+        import whisper
+        import yt_dlp
         with tempfile.TemporaryDirectory() as tmpdir:
-            audio_path = os.path.join(tmpdir, "audio.mp3")
-            cmd = [
-                "yt-dlp", "-f", "bestaudio[ext=m4a]/bestaudio/best",
-                "--extract-audio", "--audio-format", "mp3",
-                "-o", audio_path, youtube_url
-            ]
-            subprocess.run(cmd, check=True, capture_output=True)
+            ydl_opts = {'format': 'bestaudio/best', 'outtmpl': os.path.join(tmpdir, 'audio.%(ext)s')}
+            with yt_dlp.YoutubeDL(ydl_opts) as ydl:
+                ydl.download([url])
+            audio_path = [os.path.join(tmpdir, f) for f in os.listdir(tmpdir) if f.startswith("audio")][0]
             model = whisper.load_model("base")
             result = model.transcribe(audio_path)
-            return safe_strip(result.get("text", ""))
-    except Exception:
-        return ""
-
-def extract_file_text(file_bytes, content_type, task_id=""):
-    if "image" in content_type:
-        return ocr_image(file_bytes)
-    if "spreadsheet" in content_type or "excel" in content_type or task_id.endswith(".xlsx"):
-        return read_excel(file_bytes)
-    if "pdf" in content_type or task_id.endswith(".pdf"):
-        return read_pdf(file_bytes)
-    if "audio" in content_type or task_id.endswith(".mp3") or task_id.endswith(".wav"):
-        return transcribe_audio(file_bytes)
-    if "text" in content_type or "csv" in content_type or "json" in content_type or task_id.endswith(".csv") or task_id.endswith(".json") or task_id.endswith(".txt"):
-        return safe_strip(file_bytes[:10000])
-    return ""
-
-def guess_youtube_link(question):
-    matches = re.findall(r"(https?://[^\s]+)", question)
-    for url in matches:
-        if "youtube.com" in url or "youtu.be" in url:
-            return url
-    return None
-
-def format_gaia_answer(answer, question=None):
-    if not answer or not isinstance(answer, str):
-        return ""
-    answer = re.sub(r"(?i)i'?m sorry[,\.]?|i cannot|i can't|unable to|please provide.*|information not available|I can't assist.*|I'm unable.*|process the file directly", "", answer)
-    answer = re.sub(r'(?i)final answer:?\s*', '', answer).strip()
-    if answer.startswith('"') and answer.endswith('"'):
-        answer = answer[1:-1]
-    if answer.startswith('[') and answer.endswith(']'):
-        answer = answer[1:-1]
-    if not re.match(r'^[A-Za-z]+\.$', answer):
-        answer = re.sub(r'\.$', '', answer)
-
-    if question:
-        # Pure number answers
-        if re.search(r'how many|number of|at bats|total sales|albums|output.*python|highest number', question, re.I):
-            num_match = re.search(r'(\$?\d[\d,\.]*)', answer)
-            if num_match:
-                return num_match.group(1).replace(',', '')
-
-        if re.search(r'first name', question, re.I):
-            return answer.strip().split()[0]
-        if re.search(r'surname', question, re.I):
-            return answer.strip().split()[-1]
-        if re.search(r'city', question, re.I):
-            return answer.strip().split()[0]
-        if re.search(r'IOC country code|award number|NASA', question, re.I):
-            code_match = re.search(r'[A-Z0-9]{3,}', answer)
-            if code_match:
-                return code_match.group(0)
-        if 'algebraic notation' in question or 'chess' in question:
-            move_match = re.search(r'[A-Za-z0-9]+[#\+]?$', answer)
-            if move_match:
-                return move_match.group(0)
-        if "what does teal'c say" in question.lower():
-            qmatch = re.search(r'"(Indeed\.)"', answer)
-            if qmatch:
-                return qmatch.group(1)
-            if "Indeed." in answer:
-                return "Indeed."
-            return answer
-        if re.search(r'list|comma.*separated|page numbers', question, re.I):
-            items = [x.strip('",.').lower() for x in re.split(r'[,\n]', answer) if x.strip()]
-            items = [item for item in items if item and not re.match(r'(and|or|to|with|for|a|the)$', item)]
-            if 'page numbers' in question:
-                nums = [int(x) for x in re.findall(r'\d+', answer)]
-                return ', '.join(str(n) for n in sorted(nums))
-            if 'ingredient' in question or 'vegetable' in question or 'grocery' in question:
-                merged = []
-                skip = False
-                for i, item in enumerate(items):
-                    if skip:
-                        skip = False
-                        continue
-                    if i+1 < len(items) and item in ['sweet', 'green', 'lemon', 'ripe', 'whole', 'fresh']:
-                        merged.append(f"{item} {items[i+1]}")
-                        skip = True
-                    else:
-                        merged.append(item)
-                merged = sorted(set(merged))
-                return ', '.join(merged)
-            return ', '.join(items)
-        if re.search(r'pitcher.*before.*after', question, re.I):
-            names = re.findall(r'\b[A-Z][a-z]+', answer)
-            return ', '.join(names[:2])
-
-    return answer.strip().rstrip('.').strip()
-
-class GaiaAgent:
-    def __init__(self):
-        self.client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
-        self.instructions = (
-            "You are a top-tier research assistant for the GAIA benchmark. "
-            "You analyze documents, reason step by step, and always provide a single, concise, and correct answer. "
-            "If a file is provided, extract all relevant information. Use only information from the question and file. "
-            "If the question refers to a video/audio file or YouTube link, always try to transcribe it. "
-            "If you need additional facts, summarize web search results provided. "
-            "Never apologize, never say you are unable, never output placeholders. "
-            "Always output the answer only—no explanations, no extra text."
-        )
-
-    def __call__(self, question: str, task_id: str = None) -> str:
-        file_text = ""
-        prompt_parts = [self.instructions]
-        # 1. File (image, Excel, etc)
-        if task_id:
-            file_bytes, content_type = fetch_file(task_id)
-            if file_bytes and content_type:
-                file_text = extract_file_text(file_bytes, content_type, task_id)
-                if file_text:
-                    prompt_parts.append(f"Here is the extracted file content:\n{file_text}\n")
-        # 2. YouTube/video
-        youtube_url = guess_youtube_link(question)
-        if youtube_url:
-            transcript = transcribe_youtube_audio(youtube_url)
-            if transcript:
-                prompt_parts.append(f"Here is the transcript of the video:\n{transcript}\n")
-        # 3. Web search for open facts
-        search_needed = not file_text and not youtube_url
-        search_keywords = [
-            "who", "what", "when", "where", "name", "number", "how many",
-            "first", "last", "award", "recipient", "code", "surname", "year", "album", "actor", "winner"
-        ]
-        if search_needed or any(kw in question.lower() for kw in search_keywords):
-            search_results = run_web_search(question)
-            if search_results:
-                prompt_parts.append(f"Here are relevant web search results:\n{search_results}\n")
-        prompt_parts.append(f"Question: {question}\nAnswer strictly and concisely.")
-        prompt = "\n".join(prompt_parts)
-        response = self.client.chat.completions.create(
-            model="gpt-4o",
-            messages=[
-                {"role": "system", "content": self.instructions},
-                {"role": "user", "content": prompt}
-            ],
-            temperature=0.0,
-            max_tokens=512,
-        )
-        raw_output = safe_strip(response.choices[0].message.content)
-        formatted = format_gaia_answer(raw_output, question)
-        if not formatted or formatted.lower() in ('', 'unknown', 'none', 'n/a') or 'apolog' in formatted.lower():
-            web_info = run_web_search(question)
-            if web_info:
-                prompt2 = (
-                    f"{self.instructions}\n\n"
-                    f"Here are relevant web search results:\n{web_info}\n"
-                    f"Question: {question}\nAnswer strictly and concisely."
-                )
-                response2 = self.client.chat.completions.create(
-                    model="gpt-4o",
-                    messages=[
-                        {"role": "system", "content": self.instructions},
-                        {"role": "user", "content": prompt2}
-                    ],
-                    temperature=0.0,
-                    max_tokens=256,
-                )
-                formatted = format_gaia_answer(safe_strip(response2.choices[0].message.content), question)
-        return formatted
-
-def answer_question(question, task_id=None):
-    agent = GaiaAgent()
-    return agent(question, task_id)
+            return result.get("text", "")
+    except Exception as e:
+        return f"ERROR: {e}"
+
+# ---- LlamaIndex agent and workflow setup ----
+
+# 1. Initialize LLM
+llm = OpenAI(model="gpt-4o", api_key=os.environ.get("OPENAI_API_KEY"))
+
+# 2. Register tools
+tools = [
+    DuckDuckGoSearchTool(),
+    FunctionTool.from_defaults(
+        eval_python_code,
+        name="python_eval",
+        description="Evaluate simple Python code and return result as string. Use for math or code output."
+    ),
+    FunctionTool.from_defaults(
+        ocr_image,
+        name="ocr_image",
+        description="Extract text from an image file (provide file path)."
+    ),
+    FunctionTool.from_defaults(
+        transcribe_audio,
+        name="transcribe_audio",
+        description="Transcribe an audio file using Whisper (provide file path)."
+    ),
+    FunctionTool.from_defaults(
+        transcribe_youtube,
+        name="transcribe_youtube",
+        description="Download a YouTube video, extract and transcribe its audio using Whisper."
+    ),
+    FunctionTool.from_defaults(
+        format_gaia_answer,
+        name="format_gaia_answer",
+        description="Postprocess and enforce strict GAIA format on answers given a question."
+    ),
+]
+
+# 3. Agent setup (ReAct, so can reason with tools)
+agent = ReActAgent.from_tools(
+    tools=tools,
+    llm=llm,
+    system_prompt="You are a helpful GAIA benchmark agent. For every question, use the best tools available and always return only the final answer in the strict GAIA-required format—never explain, never apologize.",
+    verbose=False
+)
+
+# 4. Async entrypoint, suitable for HuggingFace Spaces or Gradio
+async def answer_question(question: str, task_id: str = None, file_path: str = None) -> str:
+    """
+    Main async function for the agent. 
+    Passes the question and uses tools as needed.
+    - task_id: for future use, if you want to fetch files from a remote API.
+    - file_path: if a file (image, audio, etc) is present locally, pass it.
+    """
+    # Example: if you want to always try OCR/audio on a file before reasoning, you could do:
+    # If question contains "image" or "chess" and file_path is set, run OCR first
+    if file_path and any(word in question.lower() for word in ["image", "chess", "screenshot"]):
+        ocr_text = ocr_image(file_path)
+        question = f"Extracted text from image: {ocr_text}\n\n{question}"
+    if file_path and any(word in question.lower() for word in ["audio", "mp3", "transcribe"]):
+        audio_text = transcribe_audio(file_path)
+        question = f"Transcribed audio: {audio_text}\n\n{question}"
+
+    # Run agent
+    result = await agent.achat(question)
+    return result.response
+
+# Synchronous wrapper for legacy compat
+def answer_question_sync(question: str, task_id: str = None, file_path: str = None) -> str:
+    return asyncio.run(answer_question(question, task_id, file_path))