hkuds--lightrag/lightrag/chunk_schema.py

"""Chunk schema helpers shared across the chunking + extraction pipeline.

Three responsibilities live here so chunker implementations and the pipeline
both consume identical normalization rules:

- :func:`normalize_chunk_heading` collapses the legacy flat
  ``heading``/``parent_headings``/``level`` triple and the new nested form
  into the canonical ``{"level", "heading", "parent_headings"}`` dict.
- :func:`normalize_chunk_sidecar` validates the new ``sidecar`` payload and
  ensures ``refs`` is always present as a list (single-source items may omit
  it before normalization; we materialize a single-element list for the
  storage layer).
- :func:`strip_internal_multimodal_markup_for_extraction` rewrites
  ``<cite>`` / ``<drawing>`` / ``<equation>`` markup so the entity-extraction
  LLM sees a clean text body. The original ``chunk["content"]`` is never
  mutated; the cleaned string is only used to build the extraction prompt.

The clean function is intentionally conservative: it only strips
parser-emitted identifier attributes that have no business reaching the LLM
(``id``, ``refid``, ``path``, ``src``). Visible captions and equation bodies
are preserved so the extracted entities can still ground against them.
"""

from __future__ import annotations

import re
import unicodedata
from typing import Any

from lightrag.constants import DEFAULT_HEADING_LEVEL_MAX_CHARS


_SIDECAR_TYPES = frozenset({"block", "drawing", "table", "equation"})

# Separator joining heading levels into a single breadcrumb line. Shared so the
# extraction-side token budgeter (see ``operate._truncate_section_context``) can
# split the rendered breadcrumb back into levels without a drifting magic string.
HEADING_BREADCRUMB_SEP = " → "


def normalize_chunk_heading(dp: dict[str, Any]) -> dict[str, Any] | None:
    """Return the canonical nested heading dict or ``None`` when absent.

    Accepts:

    - ``dp["heading"]`` already a dict ``{"level", "heading", "parent_headings"}``.
    - Legacy flat fields ``heading: str`` + ``parent_headings: list[str]`` +
      ``level: int``.

    Empty / missing inputs collapse to ``None`` so callers can simply omit
    the field when writing the chunk record.
    """
    nested = dp.get("heading")
    if isinstance(nested, dict):
        heading_text = str(nested.get("heading") or "").strip()
        parents_raw = nested.get("parent_headings") or []
        level_raw = nested.get("level", 0)
    else:
        heading_text = str(nested or "").strip()
        parents_raw = dp.get("parent_headings") or []
        level_raw = dp.get("level", 0)

    parent_headings: list[str] = []
    if isinstance(parents_raw, list):
        for entry in parents_raw:
            text = str(entry or "").strip()
            if text:
                parent_headings.append(text)

    try:
        level = int(level_raw or 0)
    except (TypeError, ValueError):
        level = 0

    if not heading_text and not parent_headings and level == 0:
        return None

    return {
        "level": level,
        "heading": heading_text,
        "parent_headings": parent_headings,
    }


_HEADING_WHITESPACE_RE = re.compile(r"\s+")
# Unicode categories stripped from a heading: control (Cc) and format (Cf)
# chars — NULs, zero-width marks (ZWSP/ZWNJ/ZWJ/WORD JOINER/BOM are all Cf),
# directional/format codes — that only add token noise to the LLM prompt.
_HEADING_STRIP_CATEGORIES = frozenset({"Cc", "Cf"})
# The only Cc chars that ``\s`` folds into a space; keep them through the strip
# pass so they become a space below instead of gluing adjacent words. (Cannot
# use ``str.isspace()`` for this: Python treats \x1c-\x1f as whitespace but
# ``\s`` does not, so those would survive the strip — an explicit set matches
# the regex.)
_HEADING_KEEP_WS = frozenset("\t\n\r\f\v")


