Coverage for src/local_deep_research/web_search_engines/engines/local_embedding_manager.py: 100%
76 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-06 15:42 +0000
1import threading
2from typing import Any, Dict, Optional
5from ...database.models.library import DocumentChunk
6from ...database.session_context import get_user_db_session
7from ...constants import (
8 DEFAULT_LOCAL_SEARCH_MODEL,
9 DEFAULT_LOCAL_SEARCH_PROVIDER,
10)
11from ...security.log_sanitizer import scrub_error
12from ...security.secure_logging import logger
13from ...utilities.url_utils import normalize_url
16class LocalEmbeddingManager:
17 """Handles embedding generation and storage for local document search"""
19 def __init__(
20 self,
21 embedding_model: str = DEFAULT_LOCAL_SEARCH_MODEL,
22 embedding_device: str = "cpu",
23 embedding_model_type: str = DEFAULT_LOCAL_SEARCH_PROVIDER, # or 'ollama'
24 ollama_base_url: Optional[str] = None,
25 settings_snapshot: Optional[Dict[str, Any]] = None,
26 ):
27 """
28 Initialize the embedding manager for local document search.
30 Args:
31 embedding_model: Name of the embedding model to use
32 embedding_device: Device to run embeddings on ('cpu' or 'cuda')
33 embedding_model_type: Type of embedding model ('sentence_transformers' or 'ollama')
34 ollama_base_url: Base URL for Ollama API if using ollama embeddings
35 settings_snapshot: Optional settings snapshot for background threads
36 """
38 self.embedding_model = embedding_model
39 self.embedding_device = embedding_device
40 self.embedding_model_type = embedding_model_type
41 self.ollama_base_url = ollama_base_url
42 self.settings_snapshot = settings_snapshot or {}
44 # Username for database access (extracted from settings if available)
45 self.username = (
46 settings_snapshot.get("_username") if settings_snapshot else None
47 )
48 # Password for encrypted database access (can be set later)
49 self.db_password = None
51 # Initialize the embedding model (with lock for thread-safe lazy init)
52 self._embeddings = None
53 self._embedding_lock = threading.Lock()
55 # Track if this manager has been closed
56 self._closed = False
58 def close(self):
59 """Release embedding model resources.
61 For Ollama embeddings, this also closes the underlying per-instance
62 ``httpx.Client`` / ``httpx.AsyncClient`` pair. langchain_ollama's
63 ``OllamaEmbeddings`` eagerly constructs both clients in its Pydantic
64 ``@model_validator(mode="after")``, so dropping the Python reference
65 alone leaks ~2 FDs per instance — see the migration regression note
66 in docs/developing/resource-cleanup.md. Non-Ollama providers
67 (sentence_transformers, OpenAI's lru_cache'd shared client) are
68 no-ops via the module-prefix check inside ``_close_base_llm``.
69 """
70 if self._closed:
71 return
72 self._closed = True
73 if self._embeddings is not None:
74 from ...utilities.llm_utils import _close_base_llm
76 _close_base_llm(self._embeddings)
77 self._embeddings = None
78 logger.debug("LocalEmbeddingManager closed")
80 def __enter__(self):
81 """Context manager entry."""
82 return self
84 def __exit__(self, exc_type, exc_val, exc_tb):
85 """Context manager exit - ensures resources are released."""
86 self.close()
87 return False
89 @property
90 def embeddings(self):
91 """
92 Lazily initialize embeddings when first accessed.
93 This allows the LocalEmbeddingManager to be created without
94 immediately loading models, which is helpful when no local search is performed.
96 Uses double-checked locking to ensure thread-safe initialization.
97 Concurrent SentenceTransformer model loading causes meta tensor errors
98 in PyTorch when multiple threads call model.to(device) simultaneously.
99 """
100 if self._embeddings is None:
101 with self._embedding_lock:
102 if self._embeddings is None:
103 logger.info("Initializing embeddings on first use")
104 self._embeddings = self._initialize_embeddings()
105 return self._embeddings
107 def _initialize_embeddings(self):
108 """Initialize the embedding model based on configuration"""
109 try:
110 # Use the new unified embedding system
111 from ...embeddings import get_embeddings
113 # Prepare kwargs for provider-specific parameters
114 kwargs = {}
116 # Add device for sentence transformers
117 if self.embedding_model_type == "sentence_transformers":
118 kwargs["device"] = self.embedding_device
120 # Add base_url for ollama if specified
121 if self.embedding_model_type == "ollama" and self.ollama_base_url:
122 kwargs["base_url"] = normalize_url(self.ollama_base_url)
124 logger.info(
125 f"Initializing embeddings with provider={self.embedding_model_type}, model={self.embedding_model}"
126 )
128 return get_embeddings(
129 provider=self.embedding_model_type,
130 model=self.embedding_model,
131 settings_snapshot=self.settings_snapshot,
132 **kwargs,
133 )
134 except ImportError as exc:
135 # Only fall back when the configured provider's dependency
136 # genuinely isn't installed — that's a deployment shape, not
137 # a transient runtime error. Any OTHER exception (Ollama
138 # DNS hiccup, provider validation, policy denial) must
139 # propagate so we don't silently fetch from huggingface.co
140 # when the user has explicitly opted into local embeddings.
141 safe_msg = scrub_error(exc)
142 logger.exception(
143 f"Embedding provider import failed ({type(exc).__name__}) — falling back to local SBERT: {safe_msg}"
144 )
145 # Route the fallback through get_embeddings(sentence_transformers)
146 # rather than constructing HuggingFaceEmbeddings directly: the SBERT
147 # model is fetched from huggingface.co on a cache miss, which the
148 # provider gate refuses under PRIVATE_ONLY / embeddings.require_local.
149 # Constructing it raw here would bypass that gate and leak an
150 # outbound HF download in offline mode. PolicyDeniedError from the
151 # gate propagates (fail closed).
152 from ...embeddings import get_embeddings
154 return get_embeddings(
155 provider="sentence_transformers",
156 model=None, # provider default (all-MiniLM-L6-v2)
157 settings_snapshot=self.settings_snapshot,
158 )
160 # NOTE(cutover): the old ``_store_chunks_to_db`` (LangChain-Document-list
161 # -> DocumentChunk rows, with chunk-hash+collection dedup/reuse for a
162 # FAISS uuid-keyed docstore) was deleted here. Its caller
163 # ``LibraryRAGService.index_document`` now writes chunk rows via
164 # ``vector_stores.facade.VectorIndex.index`` instead (int-id keyed,
165 # one row per source's chunk, no cross-source hash dedup/reuse; see
166 # the "One row = one vector" note in vector_stores/facade.py). Grep
167 # confirmed no other production callers before deletion.
169 def _delete_chunks_from_db(
170 self,
171 collection_name: str,
172 source_path: Optional[str] = None,
173 source_id: Optional[int] = None,
174 ) -> int:
175 """
176 Delete chunks from database.
178 Args:
179 collection_name: Name of the collection
180 source_path: Path to source file (for local files)
181 source_id: ID of source document (for library documents)
183 Returns:
184 Number of chunks deleted
185 """
186 if not self.username:
187 logger.warning(
188 "No username available, cannot delete chunks from database"
189 )
190 return 0
192 try:
193 with get_user_db_session(
194 self.username, self.db_password
195 ) as session:
196 query = session.query(DocumentChunk).filter_by(
197 collection_name=collection_name
198 )
200 if source_path:
201 query = query.filter_by(source_path=str(source_path))
202 if source_id:
203 query = query.filter_by(source_id=source_id)
205 count = int(query.delete())
206 session.commit()
208 logger.info(
209 f"Deleted {count} chunks from database for collection '{collection_name}'"
210 )
211 return count
213 except Exception as e:
214 safe_msg = scrub_error(e, self.db_password)
215 logger.exception(
216 f"Error deleting chunks from database for collection '{collection_name}' ({type(e).__name__}): {safe_msg}"
217 )
218 return 0