rag-from-scratch/chunks/segmenter.py

from __future__ import annotations
from markdown_it.token import Token
from chunks.models import Block
from chunks.config import ChunkerConfig

_COUNTER = 0


def _reset_counter() -> None:
    global _COUNTER
    _COUNTER = 0


def _make_id() -> str:
    global _COUNTER
    _COUNTER += 1
    return f"blk_{_COUNTER:04d}"


def _inline_to_plain(token: Token) -> str:
    if token.children:
        parts = []
        for c in token.children:
            if c.type in ("text", "code_inline"):
                parts.append(c.content)
            elif c.type == "softbreak":
                parts.append(" ")
        return "".join(parts).strip()
    return (token.content or "").strip()


def _lines_content(lines: list[str], start: int, end: int) -> str:
    return "\n".join(lines[start:end]).strip()


def _current_path(stack: dict[int, str], depth: int) -> list[dict]:
    path = []
    for level in sorted(stack.keys()):
        path.append({"level": level, "text": stack[level]})
    return path[:depth]


def _is_skip_heading(text: str, skip_set: set[str]) -> bool:
    t = text.lower().strip()
    return any(t == s or t.startswith(s) for s in {h.lower() for h in skip_set})


def _find_close(tokens: list[Token], start: int, open_type: str, close_type: str) -> int:
    """Ritorna indice del token close_type corrispondente a tokens[start]."""
    depth = 1
    i = start + 1
    while i < len(tokens) and depth > 0:
        if tokens[i].type == open_type:
            depth += 1
        elif tokens[i].type == close_type:
            depth -= 1
        i += 1
    return i - 1


def segment(tokens: list[Token], lines: list[str], config: ChunkerConfig) -> list[Block]:
    _reset_counter()
    blocks: list[Block] = []
    heading_stack: dict[int, str] = {}
    pre_heading_done = False
    skip_mode = False
    skip_level: int | None = None

    i = 0
    while i < len(tokens):
        tok = tokens[i]

        # ── Heading ──────────────────────────────────────────────────────────
        if tok.type == "heading_open":
            level = int(tok.tag[1])
            inline = tokens[i + 1] if i + 1 < len(tokens) else None
            text = (inline.content if inline and inline.type == "inline" else "").strip()

            if _is_skip_heading(text, config.skip_headings):
                skip_mode = True
                skip_level = level
            elif skip_mode and skip_level is not None and level <= skip_level:
                skip_mode = False
                skip_level = None

            for lvl in [l for l in list(heading_stack.keys()) if l >= level]:
                del heading_stack[lvl]
            heading_stack[level] = text
            pre_heading_done = True
            i += 3  # heading_open, inline, heading_close
            continue

        # ── Skip pre-heading ─────────────────────────────────────────────────
        if config.skip_pre_heading and not pre_heading_done:
            i += 1
            continue

        # ── Skip section ─────────────────────────────────────────────────────
        if skip_mode:
            i += 1
            continue

        header_path = _current_path(heading_stack, config.context_depth)

        # ── Paragraph ────────────────────────────────────────────────────────
        if tok.type == "paragraph_open":
            inline = tokens[i + 1] if i + 1 < len(tokens) else None
            if tok.map:
                start, end = tok.map
                content = _lines_content(lines, start, end)
                plain = _inline_to_plain(inline) if inline and inline.type == "inline" else content
                blocks.append(Block(
                    id=_make_id(),
                    kind="paragraph",
                    content=content,
                    plain_text=plain,
                    atomic=False,
                    start_line=start,
                    end_line=end,
                    header_path=header_path,
                    chars=len(content),
                ))
            i += 3  # paragraph_open, inline, paragraph_close
            continue

        # ── Code fence ───────────────────────────────────────────────────────
        if tok.type == "fence":
            if tok.map:
                start, end = tok.map
                content = _lines_content(lines, start, end)
                plain = f"[codice {tok.info.strip()}]" if tok.info.strip() else "[codice]"
                blocks.append(Block(
                    id=_make_id(),
                    kind="code",
                    content=content,
                    plain_text=plain,
                    atomic=True,
                    start_line=start,
                    end_line=end,
                    header_path=header_path,
                    chars=len(content),
                ))
            i += 1
            continue

        # ── Container blocks ─────────────────────────────────────────────────
        _CONTAINERS = {
            "table_open":        ("table_close",        "table",      True),
            "bullet_list_open":  ("bullet_list_close",  "list",       True),
            "ordered_list_open": ("ordered_list_close", "list",       True),
            "blockquote_open":   ("blockquote_close",   "blockquote", False),
        }
        if tok.type in _CONTAINERS:
            close_type, kind, atomic = _CONTAINERS[tok.type]
            atomic = atomic or kind in config.atomic_types
            close_idx = _find_close(tokens, i, tok.type, close_type)
            close_tok = tokens[close_idx]

            start = tok.map[0] if tok.map else 0
            end = (close_tok.map[1] if close_tok.map else None) or (tok.map[1] if tok.map else start + 1)
            content = _lines_content(lines, start, end)
            blocks.append(Block(
                id=_make_id(),
                kind=kind,
                content=content,
                plain_text=content,
                atomic=atomic,
                start_line=start,
                end_line=end,
                header_path=header_path,
                chars=len(content),
            ))
            i = close_idx + 1
            continue

        # ── Thematic break ────────────────────────────────────────────────────
        if tok.type == "hr":
            start = tok.map[0] if tok.map else 0
            end = tok.map[1] if tok.map else start + 1
            blocks.append(Block(
                id=_make_id(),
                kind="thematic_break",
                content="---",
                plain_text="",
                atomic=False,
                start_line=start,
                end_line=end,
                header_path=header_path,
                chars=3,
            ))
            i += 1
            continue

        # ── Math block (dollarmath plugin) ───────────────────────────────────
        if tok.type == "math_block":
            start = tok.map[0] if tok.map else 0
            end = tok.map[1] if tok.map else start + 1
            content = _lines_content(lines, start, end)
            blocks.append(Block(
                id=_make_id(),
                kind="math",
                content=content,
                plain_text="[formula]",
                atomic=True,
                start_line=start,
                end_line=end,
                header_path=header_path,
                chars=len(content),
            ))
            i += 1
            continue

        # ── HTML block ────────────────────────────────────────────────────────
        if tok.type == "html_block":
            start = tok.map[0] if tok.map else 0
            end = tok.map[1] if tok.map else start + 1
            content = tok.content.strip()
            blocks.append(Block(
                id=_make_id(),
                kind="html",
                content=content,
                plain_text="",
                atomic="html" in config.atomic_types,
                start_line=start,
                end_line=end,
                header_path=header_path,
                chars=len(content),
            ))
            i += 1
            continue

        i += 1

    return blocks