Mahdi Bazrafshan
16c918538b
Why: - Semantic strategies were silently falling back to fixed-count grouping when boundary embeds were missing. Changes: - Fail hard without aligned unit embeddings; orchestrator supplies Boundary embeds; Farsi-aware sentence split with line/paragraph fallback. Co-authored-by: Cursor <cursoragent@cursor.com>
119 lines
3.6 KiB
Python
119 lines
3.6 KiB
Python
"""Base chunking strategy interface and shared utilities.
|
||
|
||
Every strategy inherits from ChunkingStrategy and implements chunk().
|
||
The base class provides token counting, chunk ID generation, and
|
||
the standard Chunk construction path.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import hashlib
|
||
import re
|
||
from abc import ABC, abstractmethod
|
||
|
||
import tiktoken
|
||
|
||
from src.core.models import Chunk, DocumentTree, StrategyName
|
||
|
||
# cl100k_base is the encoding used by text-embedding-3-small and gpt-4o-mini
|
||
_encoder = tiktoken.get_encoding("cl100k_base")
|
||
|
||
|
||
def count_tokens(text: str) -> int:
|
||
"""Return the token count for a string."""
|
||
return len(_encoder.encode(text))
|
||
|
||
|
||
def make_chunk_id(strategy: StrategyName, doc_name: str, index: int) -> str:
|
||
"""Generate a deterministic chunk ID: {strategy}_{doc}_{index:06d}."""
|
||
safe_doc = re.sub(r"[^a-zA-Z0-9]", "_", doc_name)[:32]
|
||
return f"{strategy.value}_{safe_doc}_{index:06d}"
|
||
|
||
|
||
def build_chunk(
|
||
*,
|
||
strategy: StrategyName,
|
||
doc_name: str,
|
||
index: int,
|
||
text: str,
|
||
parent_id: str | None = None,
|
||
enriched_content: str | None = None,
|
||
) -> Chunk:
|
||
"""Construct a Chunk with token/character counts pre-filled."""
|
||
return Chunk(
|
||
document_name=doc_name,
|
||
chunk_id=make_chunk_id(strategy, doc_name, index),
|
||
strategy_name=strategy,
|
||
chunk_index=index,
|
||
text=text,
|
||
token_count=count_tokens(text),
|
||
character_count=len(text),
|
||
parent_id=parent_id,
|
||
enriched_content=enriched_content,
|
||
)
|
||
|
||
|
||
# ── Sentence splitting ────────────────────────────────────────────
|
||
|
||
# After . ! ? or Persian/fullwidth ؟ ！, split on following whitespace.
|
||
# Does NOT require a Latin capital next (that broke Farsi documents).
|
||
# Periods inside numbers (12.5) are safe because there is no whitespace after.
|
||
_SENTENCE_RE = re.compile(r"(?<=[.!?؟！])\s+")
|
||
|
||
|
||
def split_sentences(text: str) -> list[str]:
|
||
"""Split text into sentence-like units for Semantic Boundary Detection.
|
||
|
||
Primary: punctuation-based splits (English + Farsi terminators).
|
||
Fallback: if that yields a single unit, use non-empty lines, then
|
||
blank-line paragraphs — so table/list docs still get multiple units.
|
||
"""
|
||
text = text.strip()
|
||
if not text:
|
||
return []
|
||
|
||
sentences = [s.strip() for s in _SENTENCE_RE.split(text) if s.strip()]
|
||
if len(sentences) > 1:
|
||
return sentences
|
||
|
||
lines = [ln.strip() for ln in text.splitlines() if ln.strip()]
|
||
if len(lines) > 1:
|
||
return lines
|
||
|
||
paragraphs = [p.strip() for p in re.split(r"\n\s*\n", text) if p.strip()]
|
||
if len(paragraphs) > 1:
|
||
return paragraphs
|
||
|
||
return sentences if sentences else [text]
|
||
|
||
|
||
# ── Abstract base ─────────────────────────────────────────────────
|
||
|
||
class ChunkingStrategy(ABC):
|
||
"""Base class for all chunking strategies.
|
||
|
||
Subclasses implement chunk() which receives the full document
|
||
context and returns a list of Chunks conforming to the unified model.
|
||
"""
|
||
|
||
name: StrategyName
|
||
|
||
@abstractmethod
|
||
def chunk(
|
||
self,
|
||
*,
|
||
doc_name: str,
|
||
tree: DocumentTree,
|
||
markdown: str,
|
||
) -> list[Chunk]:
|
||
"""Produce chunks from a parsed document.
|
||
|
||
Args:
|
||
doc_name: Original filename (for metadata).
|
||
tree: Hierarchical DocumentTree from the parser.
|
||
markdown: Flat markdown rendering of the document.
|
||
|
||
Returns:
|
||
List of Chunk objects (unified model).
|
||
"""
|