def _clean_heading_text(text: str) -> str:
    """Flatten a heading into one clean line for the LLM.

    Converts the breadcrumb separator ``→`` to a space, drops every Unicode
    control (Cc) / format (Cf) char (zero-width marks, NULs, directional/format
    codes), and collapses every run of whitespace (tab, newline, NBSP,
    full-width / ideographic space, ...) into a single regular space as the
    final step. Normal CJK, Latin, digits, and punctuation are left untouched,
    and no spaces are inserted between adjacent CJK characters, so it is safe
    for Chinese headings.
    """
    # ``→`` (U+2192, the breadcrumb separator char) must never survive inside a
    # single heading, or it would forge an extra level when the breadcrumb is
    # split back on " → " (see operate._truncate_section_context).
    text = text.replace("→", " ")
    text = "".join(
        ch
        for ch in text
        if ch in _HEADING_KEEP_WS
        or unicodedata.category(ch) not in _HEADING_STRIP_CATEGORIES
    )
    # Collapse LAST so the spaces introduced above (→, kept whitespace controls,
    # and any gap left around a removed control char) all fold into one space.
    text = _HEADING_WHITESPACE_RE.sub(" ", text)
    return text.strip()


def _truncate_heading_level(text: str, max_chars: int) -> str:
    """Hard-cap a single heading level, marking elision with an ellipsis."""
    if max_chars <= 0 or len(text) <= max_chars:
        return text
    # Reserve one char for the ellipsis so the result length stays <= max_chars.
    return text[: max_chars - 1].rstrip() + "…"


def _clean_and_cap_headings(headings: list[str], max_heading_len: int) -> list[str]:
    """Clean each heading, drop empties, then hard-cap each level's length.

    Shared by :func:`format_parent_headings` and :func:`format_heading_context`
    so both the query-stage and extraction-stage breadcrumbs apply identical
    cleaning (:func:`_clean_heading_text`) and per-level truncation, and cannot
    drift apart. Order matters: clean → drop empties → cap.
    """
    return [
        _truncate_heading_level(c, max_heading_len)
        for c in (_clean_heading_text(h) for h in headings)
        if c
    ]


def format_parent_headings(
    dp: dict[str, Any],
    *,
    max_heading_len: int = DEFAULT_HEADING_LEVEL_MAX_CHARS,
) -> str:
    """Join a chunk's parent heading chain into ``h1 → h2 → h3``.

    Reuses :func:`normalize_chunk_heading` so both the nested and legacy flat
    heading shapes are handled, then cleans each heading via
    :func:`_clean_heading_text` so the string sent to the LLM is a single tidy
    line. Returns an empty string when the chunk has no (non-empty) parent
    headings, so callers can simply omit the field when it is empty.

    Each individual heading level is capped at ``max_heading_len`` characters
    (set ``<= 0`` to disable), matching :func:`format_heading_context`, so one
    runaway title cannot bloat the query context before its token truncation.
    """
    normalized = normalize_chunk_heading(dp)
    if not normalized:
        return ""
    cleaned = _clean_and_cap_headings(normalized["parent_headings"], max_heading_len)
    return HEADING_BREADCRUMB_SEP.join(cleaned)


def format_heading_context(
    dp: dict[str, Any],
    *,
    max_heading_len: int = DEFAULT_HEADING_LEVEL_MAX_CHARS,
) -> str:
    """Join a chunk's full heading chain (parents + current) into ``h1 → h2 → h3``.

    Like :func:`format_parent_headings` but appends the chunk's own section
    heading after the parent chain, so the entity-extraction LLM sees the
    complete breadcrumb of the section the input text belongs to. Reuses
    :func:`normalize_chunk_heading` (handles both nested and legacy flat shapes)
    and :func:`_clean_heading_text`. Returns an empty string when the chunk
    carries no (non-empty) heading information, so callers can simply omit the
    field when it is empty.

    Each individual heading level is capped at ``max_heading_len`` characters
    (set ``<= 0`` to disable) so one runaway title cannot bloat the prompt; the
    caller is still responsible for token-budgeting the joined breadcrumb.
    """
    normalized = normalize_chunk_heading(dp)
    if not normalized:
        return ""
    chain = list(normalized["parent_headings"])
    if normalized["heading"]:
        chain.append(normalized["heading"])
    cleaned = _clean_and_cap_headings(chain, max_heading_len)
    return HEADING_BREADCRUMB_SEP.join(cleaned)


