From 143351b96bf9375d51d88103d3493977ffdc4e04 Mon Sep 17 00:00:00 2001 From: Mahdi Bazrafshan Date: Sun, 26 Jul 2026 09:37:46 +0330 Subject: [PATCH] feat(chunking): add base strategy interface, embedding service, and orchestration MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- src/chunking/__init__.py | 1 + src/chunking/base.py | 95 ++++++++++++++++++++ src/chunking/embedding.py | 68 +++++++++++++++ src/chunking/service.py | 179 ++++++++++++++++++++++++++++++++++++++ 4 files changed, 343 insertions(+) create mode 100644 src/chunking/__init__.py create mode 100644 src/chunking/base.py create mode 100644 src/chunking/embedding.py create mode 100644 src/chunking/service.py diff --git a/src/chunking/__init__.py b/src/chunking/__init__.py new file mode 100644 index 0000000..e41ff51 --- /dev/null +++ b/src/chunking/__init__.py @@ -0,0 +1 @@ +"""Chunking strategies — interface, implementations, embedding, orchestration.""" diff --git a/src/chunking/base.py b/src/chunking/base.py new file mode 100644 index 0000000..c4feebb --- /dev/null +++ b/src/chunking/base.py @@ -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). + """ diff --git a/src/chunking/embedding.py b/src/chunking/embedding.py new file mode 100644 index 0000000..4248562 --- /dev/null +++ b/src/chunking/embedding.py @@ -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] diff --git a/src/chunking/service.py b/src/chunking/service.py new file mode 100644 index 0000000..3986ba1 --- /dev/null +++ b/src/chunking/service.py @@ -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