feat(chunking): add base strategy interface, embedding service, and orchestration

Why:
- Need abstract base class for all chunking strategies
- Need OpenAI embedding service (text-embedding-3-small)
- Need orchestration to run chunk → embed → store pipeline

Changes:
- Base: ChunkingStrategy ABC, token counting, chunk ID generation, sentence splitting
- Embedding: batch embedding with 2048 text limit per call
- Service: strategy registry, single/multi-strategy runners with per-strategy failure isolation
This commit is contained in:
2026-07-26 09:37:46 +03:30
parent fdc21e3316
commit 143351b96b
4 changed files with 343 additions and 0 deletions

1
src/chunking/__init__.py Normal file
View File

@@ -0,0 +1 @@
"""Chunking strategies — interface, implementations, embedding, orchestration."""

95
src/chunking/base.py Normal file
View File

@@ -0,0 +1,95 @@
"""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 ────────────────────────────────────────────
_SENTENCE_RE = re.compile(r"(?<=[.!?])\s+(?=[A-Z0-9])")
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()]
# ── 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).
"""

68
src/chunking/embedding.py Normal file
View File

@@ -0,0 +1,68 @@
"""OpenAI embedding service using text-embedding-3-small.
All strategies share the same embedding model (fixed, not configurable)
to ensure fair comparison. Batch support up to 2048 texts per call.
"""
from __future__ import annotations
import logging
from openai import OpenAI
from src.core.config import settings
from src.core.dependencies import get_openai_client
from src.core.exceptions import EmbeddingError
logger = logging.getLogger(__name__)
# OpenAI batch limit for text-embedding-3-small
_BATCH_SIZE = 2048
def embed_texts(texts: list[str]) -> list[list[float]]:
"""Embed a list of texts and return their vectors.
For the contextual_structure strategy, these are the enriched texts
(not raw content) — this is by design (ADR 0011).
Args:
texts: List of strings to embed.
Returns:
List of embedding vectors (same order as input).
Raises:
EmbeddingError: If the OpenAI API call fails.
"""
if not texts:
return []
client = get_openai_client()
all_embeddings: list[list[float]] = []
try:
for start in range(0, len(texts), _BATCH_SIZE):
batch = texts[start:start + _BATCH_SIZE]
response = client.embeddings.create(
model=settings.embedding_model,
input=batch,
)
# Sort by index to guarantee order matches input
sorted_data = sorted(response.data, key=lambda x: x.index)
all_embeddings.extend([item.embedding for item in sorted_data])
logger.debug(
"Embedded batch %d-%d (%d texts)",
start, start + len(batch), len(batch),
)
return all_embeddings
except Exception as exc:
raise EmbeddingError(f"Embedding failed: {exc}") from exc
def embed_single(text: str) -> list[float]:
"""Embed a single text (convenience wrapper)."""
results = embed_texts([text])
return results[0]

179
src/chunking/service.py Normal file
View File

@@ -0,0 +1,179 @@
"""Chunking orchestration service.
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.
"""
from __future__ import annotations
import logging
import time
from src.chunking.base import ChunkingStrategy
from src.chunking.embedding import embed_texts
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.core.exceptions import ChunkingError
from src.core.models import (
Chunk,
DocumentTree,
StrategyName,
)
from src.storage import qdrant as qdr
from src.storage import sqlite as db
logger = logging.getLogger(__name__)
# ── Strategy registry ─────────────────────────────────────────────
_STRATEGIES: dict[StrategyName, ChunkingStrategy] = {
StrategyName.RECURSIVE: RecursiveStrategy(),
StrategyName.FIXED_SIZE: FixedSizeStrategy(),
StrategyName.SEMANTIC: SemanticStrategy(),
StrategyName.CONTEXTUAL_RETRIEVAL: ContextualRetrievalStrategy(),
StrategyName.SEMANTIC_PARENT_CHILD: SemanticParentChildStrategy(),
}
def _get_strategy(name: StrategyName) -> ChunkingStrategy:
s = _STRATEGIES.get(name)
if s is None:
raise ChunkingError(f"Unknown strategy: {name}")
return s
# ── Single-strategy runner ────────────────────────────────────────
def _run_strategy(
strategy_name: StrategyName,
doc_id: str,
doc_name: str,
tree: DocumentTree,
markdown: str,
) -> int:
"""Run one strategy: chunk → embed → store in Qdrant.
Returns the number of chunks produced.
Raises on any failure (caller handles isolation).
"""
strategy = _get_strategy(strategy_name)
# Ensure Qdrant collection exists
qdr.ensure_collection(strategy_name)
t0 = time.time()
# Step 1: Chunk
chunks = strategy.chunk(
doc_name=doc_name,
tree=tree,
markdown=markdown,
)
if not chunks:
logger.warning("Strategy %s produced 0 chunks for doc %s", strategy_name.value, doc_id)
return 0
t_chunk = time.time() - t0
logger.info(
"Strategy %s: %d chunks in %.2fs",
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:
texts_to_embed.append(chunk.enriched_content)
else:
texts_to_embed.append(chunk.text)
t1 = time.time()
embeddings = embed_texts(texts_to_embed)
t_embed = time.time() - t1
logger.info(
"Strategy %s: embedded %d texts in %.2fs",
strategy_name.value, len(embeddings), t_embed,
)
# Step 3: Upsert to Qdrant
t2 = time.time()
stored = qdr.upsert_chunks(chunks, embeddings)
t_store = time.time() - t2
logger.info(
"Strategy %s: stored %d vectors in %.2fs",
strategy_name.value, stored, t_store,
)
return len(chunks)
# ── Multi-strategy orchestrator ───────────────────────────────────
def run_strategies(
doc_id: str,
strategies: list[StrategyName],
) -> tuple[list[dict], list[dict]]:
"""Run multiple strategies on a document with per-strategy failure isolation.
Returns:
(completed, failed) — lists of result dicts.
"""
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)
tree_raw = doc["document_tree"]
if isinstance(tree_raw, str):
tree = DocumentTree.model_validate_json(tree_raw)
else:
tree = DocumentTree.model_validate(tree_raw)
markdown = doc["parsed_text"]
doc_name = doc["filename"]
completed: list[dict] = []
failed: list[dict] = []
for strategy_name in strategies:
try:
t0 = time.time()
chunks_produced = _run_strategy(
strategy_name=strategy_name,
doc_id=doc_id,
doc_name=doc_name,
tree=tree,
markdown=markdown,
)
elapsed = time.time() - t0
completed.append({
"strategy": strategy_name,
"status": "completed",
"chunks_produced": chunks_produced,
"elapsed_seconds": round(elapsed, 2),
})
except Exception as exc:
logger.error(
"Strategy %s failed for doc %s: %s",
strategy_name.value, doc_id, exc,
)
failed.append({
"strategy": strategy_name,
"status": "failed",
"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)
return completed, failed