def normalize_chunk_sidecar(dp: dict[str, Any]) -> dict[str, Any] | None:
    """Return the canonical sidecar dict or ``None`` when absent / invalid.

    Output shape::

        {"type": <one of block|drawing|table|equation>,
         "id":   <primary source id>,
         "refs": [{"type": ..., "id": ...}, ...]}

    ``refs`` is always materialized as a list with at least the primary id.
    Single-source chunks therefore land in storage with ``refs=[{type,id}]``
    so downstream consumers don't need to special-case the field's presence.
    """
    sidecar = dp.get("sidecar")
    if not isinstance(sidecar, dict):
        return None
    sidecar_type = str(sidecar.get("type") or "").strip()
    sidecar_id = str(sidecar.get("id") or "").strip()
    if sidecar_type not in _SIDECAR_TYPES or not sidecar_id:
        return None

    refs_raw = sidecar.get("refs")
    refs: list[dict[str, str]] = []
    if isinstance(refs_raw, list):
        for entry in refs_raw:
            if not isinstance(entry, dict):
                continue
            ref_type = str(entry.get("type") or "").strip()
            ref_id = str(entry.get("id") or "").strip()
            if ref_type in _SIDECAR_TYPES and ref_id:
                refs.append({"type": ref_type, "id": ref_id})
    if not refs:
        refs = [{"type": sidecar_type, "id": sidecar_id}]

    return {"type": sidecar_type, "id": sidecar_id, "refs": refs}


# `<cite type="..." refid="...">visible text</cite>` → `visible text`.
_CITE_RE = re.compile(
    r"<cite\b[^>]*>(.*?)</cite>",
    flags=re.IGNORECASE | re.DOTALL,
)

# Inner attribute stripper used when the caller wants to *preserve* the
# `<cite type="…">…</cite>` wrapper but drop the parser-internal `refid`.
# Matches ` refid="…"` (leading whitespace + quoted value) so the
# surrounding attribute layout (e.g. `type="table"`) stays intact.
_CITE_REFID_ATTR_RE = re.compile(
    r'\s+refid\s*=\s*"[^"]*"',
    flags=re.IGNORECASE,
)

# Self-closing `<drawing ...>` placeholder.  We keep `caption` (visible) and
# drop `id`, `path`, `src`, `format`, etc.  Tags without any caption are
# removed entirely so they don't pollute extraction input.
_DRAWING_RE = re.compile(
    r"<drawing\b([^>]*)/>",
    flags=re.IGNORECASE,
)

# Container `<equation id="..." format="...">latex</equation>`.  Strip
# identifier attributes; preserve the body and the `format` attribute so
# extraction still sees the equation is a structured element.
_EQUATION_RE = re.compile(
    r"<equation\b([^>]*)>(.*?)</equation>",
    flags=re.IGNORECASE | re.DOTALL,
)

# Container `<table id="tb-..." format="json" caption="...">rows</table>`.
# Native parser emits the internal ``tb-<doc>-NNNN`` identifier here, which
# would otherwise leak into the entity-extraction prompt and become a noisy
# entity. Strip ``id``; keep ``format`` / ``caption`` (and the body verbatim)
# so the extractor still recognizes the element as a structured table.
_TABLE_RE = re.compile(
    r"<table\b([^>]*)>(.*?)</table>",
    flags=re.IGNORECASE | re.DOTALL,
)

