| """ |
| tools/attachment_reader.py |
| |
| Tool that lets the agent read a task's attached file from a local path, |
| plus a helper that downloads the attachment from the scoring API. |
| |
| Optional dependencies (only needed for the matching file types): |
| pip install openpyxl # .xlsx |
| pip install pypdf # .pdf |
| pip install python-docx # .docx |
| """ |
|
|
| import os |
| import re |
| import mimetypes |
|
|
| import requests |
| import pandas as pd |
| from smolagents import Tool |
|
|
| ATTACHMENTS_DIR = "attachments" |
| MAX_CHARS = 20000 |
|
|
| TEXT_EXTENSIONS = {".txt", ".md", ".json", ".py", ".html", ".xml", ".yaml", ".yml", ".log", ".jsonl"} |
| IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp"} |
| AUDIO_EXTENSIONS = {".mp3", ".wav", ".m4a", ".flac", ".ogg"} |
|
|
|
|
| |
| |
| |
|
|
| def _filename_from_response(response: requests.Response, task_id: str) -> str: |
| """Work out a sensible filename from the response headers.""" |
| disposition = response.headers.get("Content-Disposition", "") |
| match = re.search(r'filename\*?=(?:UTF-8\'\')?"?([^";]+)"?', disposition) |
| if match: |
| return os.path.basename(match.group(1).strip()) |
|
|
| content_type = response.headers.get("Content-Type", "").split(";")[0].strip() |
| extension = mimetypes.guess_extension(content_type) or "" |
| return f"{task_id}{extension}" |
|
|
|
|
| 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 != 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 |
|
|
|
|
| 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"[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 |
|
|
|
|
| |
| |
| |
|
|
| def _truncate(text: str) -> str: |
| if len(text) <= MAX_CHARS: |
| return text |
| return text[:MAX_CHARS] + f"\n\n[... truncated, {len(text) - MAX_CHARS} more characters ...]" |
|
|
|
|
| def _read_spreadsheet(path: str, extension: str) -> str: |
| if extension == ".csv": |
| sheets = {"csv": pd.read_csv(path)} |
| else: |
| sheets = pd.read_excel(path, sheet_name=None) |
|
|
| parts = [] |
| for name, df in sheets.items(): |
| parts.append(f"=== Sheet: {name} ({df.shape[0]} rows x {df.shape[1]} columns) ===") |
| parts.append(df.to_csv(index=False)) |
| return "\n".join(parts) |
|
|
|
|
| def _read_pdf(path: str) -> str: |
| from pypdf import PdfReader |
|
|
| reader = PdfReader(path) |
| return "\n\n".join( |
| f"--- Page {i + 1} ---\n{page.extract_text() or ''}" |
| for i, page in enumerate(reader.pages) |
| ) |
|
|
|
|
| def _read_docx(path: str) -> str: |
| import docx |
|
|
| document = docx.Document(path) |
| return "\n".join(p.text for p in document.paragraphs) |
|
|
|
|
| class AttachmentReaderTool(Tool): |
| name = "read_attachment" |
|
|
| description = ( |
| "Reads a file attached to the current task from a local path and returns " |
| "its contents as text. Supports spreadsheets (.xlsx, .xls, .csv), PDF, " |
| "Word (.docx) and plain text files. Use this when the question mentions " |
| "an attached file and a local path is provided." |
| ) |
|
|
| inputs = { |
| "file_path": { |
| "type": "string", |
| "description": "Local path of the attached file, as given in the question.", |
| } |
| } |
|
|
| output_type = "string" |
|
|
| def forward(self, file_path: str) -> str: |
| if not os.path.exists(file_path): |
| return f"ERROR: file not found: {file_path}" |
|
|
| extension = os.path.splitext(file_path)[1].lower() |
|
|
| try: |
| if extension in {".xlsx", ".xlsm", ".xls", ".csv"}: |
| return _truncate(_read_spreadsheet(file_path, extension)) |
| if extension == ".pdf": |
| return _truncate(_read_pdf(file_path)) |
| if extension == ".docx": |
| return _truncate(_read_docx(file_path)) |
| if extension in TEXT_EXTENSIONS: |
| with open(file_path, "r", encoding="utf-8", errors="replace") as f: |
| return _truncate(f.read()) |
| if extension in IMAGE_EXTENSIONS: |
| return "ERROR: image files cannot be read as text by this tool." |
| if extension in AUDIO_EXTENSIONS: |
| return "ERROR: audio files cannot be read as text by this tool." |
|
|
| |
| with open(file_path, "r", encoding="utf-8") as f: |
| return _truncate(f.read()) |
| except Exception as e: |
| return f"ERROR: could not read {file_path}: {e}" |
