attachment_processing.py

"""Local preprocessing for the attachment types used by the evaluation set."""

from __future__ import annotations

import base64
import io
import os
import warnings
import wave
from pathlib import Path
from typing import Any

import requests

from model_config import DEFAULT_CONTEXT_SIZE, DEFAULT_OLLAMA_MODEL


MAX_EXTRACTED_CHARS = 40_000


class AttachmentProcessingError(RuntimeError):
    """Raised when a task attachment cannot be converted to text evidence."""


class AttachmentProcessor:
    """Convert task attachments into bounded text for the research agent."""

    def __init__(self) -> None:
        self.ollama_base_url = os.getenv(
            "OLLAMA_BASE_URL", "http://localhost:11434"
        ).rstrip("/")
        self.multimodal_model = os.getenv(
            "OLLAMA_MULTIMODAL_MODEL",
            os.getenv("OLLAMA_VISION_MODEL", DEFAULT_OLLAMA_MODEL),
        )
        self.context_size = int(
            os.getenv("OLLAMA_CONTEXT_SIZE", str(DEFAULT_CONTEXT_SIZE))
        )
        self.audio_transport = os.getenv(
            "OLLAMA_AUDIO_TRANSPORT", "images"
        ).lower()
        if self.audio_transport not in {"images", "audios"}:
            raise AttachmentProcessingError(
                "OLLAMA_AUDIO_TRANSPORT must be either 'images' or 'audios'."
            )
        self.audio_chunk_seconds = int(
            os.getenv("GEMMA_AUDIO_CHUNK_SECONDS", "28")
        )
        if not 1 <= self.audio_chunk_seconds <= 30:
            raise AttachmentProcessingError(
                "GEMMA_AUDIO_CHUNK_SECONDS must be between 1 and 30."
            )
        self.audio_fallback = os.getenv(
            "OLLAMA_AUDIO_FALLBACK", "whisper"
        ).lower()
        if self.audio_fallback not in {"none", "whisper"}:
            raise AttachmentProcessingError(
                "OLLAMA_AUDIO_FALLBACK must be either 'none' or 'whisper'."
            )
        self.whisper_model = os.getenv("WHISPER_MODEL", "small.en")
        self.whisper_device = os.getenv("WHISPER_DEVICE", "cpu")
        self.whisper_compute_type = os.getenv("WHISPER_COMPUTE_TYPE", "int8")

    def process(self, path: Path | None, question: str) -> str:
        if path is None:
            return "No attachment was provided for this task."
        if not path.is_file():
            raise AttachmentProcessingError(f"Attachment does not exist: {path}")

        suffix = path.suffix.lower()
        if suffix in {".png", ".jpg", ".jpeg", ".webp"}:
            result = self._describe_image(path, question)
        elif suffix in {".mp3", ".wav", ".m4a", ".flac", ".ogg"}:
            result = self._analyze_audio(path, question)
        elif suffix in {".xlsx", ".xlsm"}:
            result = self._extract_workbook(path)
        elif suffix in {
            ".py",
            ".txt",
            ".md",
            ".csv",
            ".tsv",
            ".json",
            ".html",
            ".xml",
        }:
            result = self._extract_text(path)
        else:
            raise AttachmentProcessingError(
                f"Unsupported attachment type {suffix or '<none>'}: {path.name}"
            )

        return result[:MAX_EXTRACTED_CHARS]

    def _describe_image(self, path: Path, question: str) -> str:
        encoded = base64.b64encode(path.read_bytes()).decode("ascii")
        prompt = (
            "Inspect this task image carefully. Transcribe every relevant word, "
            "number, label, axis, legend, and table cell, then describe visual "
            "relationships needed to answer the question. Distinguish direct "
            "observations from uncertainty.\n\nQuestion:\n" + question
        )
        return self._multimodal_chat(
            prompt=prompt,
            encoded_media=encoded,
            media_field="images",
            description=f"image {path.name}",
        )

    def _analyze_audio(self, path: Path, question: str) -> str:
        """Give Gemma 4 the audio itself; use transcription only as a fallback."""

        try:
            wav_chunks = self._audio_as_wav_chunks(path)
            analyses = []
            for index, encoded in enumerate(wav_chunks, start=1):
                prompt = (
                    "Listen to this audio carefully. Transcribe all intelligible "
                    "speech, preserving names, numbers, spelling, and sequence. "
                    "Also identify relevant non-speech sounds, speakers, music, "
                    "timing, or uncertainty. Use the question to focus the analysis, "
                    "but report observations rather than guessing.\n\n"
                    f"Audio chunk: {index}/{len(wav_chunks)}\n"
                    f"Question:\n{question}"
                )
                analyses.append(
                    self._multimodal_chat(
                        prompt=prompt,
                        encoded_media=encoded,
                        media_field=self.audio_transport,
                        description=f"audio {path.name} chunk {index}",
                    )
                )
            return (
                f"Gemma 4 audio analysis for {path.name} "
                f"({len(wav_chunks)} chunk(s)):\n"
                + "\n\n".join(analyses)
            )
        except AttachmentProcessingError as gemma_error:
            if self.audio_fallback == "none":
                raise
            warnings.warn(
                f"Gemma 4 audio analysis failed for {path.name}; using the "
                f"Whisper fallback. Cause: {gemma_error}",
                RuntimeWarning,
                stacklevel=2,
            )
            transcript = self._transcribe_audio(path)
            return (
                "Gemma 4 audio analysis was unavailable. Whisper fallback was "
                "used, so non-speech audio details may be absent.\n" + transcript
            )

    def _multimodal_chat(
        self,
        *,
        prompt: str,
        encoded_media: str,
        media_field: str,
        description: str,
    ) -> str:
        payload: dict[str, Any] = {
            "model": self.multimodal_model,
            "messages": [
                {"role": "user", "content": prompt, media_field: [encoded_media]}
            ],
            "stream": False,
            "think": False,
            "options": {"temperature": 0, "num_ctx": self.context_size},
        }
        try:
            response = requests.post(
                f"{self.ollama_base_url}/api/chat", json=payload, timeout=300
            )
            response.raise_for_status()
            content = response.json()["message"]["content"]
        except (requests.RequestException, KeyError, TypeError, ValueError) as exc:
            raise AttachmentProcessingError(
                f"Multimodal model {self.multimodal_model!r} failed for "
                f"{description}: {exc}"
            ) from exc
        if not isinstance(content, str) or not content.strip():
            raise AttachmentProcessingError(
                f"Multimodal model {self.multimodal_model!r} returned no content "
                f"for {description}."
            )
        return content.strip()

    def _audio_as_wav_chunks(self, path: Path) -> list[str]:
        """Decode audio without transcribing it and return bounded WAV chunks."""

        try:
            import av
        except ImportError as exc:
            raise AttachmentProcessingError(
                "Gemma 4 audio input requires PyAV. Install requirements.txt."
            ) from exc

        pcm = bytearray()
        try:
            with av.open(str(path)) as container:
                if not container.streams.audio:
                    raise AttachmentProcessingError(
                        f"No audio stream was found in {path.name}."
                    )
                resampler = av.AudioResampler(
                    format="s16", layout="mono", rate=16_000
                )
                for frame in container.decode(audio=0):
                    for converted in resampler.resample(frame):
                        pcm.extend(converted.to_ndarray().tobytes())
                for converted in resampler.resample(None):
                    pcm.extend(converted.to_ndarray().tobytes())
        except AttachmentProcessingError:
            raise
        except Exception as exc:
            raise AttachmentProcessingError(
                f"Could not decode {path.name} for Gemma 4: {exc}"
            ) from exc

        if not pcm:
            raise AttachmentProcessingError(f"Decoded audio was empty: {path.name}")

        bytes_per_second = 16_000 * 2  # mono, signed 16-bit PCM
        chunk_size = self.audio_chunk_seconds * bytes_per_second
        chunks = []
        for offset in range(0, len(pcm), chunk_size):
            buffer = io.BytesIO()
            with wave.open(buffer, "wb") as wav_file:
                wav_file.setnchannels(1)
                wav_file.setsampwidth(2)
                wav_file.setframerate(16_000)
                wav_file.writeframes(pcm[offset : offset + chunk_size])
            chunks.append(base64.b64encode(buffer.getvalue()).decode("ascii"))
        return chunks

    def _transcribe_audio(self, path: Path) -> str:
        try:
            from faster_whisper import WhisperModel
        except ImportError as exc:
            raise AttachmentProcessingError(
                "Audio transcription requires faster-whisper. Install requirements.txt."
            ) from exc

        try:
            model = WhisperModel(
                self.whisper_model,
                device=self.whisper_device,
                compute_type=self.whisper_compute_type,
            )
            segments, info = model.transcribe(
                str(path), beam_size=5, vad_filter=True
            )
            lines = [
                f"[{segment.start:.2f}-{segment.end:.2f}] {segment.text.strip()}"
                for segment in segments
                if segment.text.strip()
            ]
        except Exception as exc:  # library raises backend-specific error classes
            raise AttachmentProcessingError(
                f"Speech transcription failed for {path.name}: {exc}"
            ) from exc

        language = getattr(info, "language", "unknown")
        return (
            f"Audio transcript for {path.name} (detected language: {language}):\n"
            + "\n".join(lines)
        )

    @staticmethod
    def _extract_workbook(path: Path) -> str:
        try:
            from openpyxl import load_workbook
        except ImportError as exc:
            raise AttachmentProcessingError(
                "XLSX extraction requires openpyxl. Install requirements.txt."
            ) from exc

        try:
            workbook = load_workbook(path, read_only=True, data_only=False)
        except Exception as exc:
            raise AttachmentProcessingError(
                f"Could not open workbook {path.name}: {exc}"
            ) from exc

        output = [f"Workbook extraction for {path.name}:"]
        remaining = MAX_EXTRACTED_CHARS
        try:
            for sheet in workbook.worksheets:
                output.append(f"\nSheet: {sheet.title}")
                for row_index, row in enumerate(
                    sheet.iter_rows(max_row=500, max_col=100), start=1
                ):
                    cells = []
                    for cell in row:
                        if cell.value is not None:
                            cells.append(f"{cell.coordinate}={cell.value!r}")
                    if cells:
                        line = f"Row {row_index}: " + " | ".join(cells)
                        output.append(line)
                        remaining -= len(line)
                        if remaining <= 0:
                            output.append("[Workbook output truncated]")
                            return "\n".join(output)
        finally:
            workbook.close()
        return "\n".join(output)

    @staticmethod
    def _extract_text(path: Path) -> str:
        try:
            text = path.read_text(encoding="utf-8", errors="replace")
        except OSError as exc:
            raise AttachmentProcessingError(
                f"Could not read text attachment {path.name}: {exc}"
            ) from exc
        return f"Text extraction for {path.name}:\n{text}"