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:
1
src/chunking/__init__.py
Normal file
1
src/chunking/__init__.py
Normal file
@@ -0,0 +1 @@
|
|||||||
|
"""Chunking strategies — interface, implementations, embedding, orchestration."""
|
||||||
95
src/chunking/base.py
Normal file
95
src/chunking/base.py
Normal 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
68
src/chunking/embedding.py
Normal 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
179
src/chunking/service.py
Normal 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
|
||||||
Reference in New Issue
Block a user