rag-from-scratch/chunks/chunker.py

#!/usr/bin/env python3
"""
chunker.py — CLI + orchestrazione pipeline AST-based

Uso:
    python chunks/chunker.py --stem <stem>
    python chunks/chunker.py                        # tutti gli stem in sources/
    python chunks/chunker.py --stem <stem> --force
"""
from __future__ import annotations
import argparse
import json
import sys
from dataclasses import asdict
from datetime import datetime, timezone
from pathlib import Path

_HERE = Path(__file__).resolve().parent
_ROOT = _HERE.parent
if str(_ROOT) not in sys.path:
    sys.path.insert(0, str(_ROOT))

from chunks.config import ChunkerConfig
from chunks.parser import parse
from chunks.segmenter import segment
from chunks.packer import pack
from chunks.validator import validate
from chunks.models import ChunkingResult, Diagnostics


def find_source(stem: str, root: Path) -> Path | None:
    candidates = [
        root / "sources" / f"{stem}_output" / "auto" / f"{stem}.md",
        root / "sources" / f"{stem}.md",
    ]
    for p in candidates:
        if p.exists():
            return p
    return None


def run_pipeline(stem: str, root: Path = _ROOT,
                 config: ChunkerConfig | None = None,
                 force: bool = False) -> ChunkingResult:
    config = config or ChunkerConfig()
    out_dir = root / "chunks" / stem
    chunks_path = out_dir / "chunks.json"

    if chunks_path.exists() and not force:
        print(f"[{stem}] esiste già — usa --force per rigenerare. Skip.")
        return ChunkingResult(stem=stem, source_path="", chunks=[],
                              diagnostics=Diagnostics([], [], {}))

    source_path = find_source(stem, root)
    if source_path is None:
        print(f"[{stem}] sorgente non trovata in sources/. Saltato.", file=sys.stderr)
        return ChunkingResult(stem=stem, source_path="", chunks=[],
                              diagnostics=Diagnostics(
                                  errors=[f"sorgente non trovata per stem '{stem}'"],
                                  warnings=[], metrics={}))

    source = source_path.read_text(encoding="utf-8")
    tokens, lines = parse(source)
    blocks = segment(tokens, lines, config)
    chunks = pack(blocks, config, stem)

    result = ChunkingResult(
        stem=stem,
        source_path=str(source_path.relative_to(root)),
        chunks=chunks,
        diagnostics=Diagnostics([], [], {}),
    )
    result.diagnostics = validate(result, source, config)

    out_dir.mkdir(parents=True, exist_ok=True)

    chunks_path.write_text(
        json.dumps([asdict(c) for c in chunks], ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    (out_dir / "meta.json").write_text(
        json.dumps({
            "stem": stem,
            "source_path": result.source_path,
            "total_chunks": len(chunks),
            "total_chars": sum(c.chars for c in chunks),
            "created_at": datetime.now(timezone.utc).isoformat(),
            "config": {
                "max_chars": config.max_chars,
                "min_chars": config.min_chars,
                "target_chars": config.target_chars,
                "context_depth": config.context_depth,
            },
        }, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    (out_dir / "report.json").write_text(
        json.dumps(asdict(result.diagnostics), ensure_ascii=False, indent=2),
        encoding="utf-8",
    )

    n = len(chunks)
    e = len(result.diagnostics.errors)
    w = len(result.diagnostics.warnings)
    print(f"[{stem}] {n} chunk  |  errors={e}  warnings={w}  →  {out_dir}")
    return result


def _discover_stems(root: Path) -> list[str]:
    sources = root / "sources"
    stems: set[str] = set()
    if not sources.exists():
        return []
    for p in sources.iterdir():
        if p.is_dir() and p.name.endswith("_output"):
            stem = p.name[: -len("_output")]
            if (p / "auto" / f"{stem}.md").exists():
                stems.add(stem)
        elif p.is_file() and p.suffix == ".md":
            stems.add(p.stem)
    return sorted(stems)


def main() -> None:
    parser = argparse.ArgumentParser(description="Chunker Markdown AST-based")
    parser.add_argument("--stem", help="Stem documento (es. valute-virtuali)")
    parser.add_argument("--force", action="store_true", help="Rigenera anche se già presente")
    args = parser.parse_args()

    stems = [args.stem] if args.stem else _discover_stems(_ROOT)
    if not stems:
        print("Nessun sorgente trovato in sources/.", file=sys.stderr)
        sys.exit(1)

    for stem in stems:
        run_pipeline(stem=stem, force=args.force)


if __name__ == "__main__":
    main()