feat(strategies): implement all 5 chunking strategies

Why:
- Core evaluation targets — each strategy chunks differently for comparison

Changes:
- fixed_size: token-based splitting with overlap (baseline)
- recursive: cascade splitting (headers → newlines → sentences → words)
- semantic: sentence-level embeddings with similarity-based boundaries
- contextual_retrieval: LLM-generated context prefixes per chunk (Anthropic research)
- semantic_parent_child: paragraph clustering into parent-child hierarchy
This commit is contained in:
2026-07-26 09:37:56 +03:30
parent 143351b96b
commit 8715d9c186
6 changed files with 538 additions and 0 deletions

View File

@@ -0,0 +1 @@
"""Chunking strategy implementations."""

View File

@@ -0,0 +1,120 @@
"""Contextual Retrieval strategy.
Based on Anthropic's 2024 research: prepend a short context summary
to each chunk before embedding. Improved retrieval by 49% in their
benchmarks.
Algorithm:
1. Split markdown into chunks using fixed-size token splitting.
2. For each chunk, send surrounding text + chunk to LLM.
3. LLM generates a 1-2 sentence context prefix.
4. The enriched chunk (context + original text) is what gets embedded.
No headings or document structure needed — works on any text.
"""
from __future__ import annotations
import logging
from openai import OpenAI
from src.chunking.base import ChunkingStrategy, _encoder, build_chunk
from src.core.config import settings
from src.core.dependencies import get_openai_client
from src.core.exceptions import EnrichmentError
from src.core.models import Chunk, DocumentTree, StrategyName
from src.chunking.strategies.fixed_size import _split_by_tokens
logger = logging.getLogger(__name__)
_CONTEXT_SYSTEM_PROMPT = (
"You are a document analysis assistant. Given a section of text from "
"a document, write a short context prefix (1-2 sentences) that would "
"help someone find this section later via search. Focus on the topic "
"and key terms. Do not repeat the text itself. Output ONLY the context "
"prefix, nothing else."
)
def _enrich_chunk(
client: OpenAI,
chunk_text: str,
preceding_text: str,
following_text: str,
) -> str:
"""Generate a context prefix for a chunk using the LLM.
Returns the enriched text: context prefix + original chunk.
Raises EnrichmentError on failure.
"""
user_prompt = (
f"Preceding text:\n{preceding_text[-500:] if preceding_text else '(start of document)'}\n\n"
f"This section:\n{chunk_text}\n\n"
f"Following text:\n{following_text[:500] if following_text else '(end of document)'}"
)
try:
response = client.chat.completions.create(
model=settings.llm_model,
temperature=0.0,
max_tokens=100,
messages=[
{"role": "system", "content": _CONTEXT_SYSTEM_PROMPT},
{"role": "user", "content": user_prompt},
],
)
context = response.choices[0].message.content
if not context or not context.strip():
return chunk_text
return f"{context.strip()}\n\n{chunk_text}"
except Exception as exc:
raise EnrichmentError(f"Context enrichment failed: {exc}") from exc
class ContextualRetrievalStrategy(ChunkingStrategy):
name = StrategyName.CONTEXTUAL_RETRIEVAL
def chunk(
self,
*,
doc_name: str,
tree: DocumentTree,
markdown: str,
) -> list[Chunk]:
client = get_openai_client()
chunk_size = settings.chunk_size
overlap = settings.chunk_overlap
raw_chunks = _split_by_tokens(markdown, chunk_size, overlap)
if not raw_chunks:
return []
# Build enriched chunks with context
chunks: list[Chunk] = []
full_text = markdown
for i, chunk_text in enumerate(raw_chunks):
chunk_text = chunk_text.strip()
if not chunk_text:
continue
# Find surrounding context in the full text
pos = full_text.find(chunk_text[:100])
if pos == -1:
pos = 0
preceding = full_text[max(0, pos - 500):pos]
following = full_text[pos + len(chunk_text):pos + len(chunk_text) + 500]
# Enrich with LLM context
enriched = _enrich_chunk(client, chunk_text, preceding, following)
chunks.append(build_chunk(
strategy=self.name,
doc_name=doc_name,
index=i,
text=chunk_text,
enriched_content=enriched,
))
return chunks

View File

