payload_utils.py

"""
payload_utils.py — Space 2 helper library for querying the enriched IngestPayload.

Drop this file into Space 2's source tree. It is the *only* file Space 2 needs
to work with the new payload fields; it has no dependency on networkx, OpenCV,
or any OCR library.

All functions accept the raw payload dict (the JSON-decoded body of the POST
from Space 1) and return plain Python types.

Public API
----------
chunks_in_reading_order(payload)
    -> list[dict]   all chunks, sorted by document reading order

chunks_for_section(payload, section_title, *, partial=True)
    -> list[dict]   all chunks whose outline_path contains the given title

chunks_by_type(payload, region_type)
    -> list[dict]   all chunks of a given region_type ("paragraph"|"table"|"figure")

chunks_on_page(payload, page_num)
    -> list[dict]   all chunks whose source blocks are on page_num

outline_tree(payload)
    -> list[dict]   the nested section tree (the `outline` field, convenience accessor)

breadcrumb_for_chunk(payload, chunk_id)
    -> list[str]    the outline_path for a chunk, e.g. ["Methods","Data Collection"]

embed_text_for_chunk(chunk, payload=None)
    -> str          builds the recommended embedding string for a chunk.

                    Fixes applied vs the previous version:
                    1. Strips [Context: ...] prefix before embedding — these
                       carry-forward sentences from the chunker pollute embeddings
                       with unrelated topic signals and cause false-positive
                       retrieval (Issue #1 and #4 from the RAG issues report).
                    2. Prefixes the clean text with the full outline breadcrumb
                       so that a sentence like "it was found to be significant"
                       embeds near its section topic rather than floating in
                       generic space.
                    3. For table chunks: converts the plain-text table content
                       to a simple Markdown representation so the embedding
                       model sees structured rows rather than a flat string.
                       If table_html is present it is used directly (HTML
                       preserves row/column structure better than OCR'd text).
                    4. For figure chunks: only the caption / semantic description
                       is embedded, not raw OCR text from logos or decorative
                       elements.  The figure's section_title is preserved in the
                       breadcrumb so the figure is still reachable by section.

strip_context_prefix(text)
    -> str          LEGACY GUARD — Space 1 now stores context in chunk.context
                    (separate field) so chunk.text is always clean. This function
                    is kept as a safety net for any old payloads still in flight
                    but is a no-op on current payloads.

table_text_to_markdown(text)
    -> str          best-effort conversion of a flat OCR'd table string into
                    Markdown pipe-table format so the embedding model sees
                    column structure. Used internally by embed_text_for_chunk.
"""
from __future__ import annotations

import re
from typing import Any, Dict, List, Optional


# ── context-prefix stripping ──────────────────────────────────────────────────

_CONTEXT_PREFIX_RE = re.compile(r'^\[Context:[^\]]*\]\s*\n?', re.DOTALL)


def strip_context_prefix(text: str) -> str:
    """
    LEGACY GUARD: Space 1 now stores the carry-forward sentence in Chunk.context
    (a separate field) so Chunk.text is always clean. This function is kept only
    as a safety net for any old payloads still in flight.

    On current payloads this is effectively a no-op.
    """
    return _CONTEXT_PREFIX_RE.sub("", text).lstrip()


# ── table text → markdown ─────────────────────────────────────────────────────

def table_text_to_markdown(text: str) -> str:
    """
    Best-effort conversion of a flat OCR'd table string to Markdown pipe-table.

    The OCR pipeline produces tables as plain text with whitespace-separated
    columns and newline-separated rows.  Embedding that flat string loses all
    column structure so parameter lookups ("what is the sensing radius?") fail
    to retrieve the correct table row.

    Strategy
    --------
    1. Split by newlines into candidate rows.
    2. For each row, split on 2+ consecutive spaces (typical column separator
       in OCR'd tables) or on tab characters.
    3. Align all rows to the same column count (pad with empty strings).
    4. Emit a Markdown pipe-table with a separator row after the first (header)
       row.

    If the text has fewer than 2 lines or cannot be split into multiple columns,
    the original text is returned unchanged — no point adding Markdown noise to
    a single-column block.
    """
    lines = [l for l in text.splitlines() if l.strip()]
    if len(lines) < 2:
        return text

    # Split each line on 2+ spaces or tabs
    _col_split = re.compile(r'\t|  +')
    rows = [_col_split.split(line.strip()) for line in lines]

    # Check that at least some rows have multiple columns
    max_cols = max(len(r) for r in rows)
    if max_cols < 2:
        return text  # single-column — not a structured table

    # Pad every row to max_cols
    rows = [r + [""] * (max_cols - len(r)) for r in rows]

    def _row_to_pipe(cells: List[str]) -> str:
        return "| " + " | ".join(c.strip() for c in cells) + " |"

    header_sep = "| " + " | ".join("---" for _ in range(max_cols)) + " |"

    md_rows = [_row_to_pipe(rows[0]), header_sep]
    md_rows += [_row_to_pipe(r) for r in rows[1:]]
    return "\n".join(md_rows)


