Create Attachment reader.py

tools/Attachment reader.py

@@ -0,0 +1,158 @@
+"""
+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 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.
+    """
+    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:
+            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}")
+        return path
+    except Exception as e:
+        print(f"Could not download attachment for task {task_id}: {e}")
+        return None
+
+
+# =========================================================
+# 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}"