Coverage for src/local_deep_research/web_search_engines/engines/search_engine_collection.py: 99%
106 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +0000
1"""
2Collection-specific RAG Search Engine
4Provides semantic search within a specific document collection using RAG.
5"""
7from typing import List, Dict, Any, Optional
9from ...security.secure_logging import logger
10from .search_engine_library import LibraryRAGSearchEngine
11from ...constants import (
12 DEFAULT_LOCAL_SEARCH_CHUNK_OVERLAP,
13 DEFAULT_LOCAL_SEARCH_CHUNK_SIZE,
14 DEFAULT_LOCAL_SEARCH_DISTANCE_METRIC,
15 DEFAULT_LOCAL_SEARCH_INDEX_TYPE,
16 DEFAULT_LOCAL_SEARCH_NORMALIZE_VECTORS,
17 DEFAULT_LOCAL_SEARCH_SPLITTER_TYPE,
18 SNIPPET_LENGTH_LONG,
19)
20from ...research_library.services.library_rag_service import LibraryRAGService
21from ...database.models.library import Collection
22from ...database.session_context import get_user_db_session
23from ...utilities.chunk_anchor import extract_chunk_index, extract_document_id
24from ...utilities.type_utils import to_bool
27class CollectionSearchEngine(LibraryRAGSearchEngine):
28 """
29 Search engine for a specific document collection using RAG.
30 Directly searches only the specified collection's FAISS index.
31 Each collection uses its own embedding model that was used during indexing.
32 """
34 # Mark as local RAG engine
35 is_local = True
37 def __init__(
38 self,
39 collection_id: str,
40 collection_name: str,
41 llm: Optional[Any] = None,
42 max_filtered_results: Optional[int] = None,
43 max_results: int = 10,
44 settings_snapshot: Optional[Dict[str, Any]] = None,
45 search_snippets_only: bool = True,
46 **kwargs,
47 ):
48 """
49 Initialize the collection-specific search engine.
51 Args:
52 collection_id: UUID of the collection to search within
53 collection_name: Name of the collection for display
54 llm: Language model for relevance filtering
55 max_filtered_results: Maximum number of results to keep after filtering
56 max_results: Maximum number of search results
57 settings_snapshot: Settings snapshot from thread context
58 search_snippets_only: named explicitly because create_search_engine
59 filters its params against this signature and drops the rest
60 **kwargs: Additional engine-specific parameters
61 """
62 super().__init__(
63 llm=llm,
64 max_filtered_results=max_filtered_results,
65 max_results=max_results,
66 settings_snapshot=settings_snapshot,
67 search_snippets_only=search_snippets_only,
68 **kwargs,
69 )
70 self.collection_id = collection_id
71 self.collection_name = collection_name
72 self.collection_key = f"collection_{collection_id}"
73 # The policy registry name for this engine is known at construction
74 # (the factory only stamps registry engines), so set it here —
75 # direct instantiations (e.g. the library search route) get the
76 # runtime egress backstop without caller cooperation. NB: between
77 # super().__init__ (which stamps "library") and this line the
78 # instance briefly carries the parent's name — nothing may call
79 # _verify_egress_scope() in that window.
80 self._engine_name = self.collection_key
82 # Load collection-specific embedding settings
83 self._load_collection_embedding_settings()
85 def _load_collection_embedding_settings(self):
86 """
87 Load frozen embedding settings directly from the Collection record.
88 Uses the same embedding model that was configured when the collection was created.
89 """
90 if not self.username:
91 logger.warning("Cannot load collection settings without username")
92 return
94 try:
95 with get_user_db_session(self.username) as db_session:
96 collection = (
97 db_session.query(Collection)
98 .filter_by(id=self.collection_id)
99 .first()
100 )
102 if not collection or not collection.embedding_model:
103 logger.warning(
104 f"No stored embedding settings found for collection {self.collection_id}"
105 )
106 return
108 # Use embedding settings directly from the Collection
109 self.embedding_model = collection.embedding_model
110 provider = collection.embedding_model_type
111 if provider is not None:
112 self.embedding_provider = (
113 provider.value
114 if hasattr(provider, "value")
115 else str(provider)
116 )
117 coll_chunk_size = getattr(collection, "chunk_size", None)
118 if coll_chunk_size is not None: 118 ↛ 120line 118 didn't jump to line 120 because the condition on line 118 was always true
119 self.chunk_size = int(coll_chunk_size)
120 coll_chunk_overlap = getattr(collection, "chunk_overlap", None)
121 if coll_chunk_overlap is not None: 121 ↛ 123line 121 didn't jump to line 123 because the condition on line 121 was always true
122 self.chunk_overlap = int(coll_chunk_overlap)
123 if getattr(collection, "splitter_type", None) is not None:
124 self.splitter_type = collection.splitter_type
125 if getattr(collection, "text_separators", None) is not None:
126 self.text_separators = collection.text_separators
127 if getattr(collection, "distance_metric", None) is not None:
128 self.distance_metric = collection.distance_metric
129 if getattr(collection, "normalize_vectors", None) is not None:
130 self.normalize_vectors = to_bool(
131 collection.normalize_vectors,
132 default=DEFAULT_LOCAL_SEARCH_NORMALIZE_VECTORS,
133 )
134 if getattr(collection, "index_type", None) is not None:
135 self.index_type = collection.index_type
137 logger.info(
138 f"Collection '{self.collection_name}' using embedding: "
139 f"{self.embedding_provider}/{self.embedding_model}"
140 )
142 except Exception as e:
143 safe_msg = self._scrub_error(e)
144 logger.exception(
145 f"Error loading collection {self.collection_id} settings ({type(e).__name__}): {safe_msg}"
146 )
148 def search(
149 self,
150 query: str,
151 limit: int = 10,
152 llm_callback=None,
153 extra_params: Optional[Dict[str, Any]] = None,
154 ) -> List[Dict[str, Any]]:
155 """
156 Search within the specific collection using semantic search.
158 Directly searches only this collection's FAISS index instead of
159 searching all collections and filtering.
161 Args:
162 query: Search query
163 limit: Maximum number of results to return
164 llm_callback: Optional LLM callback for processing results
165 extra_params: Additional search parameters
167 Returns:
168 List of search results from this collection
169 """
170 if not self.username:
171 logger.error("Cannot search collection without username")
172 return []
174 # search() does not go through BaseSearchEngine.run(), so apply the
175 # runtime egress backstop here as well. Raises PolicyDeniedError on
176 # denial — deliberately BEFORE the broad try/except below.
177 self._verify_egress_scope()
179 try:
180 # Get frozen embedding settings directly from Collection record
181 with get_user_db_session(self.username) as db_session:
182 collection = (
183 db_session.query(Collection)
184 .filter_by(id=self.collection_id)
185 .first()
186 )
188 if not collection or not collection.embedding_model:
189 logger.info(
190 f"No embedding settings found for collection '{self.collection_name}'"
191 )
192 return []
194 # Get embedding settings directly from Collection
195 embedding_model = collection.embedding_model
196 provider = collection.embedding_model_type
197 if provider is not None:
198 embedding_provider = (
199 provider.value
200 if hasattr(provider, "value")
201 else str(provider)
202 )
203 else:
204 embedding_provider = self.embedding_provider
205 coll_chunk_size = getattr(collection, "chunk_size", None)
206 chunk_size = (
207 int(coll_chunk_size)
208 if coll_chunk_size is not None
209 else int(self.chunk_size or DEFAULT_LOCAL_SEARCH_CHUNK_SIZE)
210 )
211 coll_chunk_overlap = getattr(collection, "chunk_overlap", None)
212 chunk_overlap = (
213 int(coll_chunk_overlap)
214 if coll_chunk_overlap is not None
215 else int(
216 self.chunk_overlap
217 if self.chunk_overlap is not None
218 else DEFAULT_LOCAL_SEARCH_CHUNK_OVERLAP
219 )
220 )
221 splitter_type = (
222 getattr(collection, "splitter_type", None)
223 or getattr(self, "splitter_type", None)
224 or DEFAULT_LOCAL_SEARCH_SPLITTER_TYPE
225 )
226 text_separators = (
227 collection.text_separators
228 if getattr(collection, "text_separators", None) is not None
229 else getattr(self, "text_separators", None)
230 )
231 # Thread the stored normalization flag AND distance
232 # metric through, exactly like LibraryRAGSearchEngine.
233 # Without normalize_vectors, LibraryRAGService defaults
234 # True and would L2-normalize the query against a
235 # raw-vector collection, flipping the top hit. Without
236 # distance_metric, the collection is labeled "cosine"
237 # and l2/dot_product hits get the wrong [0,1] mapping in
238 # the relevance transform below. NULL → prior default.
239 coll_normalize = getattr(collection, "normalize_vectors", None)
240 normalize_vectors = (
241 to_bool(
242 coll_normalize,
243 default=DEFAULT_LOCAL_SEARCH_NORMALIZE_VECTORS,
244 )
245 if coll_normalize is not None
246 else getattr(
247 self,
248 "normalize_vectors",
249 DEFAULT_LOCAL_SEARCH_NORMALIZE_VECTORS,
250 )
251 )
252 distance_metric = (
253 getattr(collection, "distance_metric", None)
254 or getattr(self, "distance_metric", None)
255 or DEFAULT_LOCAL_SEARCH_DISTANCE_METRIC
256 )
257 # Thread the stored index_type too (like metric/
258 # normalize): a legacy NULL-index_type row that is
259 # physically HNSW would otherwise fall back to "flat",
260 # mislabelling the store. NULL → "flat".
261 index_type = (
262 getattr(collection, "index_type", None)
263 or getattr(self, "index_type", None)
264 or DEFAULT_LOCAL_SEARCH_INDEX_TYPE
265 )
267 # Create RAG service with collection's embedding settings
268 with LibraryRAGService(
269 username=self.username,
270 embedding_model=embedding_model,
271 embedding_provider=embedding_provider,
272 chunk_size=chunk_size,
273 chunk_overlap=chunk_overlap,
274 splitter_type=splitter_type,
275 text_separators=text_separators,
276 normalize_vectors=normalize_vectors,
277 distance_metric=distance_metric,
278 index_type=index_type,
279 ) as rag_service:
280 # Check if there are indexed documents
281 if not rag_service.has_indexed_documents(self.collection_id):
282 logger.info(
283 f"No documents indexed in collection '{self.collection_name}'"
284 )
285 return []
287 # Search this collection's vector index via the new
288 # int-id-keyed store (chunk text is rehydrated from the
289 # encrypted DB by id — never read from the vector store
290 # itself; see the SECURITY INVARIANT in
291 # vector_stores/base.py).
292 search_results = rag_service.search(
293 query, self.collection_id, limit
294 )
296 if not search_results:
297 logger.info(
298 f"No results found in collection '{self.collection_name}'"
299 )
300 return []
302 # Convert to search result format
303 results = []
304 for r in search_results:
305 metadata = dict(r.metadata or {})
307 # Get document ID
308 doc_id = (
309 r.source_id
310 or metadata.get("source_id")
311 or metadata.get("document_id")
312 )
314 # Get title
315 title = (
316 r.document_title
317 or metadata.get("document_title")
318 or metadata.get("title")
319 or (f"Document {doc_id}" if doc_id else "Untitled")
320 )
322 # Create snippet from content
323 snippet = (
324 r.text[:SNIPPET_LENGTH_LONG] + "..."
325 if len(r.text) > SNIPPET_LENGTH_LONG
326 else r.text
327 )
329 # Generate document URL. Validate chunk_idx and
330 # sanitise doc_id through the shared helpers so a
331 # malformed chunk index (UUID, boolean, negative) or
332 # path-traversal doc_id cannot leak into the citation.
333 chunk_idx = extract_chunk_index(metadata)
334 # Authoritative id FIRST. ``extract_document_id`` scans its
335 # first mapping before any later one, so passing ``metadata``
336 # first would let a chunk's denormalised metadata override
337 # ``r.source_id`` — the DocumentChunk FK that decides which
338 # document the chunk actually belongs to. Pre-#5381 the
339 # column always won; a divergent metadata id would otherwise
340 # point the citation at a different document.
341 sanitised_doc_id = extract_document_id(
342 {"source_id": doc_id} if doc_id else None,
343 metadata,
344 )
345 document_url = self._get_document_url(
346 sanitised_doc_id, chunk_index=chunk_idx
347 )
349 # Add collection info to metadata
350 metadata["collection_id"] = self.collection_id
351 metadata["collection_name"] = (
352 self.collection_name or "Unknown"
353 )
355 # r.distance's meaning depends on r.metric. Map both to a
356 # [0, 1] relevance that grows with similarity:
357 # - cosine/dot_product (inner product on normalized
358 # vectors): a similarity in [-1, 1] where higher is
359 # nearer, so (d+1)/2 clamped to [0, 1]. Using it raw would
360 # yield negative relevance for anti-correlated hits and
361 # break downstream [0,1] consumers (similarity %, filters).
362 # - anything else (l2, or a non-standard metric string): a
363 # (squared) distance >= 0 where LOWER is nearer, so
364 # 1/(1+d) in (0, 1].
365 # This IP test MUST match faiss_store._build_base_index,
366 # which builds an inner-product index iff the metric is
367 # cosine/dot_product and an L2 index otherwise. A bare
368 # `metric == "l2"` here would score a non-standard metric
369 # (which builds an L2 index) with the IP formula, inverting
370 # the ranking.
371 relevance = (
372 max(0.0, min(1.0, (r.distance + 1.0) / 2.0))
373 if r.metric in ("cosine", "dot_product")
374 else 1.0 / (1.0 + r.distance)
375 )
377 result = {
378 "title": title,
379 "snippet": snippet,
380 "url": document_url,
381 "link": document_url,
382 "source": "library",
383 "source_type": "library",
384 # Authoritative document id at the TOP level, not only
385 # nested in ``metadata``: SearchResultsCollector rebuilds
386 # this citation URL and would otherwise have nothing but
387 # the denormalised metadata to go on.
388 "source_id": doc_id,
389 "relevance_score": float(relevance),
390 "metadata": metadata,
391 }
392 results.append(result)
394 logger.info(
395 f"Collection '{self.collection_name}' search returned "
396 f"{len(results)} results"
397 )
399 return results
401 except Exception as e:
402 # Re-raise instead of returning [] so a failed search is not
403 # indistinguishable from "no matching documents": run() records
404 # the failure in metrics, and API callers can report the error.
405 safe_msg = self._scrub_error(e)
406 logger.exception(
407 f"Error searching collection '{self.collection_name}' ({type(e).__name__}): {safe_msg}"
408 )
409 raise