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