src/gaia_agent/attachments.py

import ast
import functools
import subprocess
import sys
from pathlib import Path
from typing import Any

import pandas as pd
from faster_whisper import WhisperModel
from pypdf import PdfReader

MAX_TEXT_CHARACTERS = 40_000
PYTHON_TIMEOUT_SECONDS = 10


def inspect_attachment(path: Path) -> str:
    """Extract bounded text from trusted benchmark attachments."""
    suffix = path.suffix.lower()
    if suffix in {".txt", ".md", ".csv", ".json"}:
        return path.read_text(encoding="utf-8", errors="replace")[:MAX_TEXT_CHARACTERS]
    if suffix in {".xlsx", ".xls"}:
        return _inspect_workbook(path)
    if suffix == ".pdf":
        return _inspect_pdf(path)
    if suffix == ".py":
        return _inspect_python(path)
    if suffix in {".mp3", ".wav", ".m4a"}:
        return f"Audio transcript:\n{transcribe_media(path)}"
    if suffix in {".png", ".jpg", ".jpeg", ".webp"}:
        return f"Image attachment available at {path}."
    return f"Unsupported attachment type: {suffix or 'unknown'}"


def _inspect_workbook(path: Path) -> str:
    sheets = pd.read_excel(path, sheet_name=None)
    sections = []
    for name, frame in sheets.items():
        sections.append(f"## Sheet: {name}\n{frame.to_csv(index=False)}")
    return "\n\n".join(sections)[:MAX_TEXT_CHARACTERS]


def _inspect_pdf(path: Path) -> str:
    text = "\n\n".join(page.extract_text() or "" for page in PdfReader(path).pages)
    return text[:MAX_TEXT_CHARACTERS]


def _inspect_python(path: Path) -> str:
    path = path.resolve()
    source = path.read_text(encoding="utf-8", errors="replace")
    ast.parse(source)
    completed = subprocess.run(
        [sys.executable, "-I", str(path)],
        check=False,
        capture_output=True,
        text=True,
        timeout=PYTHON_TIMEOUT_SECONDS,
        cwd=path.parent,
    )
    return (
        f"Source:\n{source}\n\nExit code: {completed.returncode}\n"
        f"stdout:\n{completed.stdout}\nstderr:\n{completed.stderr}"
    )[:MAX_TEXT_CHARACTERS]


@functools.lru_cache(maxsize=1)
def _whisper_model() -> Any:
    return WhisperModel("small.en", device="cpu", compute_type="int8")


def transcribe_media(path: Path) -> str:
    segments, _ = _whisper_model().transcribe(str(path), beam_size=5)
    return " ".join(segment.text.strip() for segment in segments).strip()