# Match attribute pairs like ``caption="text with \"escapes\""``.  We treat
# only the safe identifier-style attributes; complex quoting is rare in
# parser output.
_ATTR_RE = re.compile(
    r'(\w+)\s*=\s*"((?:[^"\\]|\\.)*)"',
)


def _attrs_to_dict(attr_string: str) -> dict[str, str]:
    return {
        match.group(1).lower(): match.group(2)
        for match in _ATTR_RE.finditer(attr_string)
    }


def _format_attrs(pairs: list[tuple[str, str]]) -> str:
    return "".join(f' {k}="{v}"' for k, v in pairs if v)


def _replace_drawing(match: re.Match[str]) -> str:
    attrs = _attrs_to_dict(match.group(1))
    caption = attrs.get("caption", "")
    if not caption.strip():
        return ""
    return f"<drawing{_format_attrs([('caption', caption)])} />"


def _replace_equation(match: re.Match[str]) -> str:
    attrs = _attrs_to_dict(match.group(1))
    body = match.group(2)
    keep: list[tuple[str, str]] = []
    fmt = attrs.get("format", "")
    if fmt:
        keep.append(("format", fmt))
    caption = attrs.get("caption", "")
    if caption.strip():
        keep.append(("caption", caption))
    return f"<equation{_format_attrs(keep)}>{body}</equation>"


def _replace_table(match: re.Match[str]) -> str:
    attrs = _attrs_to_dict(match.group(1))
    body = match.group(2)
    keep: list[tuple[str, str]] = []
    fmt = attrs.get("format", "")
    if fmt:
        keep.append(("format", fmt))
    caption = attrs.get("caption", "")
    if caption.strip():
        keep.append(("caption", caption))
    return f"<table{_format_attrs(keep)}>{body}</table>"


def strip_internal_multimodal_markup_for_extraction(
    content: str, *, keep_cite_tag: bool = False
) -> str:
    """Strip parser-internal identifiers from a chunk content string.

    Only the entity-extraction prompt should receive the cleaned form;
    callers must NOT mutate the stored chunk ``content`` so query-time
    citations still resolve back to the original parser output.

    Transformations always applied:

    - ``<drawing id="im-…" path="…" src="…" caption="Fig 1" />``
        → ``<drawing caption="Fig 1" />``
        (drops the entire tag when no caption is present)
    - ``<table id="tb-…" format="json" caption="…">rows</table>``
        → ``<table format="json" caption="…">rows</table>``
    - ``<equation id="eq-…" format="latex">…</equation>``
        → ``<equation format="latex">…</equation>``

    Cite-tag handling depends on ``keep_cite_tag``:

    - ``keep_cite_tag=False`` (default — entity-extraction path):
      ``<cite type="…" refid="…">Table 1</cite>`` → ``Table 1``.  The
      cite wrapper is dropped so the extractor does not surface it as
      a noisy structural entity.
    - ``keep_cite_tag=True`` (multimodal-analysis surrounding path):
      ``<cite type="table" refid="…">Table 1</cite>`` →
      ``<cite type="table">Table 1</cite>``.  Only the internal
      ``refid`` is removed; the wrapper survives so the VLM/LLM can
      tell visible reference labels (e.g. "Table 1") apart from inline
      prose.
    """
    if not content:
        return content
    if keep_cite_tag:
        cleaned = _CITE_REFID_ATTR_RE.sub("", content)
    else:
        cleaned = _CITE_RE.sub(lambda m: m.group(1), content)
    cleaned = _DRAWING_RE.sub(_replace_drawing, cleaned)
    cleaned = _TABLE_RE.sub(_replace_table, cleaned)
    cleaned = _EQUATION_RE.sub(_replace_equation, cleaned)
    return cleaned


__all__ = [
    "normalize_chunk_heading",
    "format_parent_headings",
    "format_heading_context",
    "HEADING_BREADCRUMB_SEP",
    "normalize_chunk_sidecar",
    "strip_internal_multimodal_markup_for_extraction",
]