brain-of-reese/app/rag/md_images.py

"""Markdown image-reference extraction (phase 125, LOCKED A1).

Pure functions, no I/O — fully unit-testable. Stdlib only.

Markdown documents reference images inline (``![alt](path)``). The
paragraphs AROUND the reference carry the *why* — what the diagram
illustrates, which procedure it belongs to — context that the
standalone vision description of the image itself (phase 122) does not
have. :func:`extract_image_refs` finds every image reference and
returns it with its alt text, its 1-based line number, and the
surrounding paragraph context: up to *n* paragraphs before and after
the reference paragraph (default :data:`CONTEXT_PARAGRAPHS`).

Paragraph semantics mirror the chunker (PLAN §5): paragraphs are
separated by blank lines and code fences are **atomic** — a fenced
block (possibly containing blank lines) is ONE paragraph, so it counts
as a single unit of the context window. Only the *references* inside a
fence are skipped: there they are literal text, not actual images.
"""
from __future__ import annotations

import re
from collections.abc import Sequence
from dataclasses import dataclass

#: Opening/closing code fence (three or more backticks or tildes,
#: optionally indented) — the SAME semantics as the chunker's fence
#: handling (PLAN §5), so "fence" means the same thing here as in
#: chunking.
_FENCE_RE = re.compile(r"^(`{3,}|~{3,})")

#: Markdown image reference ``![alt](path)``. The alt may be EMPTY
#: (``![](path)``); the path may contain spaces but not ``)`` — nested
#: parentheses inside a path are accepted as a miss (rare in practice,
#: phase 125 task 02).
_IMAGE_REF_RE = re.compile(r"!\[([^\]]*)\]\(([^)]+)\)")

#: Default number of context paragraphs BEFORE and AFTER the reference
#: paragraph (LOCKED A1: "alt text + N paragraphs before/after the
#: reference") — 2 gives ~2-4 sentences of surrounding text.
CONTEXT_PARAGRAPHS: int = 2


@dataclass
class ImageRef:
    """One extracted ``![alt](path)`` reference with its context window.

    * ``path`` — the image path (relative or absolute; surrounding
      whitespace stripped);
    * ``alt_text`` — the alt text (may be empty for ``![](path)``);
    * ``context`` — the surrounding paragraphs: up to *n* before and
      *n* after the reference paragraph, joined with blank lines. The
      reference paragraph itself is NOT included (the reference is
      already captured as ``alt_text`` + ``path``);
    * ``line_number`` — 1-based line number of the reference.
    """

    path: str
    alt_text: str
    context: str
    line_number: int


def _fence_flags(lines: Sequence[str]) -> list[bool]:
    """Per-line flags: ``True`` on a fence delimiter or inside a fence.

    Mirrors the chunker (PLAN §5): an opener is three or more backticks
    or tildes (indented allowed); while open, EVERY line — including
    blank ones — is in-fence until the first line starting with the
    same three-character marker.
    """
    flags: list[bool] = []
    marker: str | None = None
    for line in lines:
        stripped = line.lstrip()
        if marker is None:
            m = _FENCE_RE.match(stripped)
            if m:
                marker = m.group(1)[:3]
                flags.append(True)
            else:
                flags.append(False)
        else:
            flags.append(True)
            if stripped.startswith(marker):
                marker = None
    return flags


def _paragraph_blocks(
    lines: Sequence[str], flags: Sequence[bool]
) -> tuple[list[tuple[int, int]], list[int]]:
    """Paragraph blocks as half-open (start, end) line ranges, plus a
    per-line block-index map (``-1`` on blank outside-fence lines).

    Blocks are separated by blank lines *outside* fences; a fenced code
    block (possibly containing blank lines) is ONE block — atomic (the
    same paragraph semantics the chunker chunks with, PLAN §5).
    """
    blocks: list[tuple[int, int]] = []
    start: int | None = None
    for i, (line, in_fence) in enumerate(zip(lines, flags, strict=True)):
        if in_fence or line.strip():
            if start is None:
                start = i
        elif start is not None:
            blocks.append((start, i))
            start = None
    if start is not None:
        blocks.append((start, len(lines)))
    line_block = [-1] * len(lines)
    for bi, (s, e) in enumerate(blocks):
        line_block[s:e] = [bi] * (e - s)
    return blocks, line_block


def extract_image_refs(
    content: str, context_paragraphs: int = CONTEXT_PARAGRAPHS
) -> list[ImageRef]:
    """Extract every markdown image reference in *content*, with context.

    Each ``![alt](path)`` reference OUTSIDE a code fence becomes one
    :class:`ImageRef` (multiple references on one line are extracted
    separately, each with the same context window). ``context`` joins
    the up-to-*context_paragraphs* paragraphs BEFORE and AFTER the
    reference paragraph with blank lines; the reference paragraph
    itself is excluded (the reference is already captured as
    ``alt_text`` + ``path``), and a reference at the start/end of the
    file gets whatever paragraphs are available. Fences are atomic
    paragraphs (one fence block = one context unit), and references
    inside one are literal text — skipped.

    A negative *context_paragraphs* raises :class:`ValueError` (fail
    loud); 0 means "no surrounding context".
    """
    if context_paragraphs < 0:
        raise ValueError("context_paragraphs must be >= 0")
    lines = content.splitlines()
    if not lines:
        return []
    flags = _fence_flags(lines)
    blocks, line_block = _paragraph_blocks(lines, flags)
    refs: list[ImageRef] = []
    for lineno, (line, in_fence) in enumerate(zip(lines, flags, strict=True), start=1):
        if in_fence:
            continue
        for m in _IMAGE_REF_RE.finditer(line):
            # The reference line is non-blank, outside-fence content, so
            # it always belongs to a paragraph block (>= 0).
            bi = line_block[lineno - 1]
            before = blocks[max(0, bi - context_paragraphs) : bi]
            after = blocks[bi + 1 : bi + 1 + context_paragraphs]
            context = "\n\n".join("\n".join(lines[s:e]) for s, e in before + after)
            refs.append(
                ImageRef(
                    path=m.group(2).strip(),
                    alt_text=m.group(1).strip(),
                    context=context,
                    line_number=lineno,
                )
            )
    return refs