@@ -0,0 +1,61 @@
"""Fixed-size chunking with overlap.
The simplest possible strategy — a reliable baseline for comparison.
Splits text into chunks of N tokens with M token overlap.
If smarter strategies can't beat this baseline, they're not worth
the complexity. That's the whole point of including it.
"""
from __future__ import annotations
from src.chunking.base import ChunkingStrategy, _encoder, build_chunk
from src.core.config import settings
from src.core.models import Chunk, DocumentTree, StrategyName
def _split_by_tokens(text: str, chunk_size: int, overlap: int) -> list[str]:
"""Split text into token-sized chunks with overlap."""
tokens = _encoder.encode(text)
if len(tokens) <= chunk_size:
return [text]
chunks: list[str] = []
start = 0
while start < len(tokens):
end = min(start + chunk_size, len(tokens))
chunk_tokens = tokens[start:end]
chunks.append(_encoder.decode(chunk_tokens))
if end >= len(tokens):
break
start = end - overlap
return chunks
class FixedSizeStrategy(ChunkingStrategy):
name = StrategyName.FIXED_SIZE
def chunk(
self,
*,
doc_name: str,
tree: DocumentTree,
markdown: str,
) -> list[Chunk]:
chunk_size = settings.chunk_size
overlap = settings.chunk_overlap
raw_chunks = _split_by_tokens(markdown, chunk_size, overlap)
chunks: list[Chunk] = []
for i, text in enumerate(raw_chunks):
text = text.strip()
if text:
chunks.append(build_chunk(
strategy=self.name,
doc_name=doc_name,
index=i,
text=text,
))
return chunks

View File

@@ -0,0 +1,100 @@
"""Recursive chunking strategy.
Cascade splitting using a separator hierarchy (ADR 0015 — direct API, no frameworks):
1. Markdown headers (#, ##, ###)
2. Double newline (\\n\\n)
3. Single newline (\\n)
4. Sentence-ending punctuation (. ! ? followed by space)
5. Space (word-level, last resort)
Splitting stops when chunks reach the target size. Each chunk records
which separator level produced it.
"""
from __future__ import annotations
import re
from src.chunking.base import ChunkingStrategy, build_chunk, count_tokens
from src.core.config import settings
from src.core.models import Chunk, DocumentTree, StrategyName
# Separator patterns ordered by priority (highest first)
_SEPARATORS: list[tuple[str, re.Pattern[str]]] = [
("header", re.compile(r"(?m)^(#{1,6})\s+")),
("double_newline", re.compile(r"\n\n")),
("newline", re.compile(r"\n")),
("sentence", re.compile(r"(?<=[.!?])\s+")),
("space", re.compile(r"\s+")),
]
def _split_recursive(text: str, target_size: int) -> list[str]:
"""Recursively split text using the separator cascade."""
if not text.strip():
return []
# If text fits in target size, return as-is
if len(text) <= target_size:
return [text.strip()]
# Try each separator
for sep_name, pattern in _SEPARATORS:
parts = pattern.split(text)
if len(parts) <= 1:
continue # this separator didn't split anything
# Merge parts back up to target_size
chunks: list[str] = []
current = ""
for part in parts:
if not part:
continue
candidate = (current + " " + part).strip() if current else part
if len(candidate) <= target_size:
current = candidate
else:
if current:
chunks.append(current)
# If single part exceeds target, recurse with next separator
if len(part) > target_size:
sub_chunks = _split_recursive(part, target_size)
chunks.extend(sub_chunks)
current = ""
else:
current = part
if current:
chunks.append(current)
return chunks
# Last resort: hard split by character count
return [text[i:i + target_size].strip()
for i in range(0, len(text), target_size)
if text[i:i + target_size].strip()]
class RecursiveStrategy(ChunkingStrategy):
name = StrategyName.RECURSIVE
def chunk(
self,
*,
doc_name: str,
tree: DocumentTree,
markdown: str,
) -> list[Chunk]:
target_size = settings.chunk_size
raw_chunks = _split_recursive(markdown, target_size)
chunks: list[Chunk] = []
for i, text in enumerate(raw_chunks):
if not text.strip():
continue
chunks.append(build_chunk(
strategy=self.name,
doc_name=doc_name,
index=i,
text=text.strip(),
))
return chunks

View File

