rag-from-scratch/chunks/chunker.py

#!/usr/bin/env python3
"""
chunker.py — Chunking semantico da Markdown pulito

Input:  sources/<stem>.md  — Markdown già strutturato (H1/H2/H3 + paragrafi)
Output: chunks/<stem>/chunks.json
        chunks/<stem>/meta.json

Regole:
  - ogni paragrafo diventa un chunk; paragrafi di contesti diversi non si mescolano
  - se un paragrafo supera MAX_CHARS viene spezzato a confine di frase (mai a metà)
  - paragrafi brevi (< MIN_CHARS) vengono fusi col successivo finché non raggiungono
    MIN_CHARS, a patto che abbiano lo stesso contesto heading
  - tabelle, liste e codice sono blocchi atomici (non si spezzano)

Uso:
    python chunks/chunker.py --stem <stem>
    python chunks/chunker.py                  # tutti gli stem con .md in sources/
    python chunks/chunker.py --stem <stem> --force
"""

import argparse
import json
import re
import sys
from pathlib import Path

_HERE = Path(__file__).resolve().parent
if str(_HERE) not in sys.path:
    sys.path.insert(0, str(_HERE))
import config as cfg


# ─── Utilità ──────────────────────────────────────────────────────────────────

def split_sentences(text: str) -> list[str]:
    parts = re.split(cfg.SENTENCE_SPLIT_RE, text.strip())
    return [p.strip() for p in parts if p.strip()]


def context_to_meta(context: str) -> tuple[str, str]:
    parts = [p.strip() for p in context.split(" > ") if p.strip()]
    if len(parts) >= 2:
        return " > ".join(parts[:-1]), parts[-1]
    return (parts[0] if parts else ""), ""


# ─── Parser Markdown ──────────────────────────────────────────────────────────

_SKIP_HEADINGS_LOWER = {h.lower() for h in cfg.SKIP_HEADINGS}


def _is_skip_heading(title: str) -> bool:
    t = title.lower()
    return any(t == s or t.startswith(s) for s in _SKIP_HEADINGS_LOWER)


def parse_paragraphs(text: str) -> list[dict]:
    """Estrae blocchi dal Markdown con il loro contesto heading.

    Restituisce: [{"context": "H1 > H2 > H3", "text": "...", "kind": "text|table|list|code"}]

    - Heading senza corpo non emettono chunk: aggiornano solo il contesto.
    - Tabelle (righe |), liste (righe -/*) e code block (```) sono atomici.
    - Sezioni in SKIP_HEADINGS vengono saltate completamente.
    - Contenuto pre-heading saltato se SKIP_PRE_HEADING è True.
    """
    headings  = ["", "", ""]   # H1, H2, H3
    result:   list[dict] = []
    buf:      list[str]  = []
    cur_kind  = "text"
    in_code   = False
    skip_level: int | None = None   # livello heading che ha attivato lo skip

    def current_context() -> str:
        parts = [h for h in headings if h]
        return " > ".join(parts[:cfg.CONTEXT_DEPTH]) if parts else "documento"

    def is_skipping() -> bool:
        return skip_level is not None

    def flush() -> None:
        if is_skipping():
            buf.clear()
            return
        if cfg.SKIP_PRE_HEADING and current_context() == "documento":
            buf.clear()
            return
        body = "\n".join(buf).strip()
        if body:
            result.append({"context": current_context(), "text": body, "kind": cur_kind})
        buf.clear()

    for line in text.splitlines():
        # ── Code block toggle ─────────────────────────────────────────────────
        if line.strip().startswith("```"):
            if not in_code:
                flush()
                cur_kind = "code"
                in_code  = True
                if not is_skipping():
                    buf.append(line)
            else:
                if not is_skipping():
                    buf.append(line)
                in_code = False
                flush()
                cur_kind = "text"
            continue

        if in_code:
            if not is_skipping():
                buf.append(line)
            continue

        # ── Heading ───────────────────────────────────────────────────────────
        m = re.match(r"^(#{1,3}) (.+)", line)
        if m:
            flush()
            level = len(m.group(1))
            title = m.group(2).strip()

            # chiudi skip se torniamo a livello pari o superiore
            if skip_level is not None and level <= skip_level:
                skip_level = None

            headings[level - 1] = title
            for i in range(level, 3):
                headings[i] = ""
            cur_kind = "text"

            # apri skip se questo heading è nella lista
            if _is_skip_heading(title):
                skip_level = level

            continue

        # ── Tabella ───────────────────────────────────────────────────────────
        if line.strip().startswith("|"):
            if cur_kind != "table":
                flush()
                cur_kind = "table"
            buf.append(line)
            continue

        # ── Lista ─────────────────────────────────────────────────────────────
        if re.match(r"^\s*[-*]\s", line):
            if cur_kind != "list":
                flush()
                cur_kind = "list"
            buf.append(line)
            continue

        # ── Riga vuota: chiude il paragrafo corrente ──────────────────────────
        if line.strip() == "":
            flush()
            cur_kind = "text"
            continue

        # ── Testo normale ─────────────────────────────────────────────────────
        if cur_kind in ("table", "list", "code"):
            flush()
            cur_kind = "text"
        buf.append(line)

    flush()
    return result


