Blame for tools/attachment_reader.py

"""
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"}


# =========================================================
# Download helper
# =========================================================

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


# =========================================================
# Reader tool
# =========================================================

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."

            # Unknown extension: try as text.
            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}"