@@ -0,0 +1,122 @@
"""Semantic chunking strategy.
Sentence-level granularity (ADR 0012):
1. Split markdown into sentences.
2. Embed each sentence via OpenAI.
3. Compute cosine similarity between adjacent sentences.
4. When similarity drops below SEMANTIC_THRESHOLD, create a chunk boundary.
5. Enforce SEMANTIC_MIN_CHUNK_SIZE (minimum sentences per chunk).
6. Boundary sentence stays with the previous chunk.
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.
"""
from __future__ import annotations
import numpy as np
from src.chunking.base import ChunkingStrategy, build_chunk, split_sentences
from src.core.config import settings
from src.core.models import Chunk, DocumentTree, StrategyName
def _cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
"""Compute cosine similarity between two vectors."""
dot = np.dot(a, b)
norm = np.linalg.norm(a) * np.linalg.norm(b)
if norm == 0:
return 0.0
return float(dot / norm)
def _group_sentences_into_chunks(
sentences: list[str],
embeddings: list[list[float]],
threshold: float,
min_size: int,
) -> list[str]:
"""Group sentences into chunks based on semantic similarity.
Returns a list of chunk texts.
"""
if not sentences:
return []
if len(sentences) <= min_size:
return [" ".join(sentences)]
chunks: list[str] = []
current_group: list[str] = [sentences[0]]
for i in range(1, len(sentences)):
sim = _cosine_similarity(
np.array(embeddings[i - 1]),
np.array(embeddings[i]),
)
if sim < threshold and len(current_group) >= min_size:
# Topic shift — close current chunk
chunks.append(" ".join(current_group))
current_group = [sentences[i]]
else:
current_group.append(sentences[i])
# Flush remaining
if current_group:
# If the last group is too small, merge into previous
if chunks and len(current_group) < min_size:
last = chunks.pop()
chunks.append(last + " " + " ".join(current_group))
else:
chunks.append(" ".join(current_group))
return chunks
class SemanticStrategy(ChunkingStrategy):
name = StrategyName.SEMANTIC
def chunk(
self,
*,
doc_name: str,
tree: DocumentTree,
markdown: str,
sentence_embeddings: list[list[float]] | None = None,
) -> list[Chunk]:
"""Produce semantic chunks.
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).
"""
sentences = split_sentences(markdown)
if not sentences:
return []
threshold = 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):
if text.strip():
chunks.append(build_chunk(
strategy=self.name,
doc_name=doc_name,
index=i,
text=text.strip(),
))
return chunks

View File

@@ -0,0 +1,134 @@
"""Parent-Child via Semantic Clustering strategy.
Groups paragraphs by semantic similarity into clusters. Each cluster
is a parent; each paragraph in the cluster is a child.
At query time: the child is found via vector search, then its full
parent cluster is returned as context — giving the LLM richer
information than a single paragraph.
No headings or document structure needed — uses meaning instead.
"""
from __future__ import annotations
import logging
import numpy as np
from src.chunking.base import (
ChunkingStrategy,
build_chunk,
make_chunk_id,
count_tokens,
)
from src.core.config import settings
from src.core.models import Chunk, DocumentTree, StrategyName
logger = logging.getLogger(__name__)
def _cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
dot = np.dot(a, b)
norm = np.linalg.norm(a) * np.linalg.norm(b)
return float(dot / norm) if norm > 0 else 0.0
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()]
# If no double newlines, try single newlines
if len(paragraphs) <= 1:
parts = text.split("\n")
paragraphs = [p.strip() for p in parts if p.strip()]
return paragraphs
def _cluster_paragraphs(
paragraphs: list[str],
embeddings: list[list[float]],
threshold: float,
) -> list[list[int]]:
"""Group consecutive paragraphs into clusters by semantic similarity.
Returns a list of clusters, each a list of paragraph indices.
"""
if not paragraphs or not embeddings:
return []
clusters: list[list[int]] = [[0]]
for i in range(1, len(paragraphs)):
sim = _cosine_similarity(
np.array(embeddings[i - 1]),
np.array(embeddings[i]),
)
if sim >= threshold:
# Same cluster
clusters[-1].append(i)
else:
# New cluster
clusters.append([i])
return clusters
class SemanticParentChildStrategy(ChunkingStrategy):
name = StrategyName.SEMANTIC_PARENT_CHILD
def chunk(
self,
*,
doc_name: str,
tree: DocumentTree,
markdown: str,
paragraph_embeddings: list[list[float]] | None = None,
) -> list[Chunk]:
"""Produce parent-child chunks via semantic clustering.
If paragraph_embeddings is provided (from orchestration layer),
uses them for clustering. Otherwise, groups paragraphs by
fixed count.
"""
paragraphs = _split_paragraphs(markdown)
if not paragraphs:
return []
threshold = settings.semantic_threshold
if paragraph_embeddings and len(paragraph_embeddings) == len(paragraphs):
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
for para_idx in cluster_indices:
para_text = paragraphs[para_idx]
chunks.append(build_chunk(
strategy=self.name,
doc_name=doc_name,
index=chunk_index,
text=para_text,
parent_id=parent_id,
))
chunk_index += 1
return chunks