# ─── Merge paragrafi brevi ────────────────────────────────────────────────────

def merge_short(paragraphs: list[dict]) -> list[dict]:
    """Fonde paragrafi di testo consecutivi sotto MIN_CHARS con il successivo,
    purché condividano lo stesso contesto heading."""
    if not cfg.MERGE_SHORT_PARAGRAPHS:
        return paragraphs

    result:   list[dict] = []
    buf_para: dict | None = None
    buf_text: str         = ""

    def flush_buf() -> None:
        nonlocal buf_para, buf_text
        if buf_para is not None:
            result.append({**buf_para, "text": buf_text})
            buf_para = None
            buf_text = ""

    for para in paragraphs:
        if para["kind"] != "text":
            flush_buf()
            result.append(para)
            continue

        if buf_para is None:
            buf_para = para
            buf_text = para["text"]
        elif buf_para["context"] == para["context"] and len(buf_text) < cfg.MIN_CHARS:
            buf_text = buf_text + "\n\n" + para["text"]
        else:
            flush_buf()
            buf_para = para
            buf_text = para["text"]

    flush_buf()
    return result


# ─── Chunking ─────────────────────────────────────────────────────────────────

def make_chunks(paragraphs: list[dict]) -> list[dict]:
    """Genera chunk da parse_paragraphs (dopo eventuale merge).

    - Blocchi atomici (table, list, code): un chunk, mai spezzato.
    - Testo: un chunk per paragrafo; se supera MAX_CHARS, spezza a confine frase.
    """
    chunks: list[dict] = []
    idx    = 0

    for para in paragraphs:
        text    = para["text"]
        context = para["context"]
        kind    = para["kind"]
        sezione, titolo = context_to_meta(context)

        def emit(body: str) -> None:
            nonlocal idx
            chunk_text = f"[{context}]\n{body}"
            chunks.append({
                "chunk_id": f"c{idx}",
                "text":     chunk_text,
                "sezione":  sezione,
                "titolo":   titolo,
                "context":  context,
                "n_chars":  len(chunk_text),
            })
            idx += 1

        # ── Atomici ───────────────────────────────────────────────────────────
        if kind in ("table", "list", "code"):
            emit(text)
            continue

        # ── Testo: split a confine di frase se supera MAX_CHARS ───────────────
        sents = split_sentences(text)
        if not sents:
            continue

        current: list[str] = []
        for sent in sents:
            projected = len(" ".join(current + [sent]))
            if projected <= cfg.MAX_CHARS or not current:
                current.append(sent)
            else:
                emit(" ".join(current))
                current = [sent]

        if current:
            emit(" ".join(current))

    return chunks


# ─── Merge frasi spezzate ─────────────────────────────────────────────────────

_SENT_END = re.compile(
    "[.!?;:"
    + chr(0xBB) + ")\\\\"
    + chr(0x2018) + chr(0x2019)
    + chr(0x201C) + chr(0x201D)
    + "\"'"
    + chr(0x2014) + chr(0x2013) + chr(0x2026) + chr(0xB7)
    + "]$"
    + r"|\d[\d.,/]*$"   # numero o versione
    + r"|\$$"           # formula LaTeX inline
    + r"|\}$"           # blocco LaTeX
    + r"|>$"            # tag HTML
    + r"|\\\\$"         # \\ LaTeX
    + r"|\|$"           # riga tabella
)


def _body(chunk: dict) -> str:
    text = chunk["text"]
    nl = text.find("\n")
    return text[nl + 1:] if nl != -1 else text


