Coverage for src/local_deep_research/embeddings/providers/implementations/sentence_transformers.py: 100%
60 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-20 01:24 +0000
1"""Sentence Transformers embedding provider."""
3from pathlib import Path
4from typing import Any, Dict, List, Optional
6from langchain_core.embeddings import Embeddings
7from ....security.secure_logging import logger
9from ....config.thread_settings import get_setting_from_snapshot
10from ..base import BaseEmbeddingProvider, Exposure
13class SentenceTransformersProvider(BaseEmbeddingProvider):
14 """
15 Sentence Transformers embedding provider.
17 Uses HuggingFace sentence-transformers models for local embeddings.
18 No API key required, runs entirely locally.
19 """
21 provider_name = "Sentence Transformers"
22 provider_key = "SENTENCE_TRANSFORMERS"
23 requires_api_key = False
24 supports_local = True
25 egress_exposure = Exposure.CONTAINED
26 default_model = "all-MiniLM-L6-v2" # type: ignore[assignment]
28 # Available models with metadata
29 AVAILABLE_MODELS = {
30 "all-MiniLM-L6-v2": {
31 "dimensions": 384,
32 "description": "Fast, lightweight model. Good for general use.",
33 "max_seq_length": 256,
34 },
35 "all-mpnet-base-v2": {
36 "dimensions": 768,
37 "description": "Higher quality, slower. Best accuracy.",
38 "max_seq_length": 384,
39 },
40 "multi-qa-MiniLM-L6-cos-v1": {
41 "dimensions": 384,
42 "description": "Optimized for question-answering tasks.",
43 "max_seq_length": 512,
44 },
45 "paraphrase-multilingual-MiniLM-L12-v2": {
46 "dimensions": 384,
47 "description": "Supports multiple languages.",
48 "max_seq_length": 128,
49 },
50 }
52 @classmethod
53 def create_embeddings(
54 cls,
55 model: Optional[str] = None,
56 settings_snapshot: Optional[Dict[str, Any]] = None,
57 **kwargs,
58 ) -> Embeddings:
59 """
60 Create Sentence Transformers embeddings instance.
62 Args:
63 model: Model name (defaults to all-MiniLM-L6-v2)
64 settings_snapshot: Optional settings snapshot
65 **kwargs: Additional parameters (device, etc.)
67 Returns:
68 SentenceTransformerEmbeddings instance
69 """
70 from langchain_community.embeddings import (
71 SentenceTransformerEmbeddings,
72 )
74 # Get model from settings if not specified
75 if model is None:
76 model = get_setting_from_snapshot(
77 "embeddings.sentence_transformers.model",
78 default=cls.default_model,
79 settings_snapshot=settings_snapshot,
80 )
82 # Get device setting (cpu or cuda)
83 device = kwargs.get("device")
84 if device is None:
85 device = get_setting_from_snapshot(
86 "embeddings.sentence_transformers.device",
87 default="cpu",
88 settings_snapshot=settings_snapshot,
89 )
91 logger.info(
92 f"Creating SentenceTransformerEmbeddings with model={model}, device={device}"
93 )
95 # Egress policy: if the user opted into local-only embeddings,
96 # refuse to trigger a HuggingFace download on first use. The
97 # SentenceTransformer constructor reaches out to huggingface.co
98 # when the requested model isn't cached locally — silent
99 # outbound traffic that violates ``embeddings.require_local=True``.
100 #
101 # Resolve the requirement through the egress CONTEXT, not the raw
102 # ``embeddings.require_local`` flag: under PRIVATE_ONLY,
103 # context_from_snapshot forces require_local_embeddings=True even
104 # when the user left the flag at its default False. Reading the raw
105 # flag here would let a PRIVATE_ONLY (offline) run silently download
106 # an uncached model from HuggingFace. An unknown/corrupt scope
107 # raises PolicyDeniedError out of context_from_snapshot — fail
108 # closed, do not download.
109 require_local = False
110 if settings_snapshot is not None:
111 try:
112 from ....security.egress.policy import (
113 context_from_snapshot,
114 resolve_run_primary_engine,
115 )
117 # Single source of truth for the primary (was: search.tool +
118 # searxng fallback, a fail-OPEN that could permit a remote model
119 # download for a primary-less private run).
120 _primary = resolve_run_primary_engine(settings_snapshot)
121 require_local = context_from_snapshot(
122 settings_snapshot, _primary
123 ).require_local_embeddings
124 except ValueError:
125 # No usable primary / invalid scope: fail CLOSED to local-only
126 # (block any remote model download) rather than reading the raw
127 # opt-in flag. The get_embeddings PEP already refuses a
128 # primary-less snapshot upstream, so this is defense-in-depth.
129 require_local = True
130 model_kwargs = {"device": device}
131 if require_local:
132 if not cls._is_model_cached_locally(model):
133 from ....security.egress.policy import (
134 Decision,
135 PolicyDeniedError,
136 )
138 # Render the model into the message (loguru brace
139 # formatting), not as a bound kwarg: kwargs land in
140 # ``record["extra"]``, which none of our sinks render,
141 # so a bound value is invisible to operators. ``{!r}``
142 # also makes a degenerate empty/whitespace config
143 # self-evident (shows as ``''``).
144 logger.bind(policy_audit=True).warning(
145 "refusing SentenceTransformer download for {!r} "
146 "under embeddings.require_local=True",
147 model,
148 )
149 raise PolicyDeniedError(
150 Decision(False, "embeddings_model_not_cached"),
151 target=model,
152 )
153 # Force the inner ``transformers``/``sentence_transformers``
154 # call to use cached files only. Defence in depth for the
155 # HF-cache branch — but LOAD-BEARING for the local-path admit
156 # in ``_is_model_cached_locally``: that admit only checks the
157 # path exists, doing zero content validation, so an existing
158 # model dir whose config references an UNCACHED remote base
159 # would otherwise fetch it here. Keep this set for every
160 # require_local admit; do not relax it on the assumption that
161 # the pre-flight already proved the model fully local.
162 model_kwargs["local_files_only"] = True
164 return SentenceTransformerEmbeddings(
165 model_name=model,
166 model_kwargs=model_kwargs,
167 )
169 @staticmethod
170 def _is_model_cached_locally(model_name: str) -> bool:
171 """Best-effort check whether ``model_name`` is already cached.
173 Probes the HuggingFace hub cache for both the bare and the
174 ``sentence-transformers/``-namespaced forms of the name (the
175 two cache keys the SentenceTransformer loader can request for
176 a bare input). Returns False if the lookup itself fails, which
177 fails closed under ``require_local=True``. A degenerate
178 model_name (None, empty, whitespace) also fails closed: the
179 cache probe below either raises — caught by the broad except —
180 or misses, so no special-casing is needed.
181 """
182 try:
183 # A model configured as an existing local path is already on
184 # disk: ``SentenceTransformer.__init__`` takes its
185 # ``os.path.exists`` branch and never touches the network, so
186 # it's admissible under require_local. Mirror that guard first
187 # — before treating the name as an HF repo_id, since a path
188 # with "/" would otherwise be probed as a repo_id and rejected
189 # as malformed (a false ``embeddings_model_not_cached``).
190 #
191 # The ``model_name and`` guard is load-bearing: ``Path("")``
192 # normalises to ``.`` (the cwd, which exists), so a blank
193 # config would wrongly admit — whereas the loader's
194 # ``os.path.exists("")`` returns False. The truthiness check
195 # restores fail-closed parity for the degenerate empty name.
196 if model_name and Path(model_name).exists():
197 return True
199 from huggingface_hub import try_to_load_from_cache
201 # Probe both the bare and the namespaced cache keys when the
202 # input has no "/", because ``SentenceTransformer.__init__``
203 # resolves bare names two different ways:
204 # - names in its ``basic_transformer_models`` allowlist
205 # (bert-base-uncased, gpt2, t5-base, ...) are requested
206 # from the BARE repo_id;
207 # - everything else is prefixed with
208 # ``__MODEL_HUB_ORGANIZATION__`` (e.g. "all-MiniLM-L6-v2"
209 # -> "sentence-transformers/all-MiniLM-L6-v2").
210 # The HF hub cache is keyed on whichever form the loader
211 # requests, so probing both is the only way to be correct
212 # for both classes without mirroring the upstream allowlist
213 # (a local list inside ``SentenceTransformer.__init__``).
214 if "/" in model_name:
215 candidates = [model_name]
216 else:
217 from sentence_transformers import (
218 __MODEL_HUB_ORGANIZATION__,
219 )
221 candidates = [
222 model_name,
223 f"{__MODEL_HUB_ORGANIZATION__}/{model_name}",
224 ]
226 # try_to_load_from_cache returns a path string when cached,
227 # None when missing, and the sentinel _CACHED_NO_EXIST for
228 # known-absent. Treat anything but a string path as a miss.
229 for repo_id in candidates:
230 cached = try_to_load_from_cache(
231 repo_id=repo_id, filename="config.json"
232 )
233 if isinstance(cached, str) and bool(cached):
234 return True
235 return False
236 except Exception: # pragma: no cover - defensive
237 return False
239 @classmethod
240 def is_available(
241 cls, settings_snapshot: Optional[Dict[str, Any]] = None
242 ) -> bool:
243 """
244 Check if Sentence Transformers is available.
246 Since sentence-transformers is a required dependency, this always returns True.
247 This method exists for API consistency with other providers.
248 """
249 return True
251 @classmethod
252 def get_available_models(
253 cls, settings_snapshot: Optional[Dict[str, Any]] = None
254 ) -> List[Dict[str, Any]]:
255 """
256 Get list of available Sentence Transformer models.
258 Note: Since there's no centralized API for Sentence Transformers,
259 we return a curated list of commonly used models. Users can also
260 specify any model name from HuggingFace directly in settings.
261 """
262 return [
263 {
264 "value": model,
265 "label": f"{model} ({info['dimensions']}d) - {info['description']}",
266 }
267 for model, info in cls.AVAILABLE_MODELS.items()
268 ]