feat(embedding): support boundary and corpus model roles
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
@@ -1,68 +1,219 @@
|
||||
"""OpenAI embedding service using text-embedding-3-small.
|
||||
"""Embedding service with Cloud (OpenAI) and Local (Ollama) Providers.
|
||||
|
||||
All strategies share the same embedding model (fixed, not configurable)
|
||||
to ensure fair comparison. Batch support up to 2048 texts per call.
|
||||
ADR-0024: two roles — Boundary (semantic cuts) and Corpus (storage + query).
|
||||
Callers snapshot both at operation start so a mid-flight Admin switch cannot mix models.
|
||||
Legacy Active Embedding Model maps to Corpus (and migrates into both defaults).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Literal
|
||||
|
||||
from openai import OpenAI
|
||||
|
||||
from src.core.config import settings
|
||||
from src.core.dependencies import get_openai_client
|
||||
from src.chunking.embedding_models import (
|
||||
EmbeddingModelSpec,
|
||||
Provider,
|
||||
apply_task_prefixes,
|
||||
default_model_id,
|
||||
get_model,
|
||||
)
|
||||
from src.core.dependencies import get_ollama_client, get_openai_client
|
||||
from src.core.exceptions import EmbeddingError
|
||||
from src.storage import sqlite as db
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# OpenAI batch limit for text-embedding-3-small
|
||||
_BATCH_SIZE = 2048
|
||||
_OPENAI_BATCH_SIZE = 2048
|
||||
_OLLAMA_BATCH_SIZE = 64
|
||||
|
||||
_CORPUS_SETTING = "corpus_embedding_model_id"
|
||||
_BOUNDARY_SETTING = "boundary_embedding_model_id"
|
||||
_LEGACY_ACTIVE_SETTING = "active_embedding_model_id"
|
||||
|
||||
Purpose = Literal["document", "query"]
|
||||
|
||||
|
||||
def embed_texts(texts: list[str]) -> list[list[float]]:
|
||||
"""Embed a list of texts and return their vectors.
|
||||
def _ensure_role_defaults_migrated() -> None:
|
||||
"""One-shot: legacy Active → Corpus + Boundary when role keys unset."""
|
||||
corpus = db.get_app_setting(_CORPUS_SETTING)
|
||||
boundary = db.get_app_setting(_BOUNDARY_SETTING)
|
||||
if corpus and boundary:
|
||||
return
|
||||
legacy = db.get_app_setting(_LEGACY_ACTIVE_SETTING) or default_model_id()
|
||||
if not corpus:
|
||||
db.set_app_setting(_CORPUS_SETTING, legacy)
|
||||
if not boundary:
|
||||
db.set_app_setting(_BOUNDARY_SETTING, legacy)
|
||||
|
||||
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.
|
||||
def _resolve_role(setting_key: str) -> EmbeddingModelSpec:
|
||||
_ensure_role_defaults_migrated()
|
||||
stored = db.get_app_setting(setting_key)
|
||||
model_id = stored or default_model_id()
|
||||
try:
|
||||
return get_model(model_id)
|
||||
except KeyError:
|
||||
logger.warning(
|
||||
"Stored Embedding Model '%s' (%s) not in registry; falling back to %s",
|
||||
model_id,
|
||||
setting_key,
|
||||
default_model_id(),
|
||||
)
|
||||
return get_model(default_model_id())
|
||||
|
||||
Returns:
|
||||
List of embedding vectors (same order as input).
|
||||
|
||||
Raises:
|
||||
EmbeddingError: If the OpenAI API call fails.
|
||||
def get_corpus_embedding_model() -> EmbeddingModelSpec:
|
||||
"""Default Corpus Embedding Model (storage + query)."""
|
||||
return _resolve_role(_CORPUS_SETTING)
|
||||
|
||||
|
||||
def get_boundary_embedding_model() -> EmbeddingModelSpec:
|
||||
"""Default Boundary Embedding Model (semantic cuts)."""
|
||||
return _resolve_role(_BOUNDARY_SETTING)
|
||||
|
||||
|
||||
def set_corpus_embedding_model(model_id: str) -> EmbeddingModelSpec:
|
||||
"""Persist Default Corpus. Raises KeyError if unknown."""
|
||||
_ensure_role_defaults_migrated()
|
||||
model = get_model(model_id)
|
||||
db.set_app_setting(_CORPUS_SETTING, model.id)
|
||||
# Keep legacy key in sync for older readers
|
||||
db.set_app_setting(_LEGACY_ACTIVE_SETTING, model.id)
|
||||
logger.info("Corpus Embedding Model set to %s (%s)", model.id, model.provider.value)
|
||||
return model
|
||||
|
||||
|
||||
def set_boundary_embedding_model(model_id: str) -> EmbeddingModelSpec:
|
||||
"""Persist Default Boundary. Raises KeyError if unknown."""
|
||||
_ensure_role_defaults_migrated()
|
||||
model = get_model(model_id)
|
||||
db.set_app_setting(_BOUNDARY_SETTING, model.id)
|
||||
logger.info("Boundary Embedding Model set to %s (%s)", model.id, model.provider.value)
|
||||
return model
|
||||
|
||||
|
||||
def resolve_corpus_model(model_id: str | None = None) -> EmbeddingModelSpec:
|
||||
"""Snapshot Corpus for an operation (explicit id or Admin default)."""
|
||||
if model_id:
|
||||
return get_model(model_id)
|
||||
return get_corpus_embedding_model()
|
||||
|
||||
|
||||
def resolve_boundary_model(model_id: str | None = None) -> EmbeddingModelSpec:
|
||||
"""Snapshot Boundary for an operation (explicit id or Admin default)."""
|
||||
if model_id:
|
||||
return get_model(model_id)
|
||||
return get_boundary_embedding_model()
|
||||
|
||||
|
||||
# ── Legacy aliases (Corpus) ───────────────────────────────────────
|
||||
|
||||
def get_active_embedding_model() -> EmbeddingModelSpec:
|
||||
"""Deprecated: alias for get_corpus_embedding_model (ADR-0024)."""
|
||||
return get_corpus_embedding_model()
|
||||
|
||||
|
||||
def set_active_embedding_model(model_id: str) -> EmbeddingModelSpec:
|
||||
"""Deprecated: sets Corpus default (and legacy active key)."""
|
||||
return set_corpus_embedding_model(model_id)
|
||||
|
||||
|
||||
def snapshot_active_model() -> EmbeddingModelSpec:
|
||||
"""Deprecated: snapshot Corpus Embedding Model."""
|
||||
return get_corpus_embedding_model()
|
||||
|
||||
|
||||
def embed_texts(
|
||||
texts: list[str],
|
||||
*,
|
||||
model: EmbeddingModelSpec | None = None,
|
||||
purpose: Purpose = "document",
|
||||
) -> list[list[float]]:
|
||||
"""Embed texts with the given (or Corpus) Embedding Model.
|
||||
|
||||
For contextual_retrieval, pass enriched texts — ADR 0011.
|
||||
"""
|
||||
if not texts:
|
||||
return []
|
||||
|
||||
model = model or get_corpus_embedding_model()
|
||||
prepared = apply_task_prefixes(texts, model=model, purpose=purpose)
|
||||
|
||||
try:
|
||||
if model.provider == Provider.CLOUD:
|
||||
return _embed_openai(prepared, model)
|
||||
if model.provider == Provider.LOCAL:
|
||||
return _embed_ollama(prepared, model)
|
||||
raise EmbeddingError(f"Unsupported Provider: {model.provider}")
|
||||
except EmbeddingError:
|
||||
raise
|
||||
except Exception as exc:
|
||||
raise EmbeddingError(f"Embedding failed ({model.id}): {exc}") from exc
|
||||
|
||||
|
||||
def embed_single(
|
||||
text: str,
|
||||
*,
|
||||
model: EmbeddingModelSpec | None = None,
|
||||
purpose: Purpose = "query",
|
||||
) -> list[float]:
|
||||
"""Embed a single text (default purpose=query for retrieval)."""
|
||||
results = embed_texts([text], model=model, purpose=purpose)
|
||||
return results[0]
|
||||
|
||||
|
||||
def _embed_openai(texts: list[str], model: EmbeddingModelSpec) -> list[list[float]]:
|
||||
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]
|
||||
for start in range(0, len(texts), _OPENAI_BATCH_SIZE):
|
||||
batch = texts[start : start + _OPENAI_BATCH_SIZE]
|
||||
response = client.embeddings.create(
|
||||
model=settings.embedding_model,
|
||||
model=model.model_name,
|
||||
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])
|
||||
|
||||
vectors = [item.embedding for item in sorted_data]
|
||||
_validate_dimensions(vectors, model)
|
||||
all_embeddings.extend(vectors)
|
||||
logger.debug(
|
||||
"Embedded batch %d-%d (%d texts)",
|
||||
start, start + len(batch), len(batch),
|
||||
"OpenAI embedded batch %d-%d with %s",
|
||||
start,
|
||||
start + len(batch),
|
||||
model.model_name,
|
||||
)
|
||||
|
||||
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]
|
||||
def _embed_ollama(texts: list[str], model: EmbeddingModelSpec) -> list[list[float]]:
|
||||
client = get_ollama_client()
|
||||
all_embeddings: list[list[float]] = []
|
||||
|
||||
for start in range(0, len(texts), _OLLAMA_BATCH_SIZE):
|
||||
batch = texts[start : start + _OLLAMA_BATCH_SIZE]
|
||||
response = client.embeddings.create(
|
||||
model=model.model_name,
|
||||
input=batch,
|
||||
)
|
||||
sorted_data = sorted(response.data, key=lambda x: x.index)
|
||||
vectors = [item.embedding for item in sorted_data]
|
||||
_validate_dimensions(vectors, model)
|
||||
all_embeddings.extend(vectors)
|
||||
logger.debug(
|
||||
"Ollama embedded batch %d-%d with %s",
|
||||
start,
|
||||
start + len(batch),
|
||||
model.model_name,
|
||||
)
|
||||
|
||||
return all_embeddings
|
||||
|
||||
|
||||
def _validate_dimensions(vectors: list[list[float]], model: EmbeddingModelSpec) -> None:
|
||||
for i, vec in enumerate(vectors):
|
||||
if len(vec) != model.dimension:
|
||||
raise EmbeddingError(
|
||||
f"Embedding dimension mismatch for {model.id}: "
|
||||
f"expected {model.dimension}, got {len(vec)} (index {i})"
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user