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:
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).
|
||||
"""
|
||||
Reference in New Issue
Block a user