Update tools/attachment_reader.py

tools/attachment_reader.py

@@ -42,35 +42,75 @@ def _filename_from_response(response: requests.Response, task_id: str) -> str:
     return f"{task_id}{extension}"
 
 
-def download_attachment(api_url: str, task_id: str, timeout: int = 30) -> str | None:
-    """
-    Downloads the attachment for a task, if there is one.
-    Returns the local file path, or None if the task has no attachment
-    or the download failed. Never raises.
-    """
+def _save_bytes(task_id: str, filename: str, data: bytes) -> str:
+    task_dir = os.path.join(ATTACHMENTS_DIR, str(task_id))
+    os.makedirs(task_dir, exist_ok=True)
+    path = os.path.join(task_dir, filename)
+    with open(path, "wb") as f:
+        f.write(data)
+    return path
+
+
+def _download_from_scoring_api(api_url: str, task_id: str, timeout: int = 30) -> str | None:
+    """Try the scoring API's /files/{task_id} endpoint. Returns a path or None."""
     url = f"{api_url}/files/{task_id}"
     try:
         response = requests.get(url, timeout=timeout)
-        if response.status_code == 404:
-            return None  # task has no attachment
-        response.raise_for_status()
-        if not response.content:
+        if response.status_code != 200 or not response.content:
+            print(
+                f"[attachment] {url} -> HTTP {response.status_code} "
+                f"{response.text[:150]!r}"
+            )
             return None
+        path = _save_bytes(task_id, _filename_from_response(response, task_id), response.content)
+        print(f"[attachment] downloaded from scoring API: {path}")
+        return path
+    except Exception as e:
+        print(f"[attachment] scoring API download failed for {task_id}: {e}")
+        return None
 
-        task_dir = os.path.join(ATTACHMENTS_DIR, str(task_id))
-        os.makedirs(task_dir, exist_ok=True)
-        path = os.path.join(task_dir, _filename_from_response(response, task_id))
-
-        with open(path, "wb") as f:
-            f.write(response.content)
 
-        print(f"Downloaded attachment for task {task_id}: {path}")
+def _download_from_gaia_dataset(task_id: str, file_name: str) -> str | None:
+    """
+    Fallback: fetch the file straight from the gated GAIA dataset on the Hub.
+    Needs HF_TOKEN in the environment (a Space secret) and that your account
+    has accepted the terms at https://huggingface.co/datasets/gaia-benchmark/GAIA
+    """
+    token = os.environ.get("HF_TOKEN")
+    if not token:
+        print("[attachment] HF_TOKEN not set; cannot use the GAIA dataset fallback.")
+        return None
+    try:
+        from huggingface_hub import hf_hub_download
+
+        cached = hf_hub_download(
+            repo_id="gaia-benchmark/GAIA",
+            filename=f"2023/validation/{file_name}",
+            repo_type="dataset",
+            token=token,
+        )
+        with open(cached, "rb") as f:
+            path = _save_bytes(task_id, os.path.basename(file_name), f.read())
+        print(f"[attachment] downloaded from GAIA dataset: {path}")
         return path
     except Exception as e:
-        print(f"Could not download attachment for task {task_id}: {e}")
+        print(f"[attachment] GAIA dataset fallback failed for {file_name}: {e}")
         return None
 
 
+def download_attachment(api_url: str, task_id: str, file_name: str | None = None,
+                        timeout: int = 30) -> str | None:
+    """
+    Gets the attachment for a task. Tries the scoring API first and, if that
+    fails and the task's file_name is known, falls back to the GAIA dataset.
+    Returns the local file path, or None. Never raises.
+    """
+    path = _download_from_scoring_api(api_url, task_id, timeout)
+    if path is None and file_name:
+        path = _download_from_gaia_dataset(task_id, file_name)
+    return path
+
+
 # =========================================================
 # Reader tool
 # =========================================================