mba/chunking_strategies_evaluation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(chunking): require semantic boundary detection

Browse Source 
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>
This commit is contained in:
Mahdi Bazrafshan
2026-08-10 14:12:45 +03:30
parent 736391b137
commit 16c918538b
4 changed files with 211 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
src/chunking/base.py
View File

@@ -55,13 +55,36 @@ def build_chunk(
# ── Sentence splitting ────────────────────────────────────────────
_SENTENCE_RE = re.compile(r"(?<=[.!?])\s+(?=[A-Z0-9])")
# 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 sentences using a simple regex heuristic."""
sentences = _SENTENCE_RE.split(text.strip())
return [s.strip() for s in sentences if s.strip()]
"""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 ─────────────────────────────────────────────────

158
src/chunking/service.py
View File

@@ -4,7 +4,8 @@ Runs selected strategies on a document, embeds chunks, and stores
them in Qdrant. Per-strategy failure isolation (ADR 0003): if one
strategy fails, the others' results are still committed.
This replaces the stub in src/documents/service.py.
ADR-0024: Boundary Embedding Model for semantic cuts; Corpus Embedding
Model for finished-chunk vectors and collection scoping.
"""
from __future__ import annotations
@@ -12,16 +13,19 @@ from __future__ import annotations
import logging
import time
from src.chunking.base import ChunkingStrategy
from src.chunking.base import ChunkingStrategy, split_sentences
from src.chunking.embedding import embed_texts
from src.chunking.embedding_models import EmbeddingModelSpec, get_semantic_threshold
from src.chunking.strategies.recursive import RecursiveStrategy
from src.chunking.strategies.fixed_size import FixedSizeStrategy
from src.chunking.strategies.semantic import SemanticStrategy
from src.chunking.strategies.contextual_retrieval import ContextualRetrievalStrategy
from src.chunking.strategies.semantic_parent_child import SemanticParentChildStrategy
from src.chunking.strategies.semantic_parent_child import (
SemanticParentChildStrategy,
split_paragraphs,
)
from src.core.exceptions import ChunkingError
from src.core.models import (
Chunk,
DocumentTree,
StrategyName,
)
@@ -30,6 +34,11 @@ from src.storage import sqlite as db
logger = logging.getLogger(__name__)
_BOUNDARY_STRATEGIES = {
StrategyName.SEMANTIC,
StrategyName.SEMANTIC_PARENT_CHILD,
}
# ── Strategy registry ─────────────────────────────────────────────
_STRATEGIES: dict[StrategyName, ChunkingStrategy] = {
@@ -48,6 +57,73 @@ def _get_strategy(name: StrategyName) -> ChunkingStrategy:
return s
def _chunk_document(
strategy_name: StrategyName,
strategy: ChunkingStrategy,
*,
doc_name: str,
tree: DocumentTree,
markdown: str,
boundary_model: EmbeddingModelSpec,
) -> list:
"""Run strategy.chunk(), supplying Boundary embeds when needed."""
if strategy_name == StrategyName.SEMANTIC:
sentences = split_sentences(markdown)
if not sentences:
return []
threshold = get_semantic_threshold(boundary_model.id)
logger.info(
"Strategy semantic: embedding %d sentences for boundary detection "
"(%s, threshold=%.3f)",
len(sentences),
boundary_model.id,
threshold,
)
sentence_embeddings = embed_texts(
sentences,
model=boundary_model,
purpose="document",
)
return strategy.chunk( # type: ignore[call-arg]
doc_name=doc_name,
tree=tree,
markdown=markdown,
sentence_embeddings=sentence_embeddings,
semantic_threshold=threshold,
)
if strategy_name == StrategyName.SEMANTIC_PARENT_CHILD:
paragraphs = split_paragraphs(markdown)
if not paragraphs:
return []
threshold = get_semantic_threshold(boundary_model.id)
logger.info(
"Strategy semantic_parent_child: embedding %d paragraphs for boundary detection "
"(%s, threshold=%.3f)",
len(paragraphs),
boundary_model.id,
threshold,
)
paragraph_embeddings = embed_texts(
paragraphs,
model=boundary_model,
purpose="document",
)
return strategy.chunk( # type: ignore[call-arg]
doc_name=doc_name,
tree=tree,
markdown=markdown,
paragraph_embeddings=paragraph_embeddings,
semantic_threshold=threshold,
)
return strategy.chunk(
doc_name=doc_name,
tree=tree,
markdown=markdown,
)
# ── Single-strategy runner ────────────────────────────────────────
def _run_strategy(
@@ -56,6 +132,9 @@ def _run_strategy(
doc_name: str,
tree: DocumentTree,
markdown: str,
*,
boundary_model: EmbeddingModelSpec,
corpus_model: EmbeddingModelSpec,
) -> int:
"""Run one strategy: chunk → embed → store in Qdrant.
@@ -64,16 +143,21 @@ def _run_strategy(
"""
strategy = _get_strategy(strategy_name)
# Ensure Qdrant collection exists
qdr.ensure_collection(strategy_name)
qdr.ensure_collection(
strategy_name,
model_id=corpus_model.id,
dimension=corpus_model.dimension,
)
t0 = time.time()
# Step 1: Chunk
chunks = strategy.chunk(
chunks = _chunk_document(
strategy_name,
strategy,
doc_name=doc_name,
tree=tree,
markdown=markdown,
boundary_model=boundary_model,
)
if not chunks:
@@ -86,8 +170,6 @@ def _run_strategy(
strategy_name.value, len(chunks), t_chunk,
)
# Step 2: Embed
# For contextual strategy, embed enriched_content; for others, embed text
texts_to_embed = []
for chunk in chunks:
if chunk.enriched_content:
@@ -96,16 +178,19 @@ def _run_strategy(
texts_to_embed.append(chunk.text)
t1 = time.time()
embeddings = embed_texts(texts_to_embed)
embeddings = embed_texts(
texts_to_embed,
model=corpus_model,
purpose="document",
)
t_embed = time.time() - t1
logger.info(
"Strategy %s: embedded %d texts in %.2fs",
strategy_name.value, len(embeddings), t_embed,
"Strategy %s: embedded %d texts in %.2fs (corpus=%s)",
strategy_name.value, len(embeddings), t_embed, corpus_model.id,
)
# Step 3: Upsert to Qdrant
t2 = time.time()
stored = qdr.upsert_chunks(chunks, embeddings)
stored = qdr.upsert_chunks(chunks, embeddings, model_id=corpus_model.id)
t_store = time.time() - t2
logger.info(
"Strategy %s: stored %d vectors in %.2fs",
@@ -120,17 +205,39 @@ def _run_strategy(
def run_strategies(
doc_id: str,
strategies: list[StrategyName],
*,
boundary_model_id: str | None = None,
corpus_model_id: str | None = None,
) -> tuple[list[dict], list[dict]]:
"""Run multiple strategies on a document with per-strategy failure isolation.
Returns:
(completed, failed) — lists of result dicts.
"""
from src.chunking.embedding import resolve_boundary_model, resolve_corpus_model
doc = db.get_document(doc_id)
if doc is None:
raise ChunkingError(f"Document not found: {doc_id}")
# Parse the stored document tree (may be dict or JSON string)
corpus_model = resolve_corpus_model(corpus_model_id)
needs_boundary = any(s in _BOUNDARY_STRATEGIES for s in strategies)
boundary_model = (
resolve_boundary_model(boundary_model_id)
if needs_boundary
else corpus_model # unused for non-semantic; keep a valid spec
)
logger.info(
"Processing doc %s corpus=%s (%s) boundary=%s (%s) needs_boundary=%s",
doc_id,
corpus_model.id,
corpus_model.provider.value,
boundary_model.id if needs_boundary else "—",
boundary_model.provider.value if needs_boundary else "—",
needs_boundary,
)
tree_raw = doc["document_tree"]
if isinstance(tree_raw, str):
tree = DocumentTree.model_validate_json(tree_raw)
@@ -151,14 +258,21 @@ def run_strategies(
doc_name=doc_name,
tree=tree,
markdown=markdown,
boundary_model=boundary_model,
corpus_model=corpus_model,
)
elapsed = time.time() - t0
completed.append({
entry = {
"strategy": strategy_name,
"status": "completed",
"chunks_produced": chunks_produced,
"elapsed_seconds": round(elapsed, 2),
})
"corpus_embedding_model_id": corpus_model.id,
"embedding_model_id": corpus_model.id, # legacy alias
}
if strategy_name in _BOUNDARY_STRATEGIES:
entry["boundary_embedding_model_id"] = boundary_model.id
completed.append(entry)
except Exception as exc:
logger.error(
"Strategy %s failed for doc %s: %s",
@@ -170,10 +284,16 @@ def run_strategies(
"error": str(exc),
})
# Update chunk counts on the document
counts = doc.get("chunk_counts", {})
for result in completed:
counts[result["strategy"].value] = result["chunks_produced"]
db.update_chunk_counts(doc_id, counts)
if completed:
db.update_process_embedding_provenance(
doc_id,
corpus_model_id=corpus_model.id,
boundary_model_id=boundary_model.id if needs_boundary else None,
)
return completed, failed

40
src/chunking/strategies/semantic.py
View File

@@ -1,17 +1,15 @@
"""Semantic chunking strategy.
Sentence-level granularity (ADR 0012):
Sentence-level granularity (ADR 0012 / ADR 0020):
1. Split markdown into sentences.
2. Embed each sentence via OpenAI.
2. Orchestrator embeds each sentence (Active Embedding Model).
3. Compute cosine similarity between adjacent sentences.
4. When similarity drops below SEMANTIC_THRESHOLD, create a chunk boundary.
4. When similarity drops below the Active Embedding Model's semantic_threshold, create a chunk boundary.
5. Enforce SEMANTIC_MIN_CHUNK_SIZE (minimum sentences per chunk).
6. Boundary sentence stays with the previous chunk.
7. Orchestrator embeds finished chunks for Qdrant storage.
Note: this strategy requires embeddings at chunk-time. The chunk()
method returns text chunks WITHOUT embeddings — the embedding step
happens in the orchestration layer (service.py) which calls the
embedding service after chunking.
Semantic Boundary Detection is required — no fixed-count fallback.
"""
from __future__ import annotations
@@ -20,6 +18,7 @@ import numpy as np
from src.chunking.base import ChunkingStrategy, build_chunk, split_sentences
from src.core.config import settings
from src.core.exceptions import ChunkingError
from src.core.models import Chunk, DocumentTree, StrategyName
@@ -85,30 +84,33 @@ class SemanticStrategy(ChunkingStrategy):
tree: DocumentTree,
markdown: str,
sentence_embeddings: list[list[float]] | None = None,
semantic_threshold: float | None = None,
) -> list[Chunk]:
"""Produce semantic chunks.
"""Produce semantic chunks via Semantic Boundary Detection.
If sentence_embeddings is provided (from the orchestration layer),
uses them for boundary detection. Otherwise, falls back to
paragraph-level chunking (sentences without similarity-based splits).
Requires sentence_embeddings aligned 1:1 with split_sentences(markdown).
"""
sentences = split_sentences(markdown)
if not sentences:
return []
threshold = settings.semantic_threshold
if sentence_embeddings is None or len(sentence_embeddings) != len(sentences):
got = 0 if sentence_embeddings is None else len(sentence_embeddings)
raise ChunkingError(
f"semantic requires sentence embeddings for Semantic Boundary Detection "
f"(got {got}, need {len(sentences)}). Fixed-count fallback is disabled (ADR-0020)."
)
threshold = (
semantic_threshold
if semantic_threshold is not None
else settings.semantic_threshold
)
min_size = settings.semantic_min_chunk_size
if sentence_embeddings and len(sentence_embeddings) == len(sentences):
chunk_texts = _group_sentences_into_chunks(
sentences, sentence_embeddings, threshold, min_size
)
else:
# Fallback: group sentences into fixed-size chunks
chunk_texts = []
for i in range(0, len(sentences), min_size):
group = sentences[i:i + min_size]
chunk_texts.append(" ".join(group))
chunks: list[Chunk] = []
for i, text in enumerate(chunk_texts):

39
src/chunking/strategies/semantic_parent_child.py
View File

@@ -8,11 +8,13 @@ parent cluster is returned as context — giving the LLM richer
information than a single paragraph.
No headings or document structure needed — uses meaning instead.
Semantic Boundary Detection is required — no fixed-count fallback (ADR-0020).
"""
from __future__ import annotations
import logging
import re
import numpy as np
@@ -20,9 +22,9 @@ from src.chunking.base import (
ChunkingStrategy,
build_chunk,
make_chunk_id,
count_tokens,
)
from src.core.config import settings
from src.core.exceptions import ChunkingError
from src.core.models import Chunk, DocumentTree, StrategyName
logger = logging.getLogger(__name__)
@@ -34,10 +36,8 @@ def _cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
return float(dot / norm) if norm > 0 else 0.0
def _split_paragraphs(text: str) -> list[str]:
def split_paragraphs(text: str) -> list[str]:
"""Split markdown into paragraphs (double newline or single newline)."""
import re
# Split on double newlines first, then filter empties
parts = re.split(r"\n\s*\n", text)
paragraphs = [p.strip() for p in parts if p.strip()]
@@ -89,34 +89,35 @@ class SemanticParentChildStrategy(ChunkingStrategy):
tree: DocumentTree,
markdown: str,
paragraph_embeddings: list[list[float]] | None = None,
semantic_threshold: float | None = None,
) -> list[Chunk]:
"""Produce parent-child chunks via semantic clustering.
"""Produce parent-child chunks via Semantic Boundary Detection.
If paragraph_embeddings is provided (from orchestration layer),
uses them for clustering. Otherwise, groups paragraphs by
fixed count.
Requires paragraph_embeddings aligned 1:1 with split_paragraphs(markdown).
"""
paragraphs = _split_paragraphs(markdown)
paragraphs = split_paragraphs(markdown)
if not paragraphs:
return []
threshold = settings.semantic_threshold
if paragraph_embeddings is None or len(paragraph_embeddings) != len(paragraphs):
got = 0 if paragraph_embeddings is None else len(paragraph_embeddings)
raise ChunkingError(
f"semantic_parent_child requires paragraph embeddings for "
f"Semantic Boundary Detection (got {got}, need {len(paragraphs)}). "
f"Fixed-count fallback is disabled (ADR-0020)."
)
if paragraph_embeddings and len(paragraph_embeddings) == len(paragraphs):
threshold = (
semantic_threshold
if semantic_threshold is not None
else settings.semantic_threshold
)
clusters = _cluster_paragraphs(paragraphs, paragraph_embeddings, threshold)
else:
# Fallback: group every N paragraphs
group_size = max(3, settings.semantic_min_chunk_size)
clusters = []
for i in range(0, len(paragraphs), group_size):
clusters.append(list(range(i, min(i + group_size, len(paragraphs)))))
chunks: list[Chunk] = []
chunk_index = 0
for cluster_indices in clusters:
# Parent = full cluster text
parent_text = "\n\n".join(paragraphs[i] for i in cluster_indices)
parent_id = make_chunk_id(self.name, doc_name, chunk_index)
# Each paragraph in the cluster is a child