Coverage for src/local_deep_research/web_search_engines/engines/search_engine_library.py: 95%
151 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-20 01:24 +0000
1"""
2Library RAG Search Engine
4Provides semantic search over the user's personal research library using RAG.
5"""
7import os
8from typing import List, Dict, Any, Optional
10from ..search_engine_base import BaseSearchEngine, Exposure, Sensitivity
11from ...security.secure_logging import logger
12from ...constants import SNIPPET_LENGTH_LONG
13from ...research_library.services.library_rag_service import LibraryRAGService
14from ...research_library.services.library_service import LibraryService
15from ...config.thread_settings import get_setting_from_snapshot
16from ...utilities.llm_utils import get_server_url
17from ...utilities.type_utils import to_bool
18from ...database.models.library import RAGIndex, Document
19from ...research_library.services.pdf_storage_manager import PDFStorageManager
20from ...database.session_context import get_user_db_session
21from ...config.paths import get_library_directory
24class LibraryRAGSearchEngine(BaseSearchEngine):
25 """
26 Search engine that queries the user's research library using RAG/semantic search.
27 """
29 # Mark as local RAG engine
30 is_local = True
32 # Egress (ADR-0007): a local document store — sensitive data, kept on-box.
33 egress_sensitivity = Sensitivity.SENSITIVE
34 egress_exposure = Exposure.CONTAINED
36 def __init__(
37 self,
38 llm: Optional[Any] = None,
39 max_filtered_results: Optional[int] = None,
40 max_results: int = 10,
41 settings_snapshot: Optional[Dict[str, Any]] = None,
42 **kwargs,
43 ):
44 """
45 Initialize the Library RAG search engine.
47 Args:
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.username = (
62 settings_snapshot.get("_username") if settings_snapshot else None
63 )
64 # Self-stamp the policy registry name so direct instantiations get
65 # the runtime egress backstop without caller cooperation (the
66 # factory stamps registry engines, but routes/services may build
67 # this engine directly). CollectionSearchEngine overrides this
68 # with its collection key after calling super().__init__.
69 self._engine_name = "library"
71 if not self.username:
72 logger.warning(
73 "Library RAG search engine initialized without username"
74 )
76 # Get RAG configuration from settings. NB: keyword arguments — the
77 # signature is (key, default, username, settings_snapshot); these
78 # calls used to pass (key, settings_snapshot, default) positionally,
79 # so the snapshot was never consulted and, outside a thread settings
80 # context, each setting silently resolved to the snapshot DICT
81 # itself (the misplaced "default").
82 self.embedding_model = get_setting_from_snapshot(
83 "local_search_embedding_model",
84 default="all-MiniLM-L6-v2",
85 settings_snapshot=settings_snapshot,
86 )
87 self.embedding_provider = get_setting_from_snapshot(
88 "local_search_embedding_provider",
89 default="sentence_transformers",
90 settings_snapshot=settings_snapshot,
91 )
92 self.chunk_size = get_setting_from_snapshot(
93 "local_search_chunk_size",
94 default=1000,
95 settings_snapshot=settings_snapshot,
96 )
97 self.chunk_overlap = get_setting_from_snapshot(
98 "local_search_chunk_overlap",
99 default=200,
100 settings_snapshot=settings_snapshot,
101 )
103 # Extract server URL from settings snapshot for link generation
104 self.server_url = get_server_url(settings_snapshot)
106 # Local RAG retrieval is vector-similarity ranked and can return many
107 # near-duplicate chunks per query — enough to flood an agent's context
108 # across research iterations. Apply a RAG-specific cap (default 20),
109 # independent of the global ``search.max_results`` (which targets web
110 # engines). This OVERRIDES the inherited max_results for local search.
111 # Default to the inherited max_results so an ABSENT setting preserves
112 # the caller's value (programmatic callers, tests); real web/CLI runs
113 # carry the shipped default of 20 via default_settings.json, so they ARE
114 # capped. A malformed value also falls back to no change.
115 rag_max_results = get_setting_from_snapshot(
116 "search.rag.max_results",
117 default=self.max_results,
118 settings_snapshot=settings_snapshot,
119 )
120 try:
121 rag_max_results = int(rag_max_results)
122 except (TypeError, ValueError):
123 rag_max_results = self.max_results
124 if rag_max_results > 0: 124 ↛ 136line 124 didn't jump to line 136 because the condition on line 124 was always true
125 self.max_results = rag_max_results
127 # Embeddings already rank by cosine similarity, so the LLM relevance
128 # filter is OPT-IN here (unlike keyword engines that hardcode
129 # needs_llm_relevance_filter=True). Default on: prune
130 # semantically-near-but-irrelevant chunks. The factory only ever SETS
131 # enable_llm_relevance_filter=True (for should_filter engines) and never
132 # unsets it — and RAG engines are not should_filter (no
133 # needs_llm_relevance_filter; dynamic ``collection_<uuid>`` name) — so
134 # this assignment is authoritative. The base filter is a no-op without
135 # an llm, so it degrades gracefully.
136 self.enable_llm_relevance_filter = to_bool(
137 get_setting_from_snapshot(
138 "search.rag.enable_relevance_filter",
139 default=True,
140 settings_snapshot=settings_snapshot,
141 ),
142 default=True,
143 )
145 def search(
146 self,
147 query: str,
148 limit: int = 10,
149 llm_callback=None,
150 extra_params: Optional[Dict[str, Any]] = None,
151 ) -> List[Dict[str, Any]]:
152 """
153 Search the library using semantic search.
155 Args:
156 query: Search query
157 limit: Maximum number of results to return
158 llm_callback: Optional LLM callback for processing results
159 extra_params: Additional search parameters
161 Returns:
162 List of search results with title, url, snippet, etc.
163 """
164 if not self.username:
165 logger.error("Cannot search library without username")
166 return []
168 # search() can be called directly (bypassing run()), so apply the
169 # runtime egress backstop here too. Memoized per snapshot, so the
170 # run() -> _get_previews() -> search() path re-checks for free.
171 # Raises PolicyDeniedError on denial — deliberately BEFORE the
172 # broad try/except below.
173 self._verify_egress_scope()
175 try:
176 # Initialize services
177 library_service = LibraryService(username=self.username)
179 # Get all collections for this user
180 collections = library_service.get_all_collections()
181 if not collections:
182 logger.info("No collections found for user")
183 return []
185 # Search across all collections and merge results. Each
186 # (SearchResult, relevance, metadata) tuple is pre-scored per
187 # its OWN collection's metric before merging, so a collection
188 # indexed with cosine/dot_product doesn't get compared on the
189 # wrong scale against one indexed with l2 (see the
190 # relevance-formula comment below).
191 all_results = []
192 failed_collections = []
193 for collection in collections:
194 collection_id = collection.get("id")
195 if not collection_id:
196 continue
198 try:
199 # Get the RAG index for this collection to find embedding settings
200 with get_user_db_session(self.username) as session:
201 collection_name = f"collection_{collection_id}"
202 rag_index = (
203 session.query(RAGIndex)
204 .filter_by(
205 collection_name=collection_name,
206 is_current=True,
207 )
208 .first()
209 )
211 if not rag_index:
212 logger.debug(
213 f"No RAG index found for collection {collection_id}"
214 )
215 continue
217 # Get embedding settings from the RAG index
218 embedding_model = rag_index.embedding_model
219 embedding_provider = (
220 rag_index.embedding_model_type.value
221 )
222 chunk_size = rag_index.chunk_size or self.chunk_size
223 chunk_overlap = (
224 rag_index.chunk_overlap or self.chunk_overlap
225 )
226 # Thread the stored normalization flag AND distance
227 # metric through, exactly like CollectionSearchEngine.
228 # Without normalize_vectors, LibraryRAGService defaults
229 # True and would L2-normalize the query against a
230 # raw-vector collection, flipping the top hit. Without
231 # distance_metric, every collection is labeled "cosine"
232 # and l2/dot_product hits get the wrong [0,1] mapping in
233 # the cross-collection merge below. NULL → prior default.
234 normalize_vectors = to_bool(
235 rag_index.normalize_vectors, default=True
236 )
237 distance_metric = rag_index.distance_metric or "cosine"
238 # Thread the stored index_type too (like metric/
239 # normalize): a legacy NULL-index_type row that is
240 # physically HNSW would otherwise fall back to "flat",
241 # mislabelling the store. NULL → "flat".
242 index_type = rag_index.index_type or "flat"
244 # Create RAG service with the collection's embedding settings
245 with LibraryRAGService(
246 username=self.username,
247 embedding_model=embedding_model,
248 embedding_provider=embedding_provider,
249 chunk_size=chunk_size,
250 chunk_overlap=chunk_overlap,
251 normalize_vectors=normalize_vectors,
252 distance_metric=distance_metric,
253 index_type=index_type,
254 ) as rag_service:
255 # Get RAG stats to check if there are any indexed documents
256 stats = rag_service.get_rag_stats(collection_id)
257 if stats.get("indexed_documents", 0) == 0:
258 logger.debug(
259 f"No documents indexed in collection {collection_id}"
260 )
261 continue
263 # Search this collection's vector index via the new
264 # int-id-keyed store (chunk text is rehydrated from
265 # the encrypted DB by id — never read from the
266 # vector store itself; see the SECURITY INVARIANT
267 # in vector_stores/base.py).
268 search_results = rag_service.search(
269 query, collection_id, limit
270 )
272 # Score per THIS collection's own metric, add
273 # collection info to metadata, and append to the
274 # cross-collection merge list.
275 for r in search_results:
276 metadata = dict(r.metadata or {})
277 metadata["collection_id"] = collection_id
278 metadata["collection_name"] = collection.get(
279 "name", "Unknown"
280 )
281 # Map r.distance to a [0,1] relevance growing with
282 # similarity: cosine/dot_product is an inner product
283 # in [-1,1] (higher=nearer) -> (d+1)/2 clamped; l2 (or
284 # any non-standard metric) is a distance >=0 (lower=
285 # nearer) -> 1/(1+d) in (0,1]. Raw IP would give
286 # negative relevance and break [0,1] consumers.
287 # Scoring here (per-collection, before the merge)
288 # keeps different-metric collections comparable once
289 # merged. The IP test MUST match faiss_store.
290 # _build_base_index (inner-product iff cosine/
291 # dot_product, else L2): a bare `metric == "l2"` would
292 # score a non-standard metric — which builds an L2
293 # index — with the IP formula and invert the ranking.
294 relevance = (
295 max(0.0, min(1.0, (r.distance + 1.0) / 2.0))
296 if r.metric in ("cosine", "dot_product")
297 else 1.0 / (1.0 + r.distance)
298 )
299 all_results.append((r, relevance, metadata))
301 except Exception as e:
302 # One broken collection must not abort the others, but
303 # record the failure so it is not silently equated with
304 # "no matching documents" below.
305 safe_msg = self._scrub_error(e)
306 logger.exception(
307 f"Error searching collection {collection_id} ({type(e).__name__}): {safe_msg}"
308 )
309 failed_collections.append(
310 collection.get("name") or str(collection_id)
311 )
312 continue
314 # Sort all results by relevance, descending (higher is always
315 # better post-transform, regardless of the source collection's
316 # metric — see the per-result relevance comment above).
317 all_results.sort(key=lambda item: item[1], reverse=True)
319 # Take top results across all collections
320 top_results = all_results[:limit]
322 if not top_results:
323 if failed_collections:
324 # Zero results AND at least one collection errored:
325 # reporting "no results" here would hide the failure.
326 self._raise_collections_failed(failed_collections)
327 logger.info("No results found across any collections")
328 return []
330 if failed_collections:
331 # Partial success: results from the healthy collections are
332 # returned, but tell the user they are incomplete. WARNING
333 # level streams to the research-log UI via
334 # frontend_progress_sink.
335 logger.warning(
336 f"Library search results are incomplete: "
337 f"{len(failed_collections)} of {len(collections)} "
338 f"collection(s) failed: {failed_collections}"
339 )
341 # Convert to the search-result dict format
342 results = []
343 for r, relevance, metadata in top_results:
344 # Try both source_id and document_id for compatibility
345 doc_id = (
346 r.source_id
347 or metadata.get("source_id")
348 or metadata.get("document_id")
349 )
351 # Get title from the result, with fallbacks
352 title = (
353 r.document_title
354 or metadata.get("document_title")
355 or metadata.get("title")
356 or (f"Document {doc_id}" if doc_id else "Untitled")
357 )
359 # Content is rehydrated from the encrypted DB by chunk id
360 snippet = (
361 r.text[:SNIPPET_LENGTH_LONG] + "..."
362 if len(r.text) > SNIPPET_LENGTH_LONG
363 else r.text
364 )
366 # Generate URL to document content
367 # Default to root document page (shows all options: PDF, Text, Chunks, etc.)
368 document_url = f"/library/document/{doc_id}" if doc_id else "#"
370 if doc_id:
371 try:
372 with get_user_db_session(self.username) as session:
373 document = (
374 session.query(Document)
375 .filter_by(id=doc_id)
376 .first()
377 )
378 if document: 378 ↛ 400line 378 didn't jump to line 400
379 from pathlib import Path
381 library_root = get_setting_from_snapshot(
382 "research_library.storage_path",
383 default=str(get_library_directory()),
384 settings_snapshot=self.settings_snapshot,
385 )
386 library_root = (
387 Path(os.path.expandvars(library_root))
388 .expanduser()
389 .resolve()
390 )
391 if PDFStorageManager.pdf_exists(
392 library_root, document, session
393 ):
394 document_url = (
395 f"/library/document/{doc_id}/pdf"
396 )
397 except Exception:
398 logger.warning(f"Error querying document {doc_id}")
400 result = {
401 "title": title,
402 "snippet": snippet,
403 "url": document_url,
404 "link": document_url, # Add "link" for source extraction
405 "source": "library",
406 "source_type": "library",
407 "relevance_score": float(relevance),
408 "metadata": metadata,
409 }
411 results.append(result)
413 logger.info(
414 f"Library RAG search returned {len(results)} results for query: {query}"
415 )
416 return results
418 except Exception as e:
419 # Re-raise so run() records the failure in metrics instead of
420 # treating a failed search as "no matching documents".
421 safe_msg = self._scrub_error(e)
422 logger.exception(
423 f"Error searching library RAG ({type(e).__name__}): {safe_msg}"
424 )
425 raise
427 @staticmethod
428 def _raise_collections_failed(failed_collections: List[str]) -> None:
429 """Signal zero results caused (at least partly) by search failures."""
430 raise RuntimeError(
431 f"Library search returned no results and failed for "
432 f"{len(failed_collections)} collection(s): {failed_collections}"
433 )
435 def _get_previews(
436 self,
437 query: str,
438 limit: Optional[int] = None,
439 llm_callback=None,
440 extra_params: Optional[Dict[str, Any]] = None,
441 ) -> List[Dict[str, Any]]:
442 """
443 Get preview results for the query.
444 Delegates to the search method, using the configured max_results
445 when no explicit limit is provided.
446 """
447 if limit is None:
448 limit = self.max_results
449 return self.search(query, limit, llm_callback, extra_params)
451 def _get_full_content(
452 self, relevant_items: List[Dict[str, Any]]
453 ) -> List[Dict[str, Any]]:
454 """
455 Get full content for relevant library documents.
456 Retrieves complete document text instead of just snippets.
457 """
458 if not self.username:
459 logger.error("Cannot retrieve full content without username")
460 return relevant_items
462 try:
463 from ...database.models.library import Document
464 from ...database.session_context import get_user_db_session
466 # Retrieve full content for each document
467 for item in relevant_items:
468 doc_id = item.get("metadata", {}).get("document_id")
469 if not doc_id:
470 continue
472 # Get full document text from database
473 with get_user_db_session(self.username) as db_session:
474 document = (
475 db_session.query(Document).filter_by(id=doc_id).first()
476 )
478 if document and document.text_content:
479 # Replace snippet with full content
480 item["content"] = document.text_content
481 item["snippet"] = (
482 document.text_content[:SNIPPET_LENGTH_LONG] + "..."
483 if len(document.text_content) > SNIPPET_LENGTH_LONG
484 else document.text_content
485 )
486 logger.debug(
487 f"Retrieved full content for document {doc_id}"
488 )
490 return relevant_items
492 except Exception as e:
493 safe_msg = self._scrub_error(e)
494 logger.exception(
495 f"Error retrieving full content from library ({type(e).__name__}): {safe_msg}"
496 )
497 return relevant_items
499 def close(self):
500 """Clean up resources."""
501 pass