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