Mahdi Bazrafshan
143351b96b
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
69 lines
1.9 KiB
Python
69 lines
1.9 KiB
Python
"""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]
|