# ── internal helpers ──────────────────────────────────────────────────────────

def _chunk_index(payload: Dict[str, Any]) -> Dict[str, Dict[str, Any]]:
    """Build a chunk_id -> chunk dict from the payload."""
    return {c["chunk_id"]: c for c in payload.get("chunks", [])}


# ── reading order ─────────────────────────────────────────────────────────────

def chunks_in_reading_order(payload: Dict[str, Any]) -> List[Dict[str, Any]]:
    """
    Return all chunks in document reading order.

    Uses `reading_order` (list of block IDs) + `block_to_chunks` (block ->
    chunk IDs) to reconstruct the sequence.  Falls back to the order they
    appear in payload["chunks"] when block_to_chunks is absent (old payloads).
    """
    reading_order   = payload.get("reading_order", [])
    block_to_chunks = payload.get("block_to_chunks") or {}
    chunk_idx       = _chunk_index(payload)

    if not block_to_chunks:
        return list(payload.get("chunks", []))

    seen: set = set()
    ordered: List[Dict[str, Any]] = []

    for block_id in reading_order:
        for chunk_id in block_to_chunks.get(block_id, []):
            if chunk_id not in seen and chunk_id in chunk_idx:
                ordered.append(chunk_idx[chunk_id])
                seen.add(chunk_id)

    # Append any chunks not reached via reading_order (edge case)
    for c in payload.get("chunks", []):
        if c["chunk_id"] not in seen:
            ordered.append(c)

    return ordered


# ── section filtering ─────────────────────────────────────────────────────────

def chunks_for_section(
    payload: Dict[str, Any],
    section_title: str,
    *,
    partial: bool = True,
) -> List[Dict[str, Any]]:
    """
    Return all chunks whose outline_path contains `section_title`.

    Parameters
    ----------
    section_title : the heading text to match (case-insensitive)
    partial       : if True (default), match substrings
                    ("Data" matches "Data Collection")
                    if False, require exact match

    Results are returned in reading order.
    """
    needle = section_title.lower()
    chunk_outline_path = payload.get("chunk_outline_path") or {}

    def _matches(path: List[str]) -> bool:
        for title in path:
            t = title.lower()
            if partial and needle in t:
                return True
            if not partial and t == needle:
                return True
        return False

    ordered = chunks_in_reading_order(payload)
    return [
        c for c in ordered
        if _matches(chunk_outline_path.get(c["chunk_id"], c.get("outline_path") or []))
    ]


# ── type filtering ────────────────────────────────────────────────────────────

def chunks_by_type(
    payload: Dict[str, Any],
    region_type: str,
) -> List[Dict[str, Any]]:
    """
    Return all chunks of the given region_type in reading order.

    region_type: "paragraph" | "table" | "figure"

    Uses `content_type_index` for O(1) lookup then resolves to chunk dicts.
    Falls back to a full scan if the index is absent (backwards compatibility).
    """
    content_type_index = payload.get("content_type_index")
    chunk_idx          = _chunk_index(payload)

    if content_type_index is not None:
        ids = set(content_type_index.get(region_type, []))
        return [c for c in chunks_in_reading_order(payload) if c["chunk_id"] in ids]

    return [
        c for c in chunks_in_reading_order(payload)
        if c.get("region_type") == region_type
    ]


# ── page filtering ────────────────────────────────────────────────────────────

def chunks_on_page(
    payload: Dict[str, Any],
    page_num: int,
) -> List[Dict[str, Any]]:
    """
    Return all chunks whose source blocks appear on `page_num` (1-indexed).
    """
    page_index      = payload.get("page_index")
    block_to_chunks = payload.get("block_to_chunks") or {}
    chunk_idx       = _chunk_index(payload)

    if page_index is not None:
        block_ids = page_index.get(str(page_num), [])
        seen: set = set()
        result: List[Dict[str, Any]] = []
        for block_id in block_ids:
            for chunk_id in block_to_chunks.get(block_id, []):
                if chunk_id not in seen and chunk_id in chunk_idx:
                    result.append(chunk_idx[chunk_id])
                    seen.add(chunk_id)
        return result

    return [c for c in payload.get("chunks", []) if c.get("page_num") == page_num]


# ── outline tree accessor ─────────────────────────────────────────────────────

def outline_tree(payload: Dict[str, Any]) -> List[Dict[str, Any]]:
    """Return the nested section outline tree."""
    return payload.get("outline") or []


