Coverage for src/local_deep_research/embeddings/providers/implementations/sentence_transformers.py: 96%
112 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-06 15:42 +0000
1"""Sentence Transformers embedding provider."""
3import re
4from pathlib import Path
5from typing import Any, Dict, List, Optional, Union
7from langchain_core.embeddings import Embeddings
8from ....security.secure_logging import logger
10from ....config.thread_settings import get_setting_from_snapshot
11from ..base import BaseEmbeddingProvider, Exposure
14class SentenceTransformersProvider(BaseEmbeddingProvider):
15 """
16 Sentence Transformers embedding provider.
18 Uses HuggingFace sentence-transformers models for local embeddings.
19 No API key required, runs entirely locally.
20 """
22 provider_name = "Sentence Transformers"
23 provider_key = "SENTENCE_TRANSFORMERS"
24 requires_api_key = False
25 supports_local = True
26 egress_exposure = Exposure.CONTAINED
27 default_model = "all-MiniLM-L6-v2" # type: ignore[assignment]
29 # Available models with metadata
30 AVAILABLE_MODELS = {
31 "all-MiniLM-L6-v2": {
32 "dimensions": 384,
33 "description": "Fast, lightweight model. Good for general use.",
34 "max_seq_length": 256,
35 },
36 "all-mpnet-base-v2": {
37 "dimensions": 768,
38 "description": "Higher quality, slower. Best accuracy.",
39 "max_seq_length": 384,
40 },
41 "multi-qa-MiniLM-L6-cos-v1": {
42 "dimensions": 384,
43 "description": "Optimized for question-answering tasks.",
44 "max_seq_length": 512,
45 },
46 "paraphrase-multilingual-MiniLM-L12-v2": {
47 "dimensions": 384,
48 "description": "Supports multiple languages.",
49 "max_seq_length": 128,
50 },
51 }
53 @classmethod
54 def create_embeddings(
55 cls,
56 model: Optional[str] = None,
57 settings_snapshot: Optional[Dict[str, Any]] = None,
58 **kwargs,
59 ) -> Embeddings:
60 """
61 Create Sentence Transformers embeddings instance.
63 Args:
64 model: Model name (defaults to all-MiniLM-L6-v2)
65 settings_snapshot: Optional settings snapshot
66 **kwargs: Additional parameters (device, etc.)
68 Returns:
69 SentenceTransformerEmbeddings instance
70 """
71 from langchain_community.embeddings import (
72 SentenceTransformerEmbeddings,
73 )
75 # Get model from settings if not specified
76 if model is None:
77 model = get_setting_from_snapshot(
78 "embeddings.sentence_transformers.model",
79 default=cls.default_model,
80 settings_snapshot=settings_snapshot,
81 )
83 # Get device setting (cpu or cuda)
84 device = kwargs.get("device")
85 if device is None:
86 device = get_setting_from_snapshot(
87 "embeddings.sentence_transformers.device",
88 default="cpu",
89 settings_snapshot=settings_snapshot,
90 )
92 logger.info(
93 f"Creating SentenceTransformerEmbeddings with model={model}, device={device}"
94 )
96 # Path confinement (security): the embedding model setting is
97 # user-editable free text. ``SentenceTransformer`` loads ANY value
98 # that resolves to an existing filesystem path (its ``os.path.exists``
99 # branch), so without a guard an authenticated user could point this
100 # setting at an arbitrary absolute server path — turning it into a
101 # whole-filesystem existence/structure probe and an arbitrary
102 # local-model loader. Classify the value ONCE:
103 # * a file/dir confined UNDER the app's models directory is a
104 # legitimate local model -> load it from its safe absolute path;
105 # * a filesystem-path-shaped value that ESCAPES the models dir is
106 # refused outright (never handed to the loader);
107 # * anything else is treated as a HuggingFace repo id (default).
108 local_model_path = cls._confined_local_model_path(model)
109 if local_model_path is None and cls._looks_like_filesystem_path(model):
110 logger.bind(policy_audit=True).warning(
111 "refusing sentence-transformers embedding model {!r}: "
112 "filesystem paths must resolve under the app models directory",
113 model,
114 )
115 raise ValueError(
116 "Invalid embedding model path: local model paths must "
117 "resolve under the application's models directory."
118 )
120 # Egress policy: if the user opted into local-only embeddings,
121 # refuse to trigger a HuggingFace download on first use. The
122 # SentenceTransformer constructor reaches out to huggingface.co
123 # when the requested model isn't cached locally — silent
124 # outbound traffic that violates ``embeddings.require_local=True``.
125 #
126 # Resolve the requirement through the egress CONTEXT, not the raw
127 # ``embeddings.require_local`` flag: under PRIVATE_ONLY,
128 # context_from_snapshot forces require_local_embeddings=True even
129 # when the user left the flag at its default False. Reading the raw
130 # flag here would let a PRIVATE_ONLY (offline) run silently download
131 # an uncached model from HuggingFace. An unknown/corrupt scope
132 # raises PolicyDeniedError out of context_from_snapshot — fail
133 # closed, do not download.
134 require_local = False
135 if settings_snapshot is not None:
136 try:
137 from ....security.egress.policy import (
138 context_from_snapshot,
139 resolve_run_primary_engine,
140 )
141 from ....search_system import username_from_snapshot
143 # Single source of truth for the primary (was: search.tool +
144 # searxng fallback, a fail-OPEN that could permit a remote model
145 # download for a primary-less private run).
146 _primary = resolve_run_primary_engine(settings_snapshot)
147 require_local = context_from_snapshot(
148 settings_snapshot,
149 _primary,
150 username=username_from_snapshot(settings_snapshot),
151 ).require_local_embeddings
152 except ValueError:
153 # No usable primary / invalid scope: fail CLOSED to local-only
154 # (block any remote model download) rather than reading the raw
155 # opt-in flag. The get_embeddings PEP already refuses a
156 # primary-less snapshot upstream, so this is defense-in-depth.
157 require_local = True
158 model_kwargs = {"device": device}
159 if require_local:
160 if not cls._is_model_cached_locally(model):
161 from ....security.egress.policy import (
162 Decision,
163 PolicyDeniedError,
164 )
166 # Render the model into the message (loguru brace
167 # formatting), not as a bound kwarg: kwargs land in
168 # ``record["extra"]``, which none of our sinks render,
169 # so a bound value is invisible to operators. ``{!r}``
170 # also makes a degenerate empty/whitespace config
171 # self-evident (shows as ``''``).
172 logger.bind(policy_audit=True).warning(
173 "refusing SentenceTransformer download for {!r} "
174 "under embeddings.require_local=True",
175 model,
176 )
177 raise PolicyDeniedError(
178 Decision(False, "embeddings_model_not_cached"),
179 target=model,
180 )
181 # Force the inner ``transformers``/``sentence_transformers``
182 # call to use cached files only. Defence in depth for the
183 # HF-cache branch — but LOAD-BEARING for the local-path admit
184 # in ``_is_model_cached_locally``: that admit only checks the
185 # path exists, doing zero content validation, so an existing
186 # model dir whose config references an UNCACHED remote base
187 # would otherwise fetch it here. Keep this set for every
188 # require_local admit; do not relax it on the assumption that
189 # the pre-flight already proved the model fully local.
190 model_kwargs["local_files_only"] = True
192 # Load a confined local model from its safe absolute path; otherwise
193 # hand the (repo-id) value to the loader unchanged.
194 return SentenceTransformerEmbeddings(
195 model_name=(
196 str(local_model_path) if local_model_path is not None else model
197 ),
198 model_kwargs=model_kwargs,
199 )
201 @classmethod
202 def _is_model_cached_locally(cls, model_name: str) -> bool:
203 """Best-effort check whether ``model_name`` is available WITHOUT any
204 network access.
206 Admissible in two ordered cases:
208 1. ``model_name`` resolves to an existing file/dir CONFINED under the
209 application's models directory (see
210 :meth:`_confined_local_model_path`). Such a model is already on
211 disk and ``SentenceTransformer.__init__`` loads it via its
212 ``os.path.exists`` branch without touching the network.
214 Security: the confinement check decides containment lexically and
215 NEVER probes an arbitrary user-supplied path with
216 ``Path(model_name).exists()``. A path outside the models dir can
217 therefore no longer be used as a whole-filesystem existence oracle.
219 2. Otherwise ``model_name`` is treated as a HuggingFace repo id and the
220 hub cache is probed for both the bare and the
221 ``sentence-transformers/``-namespaced forms (the two cache keys the
222 loader can request for a bare input).
224 Returns False if the lookup itself fails, which fails closed under
225 ``require_local=True``. A degenerate model_name (None, empty,
226 whitespace) also fails closed.
227 """
228 try:
229 # 1. A legitimate local model confined under the app models dir.
230 if cls._confined_local_model_path(model_name) is not None:
231 return True
233 # 2. Treat the value as an HF repo id and probe the hub cache
234 # ONLY. try_to_load_from_cache keys on the string as a repo_id
235 # inside the HF cache directory, so a non-confined path is never
236 # reached on the real filesystem here.
237 from huggingface_hub import try_to_load_from_cache
239 # Probe both the bare and the namespaced cache keys when the
240 # input has no "/", because ``SentenceTransformer.__init__``
241 # resolves bare names two different ways:
242 # - names in its ``basic_transformer_models`` allowlist
243 # (bert-base-uncased, gpt2, t5-base, ...) are requested
244 # from the BARE repo_id;
245 # - everything else is prefixed with
246 # ``__MODEL_HUB_ORGANIZATION__`` (e.g. "all-MiniLM-L6-v2"
247 # -> "sentence-transformers/all-MiniLM-L6-v2").
248 # The HF hub cache is keyed on whichever form the loader
249 # requests, so probing both is the only way to be correct
250 # for both classes without mirroring the upstream allowlist
251 # (a local list inside ``SentenceTransformer.__init__``).
252 if "/" in model_name:
253 candidates = [model_name]
254 else:
255 from sentence_transformers import (
256 __MODEL_HUB_ORGANIZATION__,
257 )
259 candidates = [
260 model_name,
261 f"{__MODEL_HUB_ORGANIZATION__}/{model_name}",
262 ]
264 # try_to_load_from_cache returns a path string when cached,
265 # None when missing, and the sentinel _CACHED_NO_EXIST for
266 # known-absent. Treat anything but a string path as a miss.
267 for repo_id in candidates:
268 cached = try_to_load_from_cache(
269 repo_id=repo_id, filename="config.json"
270 )
271 if isinstance(cached, str) and bool(cached):
272 return True
273 return False
274 except Exception: # pragma: no cover - defensive
275 return False
277 @classmethod
278 def _confined_local_model_path(
279 cls,
280 model_name: Optional[str],
281 models_dir: Optional[Union[str, Path]] = None,
282 ) -> Optional[Path]:
283 """Resolve ``model_name`` to a real local model path IFF it denotes an
284 EXISTING location confined UNDER the application's models directory.
286 Returns the resolved :class:`~pathlib.Path` for a legitimate local
287 model (a file or directory living under the models dir), or ``None``
288 for everything else — HuggingFace repo ids and, crucially, any
289 absolute/relative path that escapes the models directory.
291 Security: a value that is NOT confined under the models dir is
292 rejected by LEXICAL containment (``is_relative_to``) and is never
293 probed on disk with ``Path(model_name).exists()``. This is what closes
294 the arbitrary-filesystem existence oracle. Only a value already
295 confined under the models dir is resolved and stat-ed.
296 """
297 if not isinstance(model_name, str):
298 return None
299 text = model_name.strip()
300 if not text:
301 return None
303 try:
304 if models_dir is None: 304 ↛ 328line 304 didn't jump to line 328 because the condition on line 304 was always true
305 from ....config.paths import get_models_directory
307 models_dir = get_models_directory()
308 # Two views of the models root are needed:
309 # * ``models_root_lexical`` -- absolute + lexically normalized
310 # but NOT symlink-resolved. Used for the LEXICAL containment
311 # check on absolute inputs so a legitimate in-tree model still
312 # matches when an ANCESTOR of the models dir is a symlink
313 # (NFS-mounted homes, a symlinked ``LDR_DATA_DIR``, macOS
314 # ``/tmp`` -> ``/private/tmp``). Resolving the root but NOT the
315 # user path is precisely the mismatch that wrongly refused
316 # legitimate models.
317 # * ``models_root`` -- fully symlink-resolved. Used ONLY for the
318 # secondary post-resolution escape re-check below, which
319 # collapses any symlink in the confined candidate and confirms
320 # it still lands inside the resolved tree (catches an in-tree
321 # symlink that points OUT of the models dir).
322 #
323 # ``.absolute()`` makes the path absolute WITHOUT resolving
324 # symlinks (pathlib already collapses ``.`` and redundant
325 # separators on construction); it is used deliberately here in
326 # place of ``.resolve()`` so the lexical comparison below stays
327 # symlink-agnostic.
328 models_root_lexical = Path(models_dir).absolute()
329 models_root = Path(models_dir).resolve()
330 except Exception: # pragma: no cover - defensive
331 return None
333 from ....security.path_validator import PathValidator
335 try:
336 if Path(text).is_absolute():
337 # Absolute reference: decide containment LEXICALLY first so a
338 # path outside the models dir is rejected without any
339 # filesystem access. Compare ABSOLUTE-but-UN-RESOLVED paths on
340 # BOTH sides (``lexical`` here vs ``models_root_lexical`` above,
341 # never ``.resolve()``) so a symlinked ANCESTOR of the models
342 # dir does not spuriously break the match for a legitimate
343 # in-tree model. The ``..`` guard runs first, on the parts,
344 # so a traversal segment is refused regardless. Only a path
345 # already lexically inside is then resolved (collapsing any
346 # symlink escape) and re-checked against the resolved root
347 # below.
348 lexical = Path(text)
349 if ".." in lexical.parts or not lexical.is_relative_to(
350 models_root_lexical
351 ):
352 return None
353 candidate = lexical.resolve()
354 else:
355 # Relative reference: safe_join confines it under the models
356 # dir, rejecting traversal / absolute inputs (returns None or
357 # raises ValueError). A bare HF id like "all-MiniLM-L6-v2" or
358 # "sentence-transformers/all-MiniLM-L6-v2" confines cleanly but
359 # simply won't exist under the models dir -> falls through to
360 # None and is handled as a repo id by the caller.
361 confined = PathValidator.validate_safe_path(text, models_root)
362 if confined is None: 362 ↛ 363line 362 didn't jump to line 363 because the condition on line 362 was never true
363 return None
364 candidate = confined.resolve()
365 except (ValueError, OSError, RuntimeError):
366 # RuntimeError: ``.resolve()`` raises this (not OSError) when it
367 # detects a symlink LOOP (e.g. a -> b -> a) while collapsing the
368 # path. Caught here so a loop fails closed with a clean refusal
369 # instead of propagating an uncaught RuntimeError out of this
370 # function / ``create_embeddings``.
371 return None
373 # Secondary post-resolution escape re-check (defence in depth): both
374 # ``candidate`` (each branch above resolved it) and ``models_root`` are
375 # fully symlink-resolved here, so this catches an in-tree symlink whose
376 # target points OUT of the models dir — which the lexical check above
377 # deliberately cannot see. Also require a STRICT subpath (never the
378 # models root itself). is_relative_to stays a pure lexical check.
379 if candidate == models_root or not candidate.is_relative_to(
380 models_root
381 ):
382 return None
384 try:
385 if candidate.exists():
386 return candidate
387 except OSError: # pragma: no cover - defensive
388 return None
389 return None
391 @staticmethod
392 def _looks_like_filesystem_path(model_name: Optional[str]) -> bool:
393 """Pure string classification: does ``model_name`` look like a
394 filesystem path rather than a HuggingFace repo id?
396 Returns True only for the unambiguous path shapes used to escape the
397 models directory — absolute paths, home-relative (``~``) refs, Windows
398 drive paths, and any reference containing a parent-traversal (``..``)
399 segment. HuggingFace repo ids such as ``all-MiniLM-L6-v2`` or
400 ``sentence-transformers/all-MiniLM-L6-v2`` return False. Performs no
401 filesystem access.
402 """
403 if not isinstance(model_name, str):
404 return False
405 text = model_name.strip()
406 if not text: 406 ↛ 407line 406 didn't jump to line 407 because the condition on line 406 was never true
407 return False
408 if text.startswith(("/", "\\", "~")) or Path(text).is_absolute():
409 return True
410 # Windows drive-letter absolute path, e.g. "C:\\models".
411 if len(text) >= 2 and text[1] == ":": 411 ↛ 412line 411 didn't jump to line 412 because the condition on line 411 was never true
412 return True
413 # Parent-directory traversal in any separator form.
414 if ".." in re.split(r"[\\/]+", text):
415 return True
416 return False
418 @classmethod
419 def is_available(
420 cls, settings_snapshot: Optional[Dict[str, Any]] = None
421 ) -> bool:
422 """
423 Check if Sentence Transformers is available.
425 Since sentence-transformers is a required dependency, this always returns True.
426 This method exists for API consistency with other providers.
427 """
428 return True
430 @classmethod
431 def get_available_models(
432 cls, settings_snapshot: Optional[Dict[str, Any]] = None
433 ) -> List[Dict[str, Any]]:
434 """
435 Get list of available Sentence Transformer models.
437 Note: Since there's no centralized API for Sentence Transformers,
438 we return a curated list of commonly used models. Users can also
439 specify any model name from HuggingFace directly in settings.
440 """
441 return [
442 {
443 "value": model,
444 "label": f"{model} ({info['dimensions']}d) - {info['description']}",
445 }
446 for model, info in cls.AVAILABLE_MODELS.items()
447 ]