Coverage for astrocyte/providers/local_embeddings.py: 96%
93 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-10-29 22:10 +0000
1"""Local embeddings LLMProvider — no API, no key.
3Embed-only provider backed by a local model, run through either of two
4backends that produce the same vectors:
6- **fastembed** (ONNX Runtime) — ~143 MB installed, no torch. Preferred.
7- **sentence-transformers** (torch) — ~800 MB on macOS and several GB on
8 Linux, where the default torch wheel pulls CUDA libraries.
10``backend="auto"`` (default) uses fastembed when installed and falls back to
11sentence-transformers. For the default ``BAAI/bge-small-en-v1.5`` the two agree
12to cosine 1.00000 on English, code and CJK text, with identical rankings even
13when the query is embedded by one backend and the corpus by the other — so a
14store built under either stays searchable under both.
16Usage programmatically::
18 from astrocyte.providers.local_embeddings import LocalEmbeddingsProvider
20 embedder = LocalEmbeddingsProvider() # BAAI/bge-small-en-v1.5
21 vecs = await embedder.embed(["hello world"]) # 1536-dim (zero-padded)
23Design notes:
25- **``pad_to=1536`` by default** — the Postgres reference schema declares
26 ``vector(1536)`` on ``summary_embedding`` / fact embeddings (sized for
27 OpenAI ``text-embedding-3-small``). Zero-padding a normalized vector
28 preserves cosine similarity ordering exactly, so a 384-dim local model
29 drops into the existing DDL with no migration. Vectors longer than
30 ``pad_to`` are rejected loudly (truncation would silently change geometry).
31- The model loads lazily in a worker thread on first ``embed()`` and is
32 cached for the provider's lifetime. First call downloads the model from
33 the HuggingFace hub (~130 MB for bge-small) unless already cached.
34- ``complete()`` is intentionally unsupported — compose with a completion
35 provider via :class:`~astrocyte.providers.composite.CompositeLLMProvider`.
37Requires ``fastembed`` (``pip install 'astrocyte[local]'``) or
38``sentence-transformers`` (``pip install 'astrocyte[rerank]'``).
39"""
41from __future__ import annotations
43import asyncio
44import importlib.util
45import logging
46import math
47import os
48import sys
49from pathlib import Path
50from typing import Any, ClassVar
52from astrocyte.types import Completion, LLMCapabilities, Message, ToolDefinition
54logger = logging.getLogger("astrocyte.providers.local_embeddings")
56_MAX_CHARS = 28_000 # parity with OpenAIProvider.embed truncation guard
58_BACKENDS = ("fastembed", "sentence-transformers")
61def _installed(module: str) -> bool:
62 # A module already imported (or injected, e.g. a test stub whose
63 # __spec__ is None, which makes find_spec raise) counts as installed.
64 if module in sys.modules:
65 return sys.modules[module] is not None
66 try:
67 return importlib.util.find_spec(module) is not None
68 except (ImportError, ValueError):
69 return False
72def _resolve_backend(requested: str) -> str:
73 if requested not in ("auto", *_BACKENDS):
74 raise ValueError(f"local_embeddings: backend must be 'auto' or one of {_BACKENDS}, got {requested!r}")
75 module = {"fastembed": "fastembed", "sentence-transformers": "sentence_transformers"}
76 if requested != "auto":
77 if not _installed(module[requested]):
78 raise ImportError(
79 f"local_embeddings backend {requested!r} is not installed. "
80 f"Install it with: pip install {requested}"
81 )
82 return requested
83 for name in _BACKENDS:
84 if _installed(module[name]):
85 return name
86 raise ImportError(
87 "LocalEmbeddingsProvider needs an embedding backend. Install the light one with "
88 "pip install 'astrocyte[local]' (fastembed), or pip install 'astrocyte[rerank]' "
89 "(sentence-transformers)."
90 )
93def _fastembed_cache_dir() -> str:
94 """fastembed defaults to the system temp dir, which macOS and many Linux
95 distros clear periodically — silently re-downloading the model. Keep it
96 with other user caches unless the operator chose a location."""
97 if explicit := os.environ.get("FASTEMBED_CACHE_PATH"):
98 return explicit
99 base = os.environ.get("XDG_CACHE_HOME") or "~/.cache"
100 return str(Path(base).expanduser() / "astrocyte" / "fastembed")
103class LocalEmbeddingsProvider:
104 """Embed-only LLMProvider backed by a local embedding model."""
106 SPI_VERSION: ClassVar[int] = 1
108 def __init__(
109 self,
110 *,
111 model_name: str = "BAAI/bge-small-en-v1.5",
112 pad_to: int | None = 1536,
113 device: str | None = None,
114 batch_size: int = 64,
115 backend: str = "auto",
116 ) -> None:
117 self._backend = _resolve_backend(backend)
118 self._model_name = model_name
119 self._pad_to = pad_to
120 self._device = device
121 self._batch_size = batch_size
122 self._model: Any | None = None
123 self._load_lock = asyncio.Lock()
125 def capabilities(self) -> LLMCapabilities:
126 return LLMCapabilities(
127 supports_multimodal_completion=False,
128 modalities_supported=("text",),
129 supports_multimodal_embedding=False,
130 supports_batch_embed=True,
131 )
133 @property
134 def backend(self) -> str:
135 return self._backend
137 async def _ensure_model(self) -> Any:
138 if self._model is not None:
139 return self._model
140 async with self._load_lock:
141 if self._model is None:
142 def _load() -> Any:
143 logger.info(
144 "local_embeddings: loading %s via %s (device=%s)",
145 self._model_name, self._backend, self._device or "auto",
146 )
147 if self._backend == "fastembed":
148 from fastembed import TextEmbedding
150 return TextEmbedding(self._model_name, cache_dir=_fastembed_cache_dir())
151 from sentence_transformers import SentenceTransformer
153 return SentenceTransformer(self._model_name, device=self._device)
155 self._model = await asyncio.to_thread(_load)
156 return self._model
158 async def embed(
159 self,
160 texts: list[str],
161 model: str | None = None, # noqa: ARG002 — single local model
162 ) -> list[list[float]]:
163 if not texts:
164 return []
165 model_obj = await self._ensure_model()
166 safe = [t[:_MAX_CHARS] if t else " " for t in texts]
168 def _encode() -> list[list[float]]:
169 if self._backend == "fastembed":
170 # Normalised explicitly, matching normalize_embeddings=True
171 # below. Pure Python keeps numpy out of the core package; the
172 # cost is negligible beside model inference.
173 out: list[list[float]] = []
174 for vec in model_obj.embed(safe, batch_size=self._batch_size):
175 vals = [float(x) for x in vec]
176 norm = math.sqrt(sum(x * x for x in vals)) or 1.0
177 out.append([x / norm for x in vals])
178 return out
179 vecs = model_obj.encode(
180 safe,
181 batch_size=self._batch_size,
182 normalize_embeddings=True,
183 show_progress_bar=False,
184 )
185 return [v.tolist() for v in vecs]
187 raw = await asyncio.to_thread(_encode)
189 # 0 means "native width" as well as None: provider config passes
190 # through config_kwargs, which drops None values, so YAML can only
191 # express "no padding" as 0.
192 if not self._pad_to:
193 return raw
194 dim = len(raw[0]) if raw else 0
195 if dim > self._pad_to:
196 raise ValueError(
197 f"local_embeddings: model dim {dim} exceeds pad_to={self._pad_to}; "
198 "truncation would change similarity geometry — pick a smaller "
199 "model or raise pad_to (requires a schema migration)."
200 )
201 if dim == self._pad_to:
202 return raw
203 pad = [0.0] * (self._pad_to - dim)
204 return [v + pad for v in raw]
206 async def complete(
207 self,
208 messages: list[Message],
209 model: str | None = None,
210 max_tokens: int = 1024,
211 temperature: float = 0.0,
212 tools: list[ToolDefinition] | None = None,
213 tool_choice: str | None = None,
214 response_format: dict | None = None,
215 ) -> Completion:
216 raise NotImplementedError(
217 "LocalEmbeddingsProvider is embed-only. Compose with a completion "
218 "provider via CompositeLLMProvider (astrocyte.providers.composite)."
219 )