def flat_outline(payload: Dict[str, Any]) -> List[Dict[str, Any]]:
    """
    Flatten the nested outline tree into a list of dicts in document order.
    Each entry: { node_id, title, h_level, depth, chunk_ids }
    """
    result: List[Dict[str, Any]] = []

    def _walk(nodes: List[Dict[str, Any]], depth: int) -> None:
        for node in nodes:
            result.append({
                "node_id":   node["node_id"],
                "title":     node["title"],
                "h_level":   node["h_level"],
                "depth":     depth,
                "chunk_ids": node["chunk_ids"],
            })
            _walk(node.get("children", []), depth + 1)

    _walk(outline_tree(payload), 0)
    return result


# ── breadcrumb helpers ────────────────────────────────────────────────────────

def breadcrumb_for_chunk(
    payload: Dict[str, Any],
    chunk_id: str,
) -> List[str]:
    """
    Return the outline_path (ancestor title breadcrumb) for a specific chunk.

    Checks `chunk_outline_path` index first, falls back to the `outline_path`
    field on the chunk dict itself.
    """
    path = (payload.get("chunk_outline_path") or {}).get(chunk_id)
    if path is not None:
        return path

    chunk_idx = _chunk_index(payload)
    chunk = chunk_idx.get(chunk_id)
    if chunk:
        return chunk.get("outline_path") or []

    return []


# ── embedding text builder ────────────────────────────────────────────────────

def embed_text_for_chunk(
    chunk: Dict[str, Any],
    payload: Optional[Dict[str, Any]] = None,
) -> str:
    """
    Build the recommended string to feed to your embedding model for a chunk.

    Fixes applied (see module docstring for full rationale):
      1. [Context: ...] prefix is stripped before embedding.
      2. Full outline breadcrumb is prepended as "[h1 > h2 > h3]".
      3. Table chunks are converted to Markdown (or HTML) for column structure.
      4. Figure chunks embed only their caption/description text, not OCR noise.

    Format (paragraph / title):
        [Section > Subsection]
        <clean chunk text>

    Format (table):
        [Section > Subsection]
        | col1 | col2 | col3 |
        | --- | --- | --- |
        | val | val | val |

    Format (figure):
        [Section > Subsection]
        Figure: <caption text>
    """
    region = chunk.get("region_type", "paragraph")

    # ── 1. resolve breadcrumb ─────────────────────────────────────────────────
    if payload is not None:
        path = breadcrumb_for_chunk(payload, chunk["chunk_id"])
    else:
        path = chunk.get("outline_path") or []

    # Fall back to section_title alone when no outline path exists
    if not path:
        section = (chunk.get("section_title") or "").strip()
        path = [section] if section else []

    breadcrumb_str = " > ".join(path) if path else ""

    # ── 2. build raw text for this region type ────────────────────────────────

    if region == "table":
        # Prefer structured HTML (preserves rows/columns perfectly for the LLM).
        # text is already clean (no [Context:] prefix — Space 1 FIX #1).
        if chunk.get("table_html"):
            body = chunk["table_html"]
        else:
            # Convert OCR'd flat text to Markdown pipe-table.
            # strip_context_prefix is a no-op on current payloads but kept as guard.
            raw = strip_context_prefix(chunk.get("text") or "")
            body = table_text_to_markdown(raw)

    elif region == "figure":
        raw = strip_context_prefix(chunk.get("text") or "")
        # Space 1 FIX #7/#8: chart OCR noise is already replaced with a semantic
        # placeholder like "[Figure on page N — chart or diagram ...]" by the
        # chunker. Detect that and embed it directly; only run caption extraction
        # on genuine figure OCR text.
        if raw.startswith("[Figure on page"):
            body = raw
        else:
            caption = _extract_figure_caption(raw)
            body = f"Figure: {caption}" if caption else f"Figure (no caption): {raw[:120]}"

    else:
        # paragraph / title / unknown
        # text is already clean — strip_context_prefix is a no-op guard only.
        body = strip_context_prefix(chunk.get("text") or "")

    # ── 3. assemble final string ──────────────────────────────────────────────
    if breadcrumb_str:
        return f"[{breadcrumb_str}]\n{body}"

    return body


def _extract_figure_caption(ocr_text: str) -> str:
    """
    Extract the human-readable caption from a figure OCR block.

    Figure OCR blocks typically contain:
      - Axis labels, tick values, legend text (noise)
      - A caption line beginning with "Figure", "Fig.", "FIG", or a digit

    Strategy: scan lines for a caption-pattern match; if none is found return
    the full text (better than nothing).
    """
    _CAPTION_RE = re.compile(
        r'^(fig(?:ure)?\.?\s*\d+|figure\s+\d+|\d+\s*[:\.])',
        re.IGNORECASE,
    )
    lines = [l.strip() for l in ocr_text.splitlines() if l.strip()]
    for line in lines:
        if _CAPTION_RE.match(line):
            return line
    # No structured caption found — return all lines that look like prose
    prose = " ".join(
        l for l in lines
        if len(l) > 20 and not re.match(r'^[\d\s\.\-\+]+$', l)
    )
    return prose or ocr_text.strip()