Update tools.py

tools.py

@@ -4,12 +4,12 @@ import json
 import base64
 import tempfile
 import subprocess
+from pathlib import Path
 
 import requests
 import pandas as pd
 
 from bs4 import BeautifulSoup
-from duckduckgo_search import DDGS
 from huggingface_hub import InferenceClient
 
 
@@ -21,20 +21,11 @@ SCORING_API = "https://agents-course-unit4-scoring.hf.space"
 
 HF_TOKEN = os.getenv("HF_TOKEN")
 
-TEXT_MODEL = os.getenv(
-    "TEXT_MODEL",
-    "Qwen/Qwen2.5-72B-Instruct"
-)
+TEXT_MODEL = "Qwen/Qwen2.5-72B-Instruct"
 
-VISION_MODEL = os.getenv(
-    "VISION_MODEL",
-    "Qwen/Qwen2.5-VL-72B-Instruct"
-)
+VISION_MODEL = "Qwen/Qwen2.5-VL-72B-Instruct"
 
-ASR_MODEL = os.getenv(
-    "ASR_MODEL",
-    "openai/whisper-large-v3"
-)
+AUDIO_MODEL = "openai/whisper-large-v3"
 
 client = InferenceClient(
     token=HF_TOKEN
@@ -46,18 +37,31 @@ client = InferenceClient(
 # ============================================================
 
 def web_search(query):
+
+    try:
+
+        from ddgs import DDGS
+
+    except ImportError:
+
+        from duckduckgo_search import DDGS
+
     results = []
 
     with DDGS() as ddgs:
-        for result in ddgs.text(
+
+        for item in ddgs.text(
             query,
-            max_results=6
+            max_results=8
         ):
-            results.append({
-                "title": result.get("title", ""),
-                "url": result.get("href", ""),
-                "snippet": result.get("body", "")
-            })
+
+            results.append(
+                {
+                    "title": item.get("title", ""),
+                    "url": item.get("href", ""),
+                    "snippet": item.get("body", "")
+                }
+            )
 
     return json.dumps(
         results,
@@ -71,11 +75,15 @@ def web_search(query):
 # ============================================================
 
 def read_webpage(url):
+
     response = requests.get(
         url,
-        timeout=25,
+        timeout=30,
         headers={
-            "User-Agent": "Mozilla/5.0"
+            "User-Agent": (
+                "Mozilla/5.0 "
+                "GAIA-Agent/1.0"
+            )
         }
     )
 
@@ -86,13 +94,14 @@ def read_webpage(url):
         "html.parser"
     )
 
-    for tag in soup([
-        "script",
-        "style",
-        "nav",
-        "footer",
-        "header"
-    ]):
+    for tag in soup(
+        [
+            "script",
+            "style",
+            "noscript",
+            "svg"
+        ]
+    ):
         tag.decompose()
 
     text = soup.get_text(
@@ -100,15 +109,16 @@ def read_webpage(url):
         strip=True
     )
 
-    # Avoid sending enormous webpages to the model.
-    return text[:40000]
+    # Prevent enormous webpages from consuming context.
+    return text[:30000]
 
 
 # ============================================================
 # DOWNLOAD GAIA ATTACHMENT
 # ============================================================
 
-def download_file(task_id):
+def download_file(task_id, filename_hint=""):
+
     url = f"{SCORING_API}/files/{task_id}"
 
     response = requests.get(
@@ -123,148 +133,193 @@ def download_file(task_id):
         ""
     ).lower()
 
-    extension = ""
+    content_disposition = response.headers.get(
+        "content-disposition",
+        ""
+    )
 
-    if "png" in content_type:
-        extension = ".png"
+    filename = None
 
-    elif "jpeg" in content_type or "jpg" in content_type:
-        extension = ".jpg"
+    match = re.search(
+        r'filename="?([^"]+)"?',
+        content_disposition
+    )
 
-    elif "mp3" in content_type or "audio" in content_type:
-        extension = ".mp3"
+    if match:
+        filename = match.group(1)
 
-    elif "spreadsheet" in content_type or "excel" in content_type:
-        extension = ".xlsx"
+    if not filename:
 
-    elif "pdf" in content_type:
-        extension = ".pdf"
+        if filename_hint:
 
-    elif "python" in content_type or "text" in content_type:
-        extension = ".txt"
+            filename = Path(
+                filename_hint
+            ).name
 
-    else:
-        # Try to infer from Content-Disposition.
-        disposition = response.headers.get(
-            "content-disposition",
-            ""
-        )
+        else:
 
-        match = re.search(
-            r'filename="?([^"]+)"?',
-            disposition
-        )
+            if "audio" in content_type:
+                extension = ".mp3"
+
+            elif "image" in content_type:
+                extension = ".png"
+
+            elif "spreadsheet" in content_type:
+                extension = ".xlsx"
+
+            elif "pdf" in content_type:
+                extension = ".pdf"
+
+            elif "python" in content_type:
+                extension = ".py"
+
+            else:
+                extension = ".bin"
 
-        if match:
-            extension = os.path.splitext(
-                match.group(1)
-            )[1]
+            filename = (
+                f"{task_id}{extension}"
+            )
+
+    safe_name = re.sub(
+        r"[^A-Za-z0-9_.-]",
+        "_",
+        filename
+    )
 
-    if not extension:
-        extension = ".bin"
+    folder = Path(
+        tempfile.gettempdir()
+    ) / "gaia_files"
 
-    path = os.path.join(
-        tempfile.gettempdir(),
-        f"gaia_{task_id}{extension}"
+    folder.mkdir(
+        parents=True,
+        exist_ok=True
     )
 
-    with open(path, "wb") as f:
-        f.write(response.content)
+    path = folder / safe_name
+
+    path.write_bytes(
+        response.content
+    )
 
-    return path
+    return str(path)
 
 
 # ============================================================
-# FILE INSPECTION
+# TEXT FILE INSPECTION
 # ============================================================
 
-def inspect_file(path):
-    extension = os.path.splitext(
-        path
-    )[1].lower()
+def inspect_text_file(path):
 
-    # ---------------- TXT / PY / CSV ----------------
+    text = Path(path).read_text(
+        errors="ignore"
+    )
 
-    if extension in [
-        ".txt",
-        ".py",
-        ".md",
-        ".csv",
-        ".tsv"
-    ]:
+    return text[:30000]
+
+
+# ============================================================
+# EXCEL INSPECTION
+# ============================================================
+
+def inspect_excel(path):
+
+    excel = pd.ExcelFile(path)
+
+    output = []
 
-        with open(
+    output.append(
+        f"Sheets: {excel.sheet_names}"
+    )
+
+    for sheet in excel.sheet_names:
+
+        df = pd.read_excel(
             path,
-            "r",
-            encoding="utf-8",
-            errors="ignore"
-        ) as f:
-            return f.read()[:50000]
+            sheet_name=sheet
+        )
 
-    # ---------------- EXCEL ----------------
+        output.append(
+            f"\n=== SHEET: {sheet} ==="
+        )
 
-    if extension in [
-        ".xlsx",
-        ".xls"
-    ]:
+        output.append(
+            f"Shape: {df.shape}"
+        )
 
-        excel = pd.ExcelFile(path)
+        output.append(
+            df.to_string(index=False)
+        )
 
-        output = []
+    return "\n".join(output)[:40000]
 
-        for sheet in excel.sheet_names:
 
-            df = pd.read_excel(
-                path,
-                sheet_name=sheet
-            )
+# ============================================================
+# PDF INSPECTION
+# ============================================================
 
-            output.append(
-                f"\n--- SHEET: {sheet} ---\n"
-            )
+def inspect_pdf(path):
 
-            output.append(
-                f"Columns: {list(df.columns)}\n"
-            )
+    try:
 
-            output.append(
-                f"Rows: {len(df)}\n"
-            )
+        from PyPDF2 import PdfReader
+
+        reader = PdfReader(path)
+
+        pages = []
 
-            output.append(
-                df.head(50).to_string(
-                    index=False
-                )
+        for page in reader.pages:
+
+            pages.append(
+                page.extract_text() or ""
             )
 
-        return "\n".join(output)[:50000]
+        return "\n".join(pages)[:40000]
 
-    # ---------------- PDF ----------------
+    except Exception as e:
 
-    if extension == ".pdf":
+        return f"PDF ERROR: {e}"
+
+
+# ============================================================
+# GENERIC FILE INSPECTION
+# ============================================================
+
+def inspect_file(path):
 
-        try:
-            import PyPDF2
+    path = str(path)
+
+    extension = (
+        Path(path)
+        .suffix
+        .lower()
+    )
+
+    if extension in [
+        ".txt",
+        ".py",
+        ".md",
+        ".csv",
+        ".tsv",
+        ".json",
+        ".xml"
+    ]:
 
-            reader = PyPDF2.PdfReader(path)
+        return inspect_text_file(path)
 
-            pages = []
+    if extension in [
+        ".xlsx",
+        ".xls"
+    ]:
 
-            for page in reader.pages:
-                pages.append(
-                    page.extract_text() or ""
-                )
+        return inspect_excel(path)
 
-            return "\n".join(pages)[:50000]
+    if extension == ".pdf":
 
-        except Exception as e:
-            return f"Could not read PDF: {e}"
+        return inspect_pdf(path)
 
     return (
-        f"File downloaded successfully at:\n"
-        f"{path}\n"
-        f"File type: {extension}\n"
-        f"Use an appropriate tool to analyze it."
+        f"Unsupported file type: "
+        f"{extension}"
     )
 
 
@@ -273,39 +328,40 @@ def inspect_file(path):
 # ============================================================
 
 def run_python(code):
-    with tempfile.NamedTemporaryFile(
-        mode="w",
-        suffix=".py",
-        delete=False,
-        encoding="utf-8"
-    ) as f:
 
-        f.write(code)
+    folder = Path(
+        tempfile.gettempdir()
+    ) / "gaia_python"
 
-        path = f.name
+    folder.mkdir(
+        parents=True,
+        exist_ok=True
+    )
 
-    try:
+    filename = folder / "script.py"
 
-        result = subprocess.run(
-            ["python", path],
-            capture_output=True,
-            text=True,
-            timeout=30
-        )
+    filename.write_text(
+        code,
+        encoding="utf-8"
+    )
 
-        return (
-            "STDOUT:\n"
-            + result.stdout
-            + "\nSTDERR:\n"
-            + result.stderr
-        )[:30000]
+    process = subprocess.run(
+        [
+            "python",
+            str(filename)
+        ],
+        capture_output=True,
+        text=True,
+        timeout=30
+    )
 
-    finally:
+    output = (
+        process.stdout
+        + "\n"
+        + process.stderr
+    )
 
-        try:
-            os.remove(path)
-        except OSError:
-            pass
+    return output[-20000:]
 
 
 # ============================================================
@@ -314,12 +370,21 @@ def run_python(code):
 
 def transcribe_audio(path):
 
-    result = client.automatic_speech_recognition(
+    with open(
         path,
-        model=ASR_MODEL
-    )
+        "rb"
+    ) as audio_file:
+
+        result = client.automatic_speech_recognition(
+            audio_file,
+            model=AUDIO_MODEL
+        )
+
+    if hasattr(result, "text"):
 
-    return result.text
+        return result.text
+
+    return str(result)
 
 
 # ============================================================
@@ -328,97 +393,121 @@ def transcribe_audio(path):
 
 def analyze_image(path, question):
 
-    with open(path, "rb") as f:
-        image_bytes = f.read()
+    data = Path(path).read_bytes()
 
     encoded = base64.b64encode(
-        image_bytes
+        data
     ).decode("utf-8")
 
-    # Determine MIME type.
-    extension = os.path.splitext(
-        path
-    )[1].lower()
+    extension = (
+        Path(path)
+        .suffix
+        .lower()
+    )
 
-    if extension == ".jpg" or extension == ".jpeg":
-        mime = "image/jpeg"
-    else:
-        mime = "image/png"
+    mime = {
+        ".png": "image/png",
+        ".jpg": "image/jpeg",
+        ".jpeg": "image/jpeg",
+        ".webp": "image/webp"
+    }.get(
+        extension,
+        "image/png"
+    )
 
-    response = client.chat.completions.create(
-        model=VISION_MODEL,
-        messages=[
-            {
-                "role": "user",
-                "content": [
-                    {
-                        "type": "text",
-                        "text": question
-                    },
-                    {
-                        "type": "image_url",
-                        "image_url": {
-                            "url":
-                            f"data:{mime};base64,{encoded}"
-                        }
+    image_url = (
+        f"data:{mime};base64,{encoded}"
+    )
+
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "text",
+                    "text": question
+                },
+                {
+                    "type": "image_url",
+                    "image_url": {
+                        "url": image_url
                     }
-                ]
-            }
-        ],
-        max_tokens=1000,
-        temperature=0
+                }
+            ]
+        }
+    ]
+
+    response = client.chat_completion(
+        model=VISION_MODEL,
+        messages=messages,
+        max_tokens=1500,
+        temperature=0.1
     )
 
-    return response.choices[0].message.content
+    return (
+        response
+        .choices[0]
+        .message
+        .content
+    )
 
 
 # ============================================================
-# TOOL DISPATCHER
+# TOOL DISPATCH
 # ============================================================
 
 def execute_tool(
     tool_name,
-    tool_input,
-    task_id,
-    question
+    tool_input
 ):
 
     if tool_name == "web_search":
-        return web_search(tool_input)
 
-    if tool_name == "read_webpage":
-        return read_webpage(tool_input)
+        return web_search(
+            str(tool_input)
+        )
 
-    if tool_name == "download_file":
-        path = download_file(task_id)
+    if tool_name == "read_webpage":
 
-        return (
-            f"FILE_PATH:\n{path}\n\n"
-            f"{inspect_file(path)}"
+        return read_webpage(
+            str(tool_input)
         )
 
     if tool_name == "inspect_file":
 
-        path = tool_input.strip()
-
-        return inspect_file(path)
+        return inspect_file(
+            str(tool_input)
+        )
 
     if tool_name == "run_python":
-        return run_python(tool_input)
 
-    if tool_name == "transcribe_audio":
+        return run_python(
+            str(tool_input)
+        )
 
-        path = tool_input.strip()
+    if tool_name == "transcribe_audio":
 
-        return transcribe_audio(path)
+        return transcribe_audio(
+            str(tool_input)
+        )
 
     if tool_name == "analyze_image":
 
-        path = tool_input.strip()
+        if isinstance(
+            tool_input,
+            dict
+        ):
 
-        return analyze_image(
-            path,
-            question
+            return analyze_image(
+                tool_input["path"],
+                tool_input["question"]
+            )
+
+        return (
+            "analyze_image requires "
+            "a path and question."
         )
 
-    return f"Unknown tool: {tool_name}"
+    return (
+        f"Unknown tool: {tool_name}"
+    )