shira-hermes/api/services/kb/chunker.py

"""Text chunking — strategies per source kind.

All strategies return list[dict] with keys:
    chunk_index, heading_path, section_ref, content, token_count
"""

from __future__ import annotations

import re
from dataclasses import dataclass
from typing import Literal

# Approx tokens per char for Hebrew/mixed text (voyage-multilingual-2 tokenizer).
_TOKENS_PER_CHAR = 0.35

# Hard safety ceiling per chunk. Voyage accepts up to 32K tokens per item,
# but a single batch is capped at 120K tokens. Keep chunks well below both.
_MAX_TOKENS_PER_CHUNK = 2000

# A "section" in Israeli statutes: "סעיף 12", "תקנה 5", "5א.", "12(ב)"
_SECTION_RE = re.compile(
    r"(?m)^\s*(?:סעיף\s+|תקנה\s+|§\s*)?(\d+[א-ת]?\.?)\s*(?:\(([א-ת\d]+)\))?\s*[\.\:]",
)

# Header in circulars — markdown-ish headings, or numbered "נושא: ..." / "מס' N. ..."
_HEADER_RE = re.compile(
    r"(?m)^\s*(?:#{1,4}\s+|מס['׳]\s*\d+\s*[\.\:]\s*|[A-Zא-ת][^\n]{0,80}\:\s*$)",
)


@dataclass
class Chunk:
    chunk_index: int
    heading_path: str
    section_ref: str
    content: str
    token_count: int


def _count_tokens(text: str) -> int:
    return max(1, int(len(text) * _TOKENS_PER_CHAR))


def _as_dicts(chunks: list[Chunk]) -> list[dict]:
    return [
        {
            "chunk_index": c.chunk_index,
            "heading_path": c.heading_path or None,
            "section_ref": c.section_ref or None,
            "content": c.content,
            "token_count": c.token_count,
        }
        for c in chunks
    ]


def _split_oversized(chunks: list[Chunk]) -> list[Chunk]:
    """Split any chunk whose token_count exceeds the ceiling into fixed-size pieces.

    Preserves heading_path/section_ref on each piece and re-indexes.
    """
    max_chars = int(_MAX_TOKENS_PER_CHUNK / _TOKENS_PER_CHAR)
    result: list[Chunk] = []
    for c in chunks:
        if c.token_count <= _MAX_TOKENS_PER_CHUNK:
            c.chunk_index = len(result)
            result.append(c)
            continue
        text = c.content
        for start in range(0, len(text), max_chars):
            piece = text[start : start + max_chars]
            result.append(
                Chunk(
                    chunk_index=len(result),
                    heading_path=c.heading_path,
                    section_ref=c.section_ref,
                    content=piece,
                    token_count=_count_tokens(piece),
                )
            )
    return result


def chunk_statute(text: str) -> list[dict]:
    """Law / regulation: one chunk per section. Preserves section numbering."""
    text = text.strip()
    if not text:
        return []

    matches = list(_SECTION_RE.finditer(text))
    if not matches:
        # Fallback: treat the whole thing as one chunk.
        return _as_dicts(
            [Chunk(0, "", "", text, _count_tokens(text))]
        )

    chunks: list[Chunk] = []
    for i, m in enumerate(matches):
        start = m.start()
        end = matches[i + 1].start() if i + 1 < len(matches) else len(text)
        content = text[start:end].strip()
        if not content:
            continue
        number = m.group(1).rstrip(".")
        sub = m.group(2)
        section_ref = f"ס' {number}" + (f"({sub})" if sub else "")
        chunks.append(
            Chunk(
                chunk_index=len(chunks),
                heading_path=section_ref,
                section_ref=section_ref,
                content=content,
                token_count=_count_tokens(content),
            )
        )
    return _as_dicts(chunks)


def chunk_circular(
    text: str,
    target_tokens: int = 600,
    overlap_tokens: int = 80,
) -> list[dict]:
    """Circular: header-aware split, then pack paragraphs up to target size with overlap."""
    text = text.strip()
    if not text:
        return []

    target_chars = int(target_tokens / _TOKENS_PER_CHAR)
    overlap_chars = int(overlap_tokens / _TOKENS_PER_CHAR)

    headers = list(_HEADER_RE.finditer(text))
    segments: list[tuple[str, str]] = []  # (heading, body)
    if not headers:
        segments.append(("", text))
    else:
        if headers[0].start() > 0:
            segments.append(("", text[: headers[0].start()].strip()))
        for i, h in enumerate(headers):
            end = headers[i + 1].start() if i + 1 < len(headers) else len(text)
            heading = h.group(0).strip().rstrip(":").lstrip("#").strip()
            body = text[h.end() : end].strip()
            if body:
                segments.append((heading, body))

    chunks: list[Chunk] = []
    for heading, body in segments:
        if len(body) <= target_chars:
            chunks.append(
                Chunk(
                    chunk_index=len(chunks),
                    heading_path=heading,
                    section_ref=heading,
                    content=body,
                    token_count=_count_tokens(body),
                )
            )
            continue

        # Pack paragraphs
        paragraphs = [p.strip() for p in body.split("\n\n") if p.strip()]
        buf = ""
        for p in paragraphs:
            if buf and len(buf) + len(p) + 2 > target_chars:
                chunks.append(
                    Chunk(
                        chunk_index=len(chunks),
                        heading_path=heading,
                        section_ref=heading,
                        content=buf.strip(),
                        token_count=_count_tokens(buf),
                    )
                )
                # Overlap: keep last N chars
                buf = (buf[-overlap_chars:] + "\n\n" + p) if overlap_chars > 0 else p
            else:
                buf = f"{buf}\n\n{p}" if buf else p
        if buf.strip():
            chunks.append(
                Chunk(
                    chunk_index=len(chunks),
                    heading_path=heading,
                    section_ref=heading,
                    content=buf.strip(),
                    token_count=_count_tokens(buf),
                )
            )

    return _as_dicts(chunks)


def chunk(text: str, kind: Literal["law", "regulation", "circular"]) -> list[dict]:
    if kind in ("law", "regulation"):
        raw = chunk_statute(text)
    else:
        raw = chunk_circular(text)
    # Reconstruct Chunk objects to apply the safety ceiling.
    chunks = [
        Chunk(
            chunk_index=r["chunk_index"],
            heading_path=r["heading_path"] or "",
            section_ref=r["section_ref"] or "",
            content=r["content"],
            token_count=r["token_count"],
        )
        for r in raw
    ]
    return _as_dicts(_split_oversized(chunks))