def merge_broken_sentences(chunks: list[dict]) -> list[dict]:
    """Fonde chunk consecutivi con lo stesso contesto quando il primo termina
    senza punteggiatura di fine frase (frase spezzata dal sorgente)."""
    result: list[dict] = []
    i = 0
    while i < len(chunks):
        c    = dict(chunks[i])
        body = _body(c)

        while (
            i + 1 < len(chunks)
            and chunks[i + 1]["context"] == c["context"]
            and not _SENT_END.search(body.rstrip())
        ):
            i += 1
            next_body = _body(chunks[i])
            body      = body.rstrip() + " " + next_body.lstrip()

        chunk_text = f"[{c['context']}]\n{body}"
        c["text"]    = chunk_text
        c["n_chars"] = len(chunk_text)
        result.append(c)
        i += 1

    for idx, c in enumerate(result):
        c["chunk_id"] = f"c{idx}"

    return result


# ─── Pipeline per documento ───────────────────────────────────────────────────

def process_stem(stem: str, project_root: Path, force: bool) -> bool:
    md_path = project_root / "sources" / f"{stem}_output" / "auto" / f"{stem}.md"
    if not md_path.exists():
        md_path = project_root / "sources" / f"{stem}.md"
    if not md_path.exists():
        print(f"  ✗ {stem}.md non trovato (cercato in sources/{stem}_output/auto/ e sources/)")
        return False

    out_dir  = project_root / "chunks" / stem
    out_file = out_dir / "chunks.json"

    if out_file.exists() and not force:
        print(f"  ↩  chunks/{stem}/chunks.json già presente — skip (usa --force per rigenerare)")
        return True

    print(f"[chunker] {stem}")

    text       = md_path.read_text(encoding="utf-8")
    paragraphs = parse_paragraphs(text)

    if not paragraphs:
        print(f"  ✗ Nessun paragrafo estratto da {md_path.name}")
        return False

    if cfg.MERGE_SHORT_PARAGRAPHS:
        paragraphs = merge_short(paragraphs)

    chunks = make_chunks(paragraphs)
    chunks = merge_broken_sentences(chunks)

    if not chunks:
        print(f"  ✗ Nessun chunk generato")
        return False

    out_dir.mkdir(parents=True, exist_ok=True)
    out_file.write_text(
        json.dumps(chunks, ensure_ascii=False, indent=2), encoding="utf-8"
    )
    (out_dir / "meta.json").write_text(
        json.dumps({
            "stem":        stem,
            "source":      str(md_path.relative_to(project_root)),
            "max_chars":   cfg.MAX_CHARS,
            "min_chars":   cfg.MIN_CHARS,
            "merge_short": cfg.MERGE_SHORT_PARAGRAPHS,
            "strategy":    "one_paragraph_per_chunk",
        }, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )

    lengths   = [c["n_chars"] for c in chunks]
    over_max  = sum(1 for l in lengths if l > cfg.MAX_CHARS)
    under_min = sum(1 for l in lengths if l < cfg.MIN_CHARS)
    avg       = int(sum(lengths) / len(lengths))

    print(f"  ✅ {len(chunks)} chunk  |  media {avg} char  |  max {max(lengths)} char")
    if over_max:
        print(f"  ⚠️  {over_max} chunk superano MAX_CHARS={cfg.MAX_CHARS}")
    if under_min:
        print(f"  ℹ️  {under_min} chunk sotto MIN_CHARS={cfg.MIN_CHARS}")
    print(f"  → chunks/{stem}/chunks.json")
    return True


# ─── Entry point ──────────────────────────────────────────────────────────────

if __name__ == "__main__":
    project_root = Path(__file__).parent.parent
    sources_dir  = project_root / "sources"

    parser = argparse.ArgumentParser(description="Markdown pulito → chunks.json")
    parser.add_argument("--stem",  help="Nome documento (es. analisi2)")
    parser.add_argument("--force", action="store_true",
                        help="Rigenera chunks.json anche se già presente")
    args = parser.parse_args()

    if args.stem:
        stems = [args.stem]
    else:
        found = set()
        for p in sources_dir.glob("*_output/auto/*.md"):
            found.add(p.stem)
        for p in sources_dir.glob("*.md"):
            found.add(p.stem)
        stems = sorted(found)
        if not stems:
            print("Nessun file .md trovato in sources/")
            sys.exit(1)

    results = [process_stem(s, project_root, args.force) for s in stems]
    ok = sum(results)
    print(f"\n{'✅' if all(results) else '⚠️ '} {ok}/{len(results)} documenti processati")
    sys.exit(0 if all(results) else 1)