Coverage for src/local_deep_research/web_search_engines/engines/search_engine_collection.py: 98%
100 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-20 01:24 +0000
1"""
2Collection-specific RAG Search Engine
4Provides semantic search within a specific document collection using RAG.
5"""
7import os
8from typing import List, Dict, Any, Optional
10from ...security.secure_logging import logger
11from .search_engine_library import LibraryRAGSearchEngine
12from ...constants import SNIPPET_LENGTH_LONG
13from ...research_library.services.library_rag_service import LibraryRAGService
14from ...database.models.library import RAGIndex, Document
15from ...research_library.services.pdf_storage_manager import PDFStorageManager
16from ...database.session_context import get_user_db_session
17from ...config.thread_settings import get_setting_from_snapshot
18from ...config.paths import get_library_directory
19from ...utilities.type_utils import to_bool
22class CollectionSearchEngine(LibraryRAGSearchEngine):
23 """
24 Search engine for a specific document collection using RAG.
25 Directly searches only the specified collection's FAISS index.
26 Each collection uses its own embedding model that was used during indexing.
27 """
29 # Mark as local RAG engine
30 is_local = True
32 def __init__(
33 self,
34 collection_id: str,
35 collection_name: str,
36 llm: Optional[Any] = None,
37 max_filtered_results: Optional[int] = None,
38 max_results: int = 10,
39 settings_snapshot: Optional[Dict[str, Any]] = None,
40 **kwargs,
41 ):
42 """
43 Initialize the collection-specific search engine.
45 Args:
46 collection_id: UUID of the collection to search within
47 collection_name: Name of the collection for display
48 llm: Language model for relevance filtering
49 max_filtered_results: Maximum number of results to keep after filtering
50 max_results: Maximum number of search results
51 settings_snapshot: Settings snapshot from thread context
52 **kwargs: Additional engine-specific parameters
53 """
54 super().__init__(
55 llm=llm,
56 max_filtered_results=max_filtered_results,
57 max_results=max_results,
58 settings_snapshot=settings_snapshot,
59 **kwargs,
60 )
61 self.collection_id = collection_id
62 self.collection_name = collection_name
63 self.collection_key = f"collection_{collection_id}"
64 # The policy registry name for this engine is known at construction
65 # (the factory only stamps registry engines), so set it here —
66 # direct instantiations (e.g. the library search route) get the
67 # runtime egress backstop without caller cooperation. NB: between
68 # super().__init__ (which stamps "library") and this line the
69 # instance briefly carries the parent's name — nothing may call
70 # _verify_egress_scope() in that window.
71 self._engine_name = self.collection_key
73 # Load collection-specific embedding settings
74 self._load_collection_embedding_settings()
76 def _load_collection_embedding_settings(self):
77 """
78 Load embedding settings from the collection's RAG index.
79 Uses the same embedding model that was used during indexing.
80 """
81 if not self.username:
82 logger.warning("Cannot load collection settings without username")
83 return
85 try:
86 with get_user_db_session(self.username) as db_session:
87 # Get RAG index for this collection
88 rag_index = (
89 db_session.query(RAGIndex)
90 .filter_by(
91 collection_name=self.collection_key,
92 is_current=True,
93 )
94 .first()
95 )
97 if not rag_index:
98 logger.warning(
99 f"No RAG index found for collection {self.collection_id}"
100 )
101 return
103 # Use embedding settings from the RAG index
104 self.embedding_model = rag_index.embedding_model
105 self.embedding_provider = rag_index.embedding_model_type.value
106 self.chunk_size = rag_index.chunk_size or self.chunk_size
107 self.chunk_overlap = (
108 rag_index.chunk_overlap or self.chunk_overlap
109 )
111 logger.info(
112 f"Collection '{self.collection_name}' using embedding: "
113 f"{self.embedding_provider}/{self.embedding_model}"
114 )
116 except Exception as e:
117 safe_msg = self._scrub_error(e)
118 logger.exception(
119 f"Error loading collection {self.collection_id} settings ({type(e).__name__}): {safe_msg}"
120 )
122 def search(
123 self,
124 query: str,
125 limit: int = 10,
126 llm_callback=None,
127 extra_params: Optional[Dict[str, Any]] = None,
128 ) -> List[Dict[str, Any]]:
129 """
130 Search within the specific collection using semantic search.
132 Directly searches only this collection's FAISS index instead of
133 searching all collections and filtering.
135 Args:
136 query: Search query
137 limit: Maximum number of results to return
138 llm_callback: Optional LLM callback for processing results
139 extra_params: Additional search parameters
141 Returns:
142 List of search results from this collection
143 """
144 if not self.username:
145 logger.error("Cannot search collection without username")
146 return []
148 # search() does not go through BaseSearchEngine.run(), so apply the
149 # runtime egress backstop here as well. Raises PolicyDeniedError on
150 # denial — deliberately BEFORE the broad try/except below.
151 self._verify_egress_scope()
153 try:
154 # Get RAG index info for this collection
155 with get_user_db_session(self.username) as db_session:
156 rag_index = (
157 db_session.query(RAGIndex)
158 .filter_by(
159 collection_name=self.collection_key,
160 is_current=True,
161 )
162 .first()
163 )
165 if not rag_index:
166 logger.info(
167 f"No RAG index for collection '{self.collection_name}'"
168 )
169 return []
171 # Get embedding settings from RAG index
172 embedding_model = rag_index.embedding_model
173 embedding_provider = rag_index.embedding_model_type.value
174 chunk_size = rag_index.chunk_size or self.chunk_size
175 chunk_overlap = rag_index.chunk_overlap or self.chunk_overlap
176 # Thread the stored normalization flag through so the index is
177 # queried with the same L2 normalization it was BUILT with.
178 # Otherwise LibraryRAGService defaults normalize_vectors=True,
179 # corrupting similarity rankings for collections indexed with
180 # local_search_normalize_vectors=False. NULL → True mirrors the
181 # to_bool(..., default=True) convention used elsewhere.
182 normalize_vectors = to_bool(
183 rag_index.normalize_vectors, default=True
184 )
185 # Thread the stored distance metric through too: it is what
186 # LibraryRAGService stamps onto each SearchResult.metric, which
187 # in turn selects the relevance formula below. Defaulting to
188 # "cosine" (the LibraryRAGService default) would mislabel an
189 # l2/dot_product collection and apply the wrong [0,1] mapping.
190 # NULL (legacy rows) → "cosine" preserves prior behavior.
191 distance_metric = rag_index.distance_metric or "cosine"
192 # Thread the stored index_type too (like metric/normalize): with
193 # it left at the "flat" default, _get_vector_index would fall
194 # back to flat for a legacy NULL-index_type row that is physically
195 # HNSW, mislabelling the store (wrong supports_delete/relevance,
196 # and now a load()-time config-mismatch rejection). NULL → "flat".
197 index_type = rag_index.index_type or "flat"
199 # Create RAG service with collection's embedding settings
200 with LibraryRAGService(
201 username=self.username,
202 embedding_model=embedding_model,
203 embedding_provider=embedding_provider,
204 chunk_size=chunk_size,
205 chunk_overlap=chunk_overlap,
206 normalize_vectors=normalize_vectors,
207 distance_metric=distance_metric,
208 index_type=index_type,
209 ) as rag_service:
210 # Check if there are indexed documents
211 stats = rag_service.get_rag_stats(self.collection_id)
212 if stats.get("indexed_documents", 0) == 0:
213 logger.info(
214 f"No documents indexed in collection '{self.collection_name}'"
215 )
216 return []
218 # Search this collection's vector index via the new
219 # int-id-keyed store (chunk text is rehydrated from the
220 # encrypted DB by id — never read from the vector store
221 # itself; see the SECURITY INVARIANT in
222 # vector_stores/base.py).
223 search_results = rag_service.search(
224 query, self.collection_id, limit
225 )
227 if not search_results:
228 logger.info(
229 f"No results found in collection '{self.collection_name}'"
230 )
231 return []
233 # Convert to search result format
234 results = []
235 for r in search_results:
236 metadata = dict(r.metadata or {})
238 # Get document ID
239 doc_id = (
240 r.source_id
241 or metadata.get("source_id")
242 or metadata.get("document_id")
243 )
245 # Get title
246 title = (
247 r.document_title
248 or metadata.get("document_title")
249 or metadata.get("title")
250 or (f"Document {doc_id}" if doc_id else "Untitled")
251 )
253 # Create snippet from content
254 snippet = (
255 r.text[:SNIPPET_LENGTH_LONG] + "..."
256 if len(r.text) > SNIPPET_LENGTH_LONG
257 else r.text
258 )
260 # Generate document URL
261 document_url = self._get_document_url(doc_id)
263 # Add collection info to metadata
264 metadata["collection_id"] = self.collection_id
265 metadata["collection_name"] = self.collection_name
267 # r.distance's meaning depends on r.metric. Map both to a
268 # [0, 1] relevance that grows with similarity:
269 # - cosine/dot_product (inner product on normalized
270 # vectors): a similarity in [-1, 1] where higher is
271 # nearer, so (d+1)/2 clamped to [0, 1]. Using it raw would
272 # yield negative relevance for anti-correlated hits and
273 # break downstream [0,1] consumers (similarity %, filters).
274 # - anything else (l2, or a non-standard metric string): a
275 # (squared) distance >= 0 where LOWER is nearer, so
276 # 1/(1+d) in (0, 1].
277 # This IP test MUST match faiss_store._build_base_index,
278 # which builds an inner-product index iff the metric is
279 # cosine/dot_product and an L2 index otherwise. A bare
280 # `metric == "l2"` here would score a non-standard metric
281 # (which builds an L2 index) with the IP formula, inverting
282 # the ranking.
283 relevance = (
284 max(0.0, min(1.0, (r.distance + 1.0) / 2.0))
285 if r.metric in ("cosine", "dot_product")
286 else 1.0 / (1.0 + r.distance)
287 )
289 result = {
290 "title": title,
291 "snippet": snippet,
292 "url": document_url,
293 "link": document_url,
294 "source": "library",
295 "source_type": "library",
296 "relevance_score": float(relevance),
297 "metadata": metadata,
298 }
299 results.append(result)
301 logger.info(
302 f"Collection '{self.collection_name}' search returned "
303 f"{len(results)} results for query: {query[:50]}..."
304 )
306 return results
308 except Exception as e:
309 # Re-raise instead of returning [] so a failed search is not
310 # indistinguishable from "no matching documents": run() records
311 # the failure in metrics, and API callers can report the error.
312 safe_msg = self._scrub_error(e)
313 logger.exception(
314 f"Error searching collection '{self.collection_name}' ({type(e).__name__}): {safe_msg}"
315 )
316 raise
318 def _get_document_url(self, doc_id: Optional[str]) -> str:
319 """Get the URL for viewing a document."""
320 if not doc_id:
321 return "#"
323 # Default to root document page (shows all options: PDF, Text, Chunks, etc.)
324 document_url = f"/library/document/{doc_id}"
326 try:
327 with get_user_db_session(self.username) as session:
328 document = session.query(Document).filter_by(id=doc_id).first()
329 if document:
330 from pathlib import Path
332 library_root = get_setting_from_snapshot(
333 "research_library.storage_path",
334 default=str(get_library_directory()),
335 settings_snapshot=self.settings_snapshot,
336 )
337 library_root = (
338 Path(os.path.expandvars(library_root))
339 .expanduser()
340 .resolve()
341 )
342 if PDFStorageManager.pdf_exists( 342 ↛ 349line 342 didn't jump to line 349
343 library_root, document, session
344 ):
345 document_url = f"/library/document/{doc_id}/pdf"
346 except Exception:
347 logger.warning(f"Error getting document URL for {doc_id}")
349 return document_url