File size: 9,038 Bytes
30e6218 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 | """Embedding providers.
Providers are constructed lazily from settings. Importing this module never
imports a provider SDK, opens a connection, or loads a model.
- ``openai``: OpenAI embeddings API (paid; requires ``OPENAI_API_KEY``).
- ``local``: sentence-transformers model loaded from the local cache. Downloads
are refused unless ``ALLOW_MODEL_DOWNLOAD=true``.
- ``fastembed``: the same MiniLM family through ONNX Runtime (no torch). Model
files come from the fastembed cache; downloads are refused unless
``ALLOW_MODEL_DOWNLOAD=true``.
- ``hash``: deterministic hashed bag-of-words vectors. Test-only: proves the
vector plumbing but carries no semantic signal.
- ``none``: dense retrieval disabled; the system runs lexical-only.
"""
from __future__ import annotations
import asyncio
import hashlib
import logging
import math
import os
import re
from typing import List, Optional, Protocol, Sequence
from .config import Settings
logger = logging.getLogger(__name__)
class EmbeddingProvider(Protocol):
identity: str
dimension: int
semantic: bool
def embed_documents(self, texts: Sequence[str]) -> List[List[float]]:
...
def embed_query(self, text: str) -> List[float]:
...
class HashEmbeddings:
"""Deterministic, dependency-free embeddings for plumbing tests."""
semantic = False
_token_re = re.compile(r"\w+")
def __init__(self, dimension: int = 64):
self.dimension = dimension
self.identity = f"hash:bow:{dimension}"
def _vector(self, text: str) -> List[float]:
vec = [0.0] * self.dimension
for token in self._token_re.findall(text.lower()):
digest = hashlib.blake2b(token.encode("utf-8"), digest_size=8).digest()
index = int.from_bytes(digest[:4], "big") % self.dimension
sign = 1.0 if digest[4] % 2 == 0 else -1.0
vec[index] += sign
norm = math.sqrt(sum(v * v for v in vec))
if norm == 0.0:
vec[0] = 1.0
return vec
return [v / norm for v in vec]
def embed_documents(self, texts: Sequence[str]) -> List[List[float]]:
return [self._vector(t) for t in texts]
def embed_query(self, text: str) -> List[float]:
return self._vector(text)
class OpenAIEmbeddings:
"""OpenAI embeddings via the official SDK (imported lazily)."""
semantic = True
_KNOWN_DIMENSIONS = {
"text-embedding-3-small": 1536,
"text-embedding-3-large": 3072,
"text-embedding-ada-002": 1536,
}
def __init__(
self,
model: str,
api_key: str,
base_url: Optional[str] = None,
timeout: float = 60.0,
batch_size: int = 100,
client=None,
):
self.model = model
self.identity = f"openai:{model}"
self.batch_size = batch_size
self.dimension = self._KNOWN_DIMENSIONS.get(model, 0)
if client is None:
try:
from openai import OpenAI
except ImportError as exc:
raise RuntimeError(
"The 'openai' package is required for embedding_provider=openai "
"(pip install -r requirements-ml.txt)"
) from exc
client = OpenAI(api_key=api_key, base_url=base_url, timeout=timeout)
self._client = client
def embed_documents(self, texts: Sequence[str]) -> List[List[float]]:
out: List[List[float]] = []
for i in range(0, len(texts), self.batch_size):
batch = [t if t.strip() else " " for t in texts[i : i + self.batch_size]]
response = self._client.embeddings.create(model=self.model, input=batch)
vectors = [item.embedding for item in sorted(response.data, key=lambda d: d.index)]
if len(vectors) != len(batch):
raise RuntimeError("Embedding provider returned an unexpected count")
out.extend(vectors)
if out and not self.dimension:
self.dimension = len(out[0])
return out
def embed_query(self, text: str) -> List[float]:
return self.embed_documents([text])[0]
class LocalEmbeddings:
"""sentence-transformers embeddings loaded from the local HF cache."""
semantic = True
def __init__(self, model_name: str, allow_download: bool, device: str = "cpu"):
self.identity = f"local:{model_name}"
if not allow_download:
os.environ.setdefault("HF_HUB_OFFLINE", "1")
os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
try:
from sentence_transformers import SentenceTransformer
except ImportError as exc:
raise RuntimeError(
"sentence-transformers is required for embedding_provider=local "
"(pip install -r requirements-ml.txt)"
) from exc
try:
self._model = SentenceTransformer(model_name, device=device)
except Exception as exc:
raise RuntimeError(
f"Could not load local embedding model '{model_name}'. Prepare it "
"with `python scripts/setup/init_models.py --embedding` or set "
"ALLOW_MODEL_DOWNLOAD=true."
) from exc
self.dimension = int(self._model.get_sentence_embedding_dimension() or 0)
def embed_documents(self, texts: Sequence[str]) -> List[List[float]]:
vectors = self._model.encode(list(texts), normalize_embeddings=True)
return [list(map(float, v)) for v in vectors]
def embed_query(self, text: str) -> List[float]:
return self.embed_documents([text])[0]
class FastEmbedEmbeddings:
"""ONNX Runtime embeddings via ``fastembed`` (imported lazily).
The model must already be present in ``cache_dir`` (prepare it with
``scripts/setup/init_models.py --fastembed``) unless downloads are allowed.
"""
semantic = True
def __init__(
self,
model_name: str,
cache_dir: Optional[str],
allow_download: bool,
model_factory=None,
):
self.identity = f"fastembed:{model_name}"
self.model_name = model_name
if not allow_download:
os.environ.setdefault("HF_HUB_OFFLINE", "1")
if model_factory is None:
try:
from fastembed import TextEmbedding
except ImportError as exc:
raise RuntimeError(
"fastembed is required for embedding_provider=fastembed "
"(pip install -r requirements-demo.txt)"
) from exc
model_factory = TextEmbedding
kwargs: dict = {"local_files_only": not allow_download}
if cache_dir:
kwargs["cache_dir"] = cache_dir
try:
self._model = model_factory(model_name, **kwargs)
except Exception as exc:
raise RuntimeError(
f"Could not load fastembed model '{model_name}'. Prepare it with "
"`python scripts/setup/init_models.py --fastembed` or set "
"ALLOW_MODEL_DOWNLOAD=true."
) from exc
self.dimension = len(self.embed_query("dimension probe"))
def embed_documents(self, texts: Sequence[str]) -> List[List[float]]:
batch = [t if t.strip() else " " for t in texts]
if not batch:
return []
vectors = [list(map(float, v)) for v in self._model.embed(batch)]
if len(vectors) != len(batch):
raise RuntimeError("fastembed returned an unexpected number of vectors")
return vectors
def embed_query(self, text: str) -> List[float]:
return self.embed_documents([text])[0]
def build_embedding_provider(settings: Settings) -> Optional[EmbeddingProvider]:
"""Instantiate the configured provider, or ``None`` for lexical-only mode."""
name = settings.resolved_embedding_provider
if name == "none":
return None
if name == "hash":
return HashEmbeddings()
if name == "openai":
return OpenAIEmbeddings(
model=settings.openai_embedding_model,
api_key=settings.openai_api_key or "",
base_url=settings.openai_base_url,
timeout=settings.openai_timeout_seconds,
)
if name == "local":
return LocalEmbeddings(
settings.local_embedding_model, allow_download=settings.allow_model_download
)
if name == "fastembed":
return FastEmbedEmbeddings(
settings.fastembed_model,
cache_dir=settings.fastembed_cache_dir,
allow_download=settings.allow_model_download,
)
raise ValueError(f"Unknown embedding provider: {name}")
async def embed_documents_async(
provider: EmbeddingProvider, texts: Sequence[str]
) -> List[List[float]]:
return await asyncio.to_thread(provider.embed_documents, texts)
async def embed_query_async(provider: EmbeddingProvider, text: str) -> List[float]:
return await asyncio.to_thread(provider.embed_query, text)
|