Coverage for src/local_deep_research/web_search_engines/engines/search_engine_retriever.py: 100%
60 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-20 01:24 +0000
1"""
2Search engine implementation that wraps any LangChain retriever.
3This allows using vector stores, databases, or any custom retriever as a search source in LDR.
4"""
6from typing import Any, Dict, List, Optional
8from langchain_core.retrievers import BaseRetriever, Document
10from ...constants import SNIPPET_LENGTH_LONG
11from ...security.secure_logging import logger
12from ..search_engine_base import BaseSearchEngine
15class RetrieverSearchEngine(BaseSearchEngine):
16 """
17 Search engine that uses any LangChain retriever.
19 This allows users to plug in any LangChain retriever (vector stores,
20 databases, custom implementations) and use it as a search engine in LDR.
21 """
23 def __init__(
24 self,
25 retriever: BaseRetriever,
26 max_results: int = 10,
27 name: str | None = None,
28 settings_snapshot: Optional[Dict[str, Any]] = None,
29 **kwargs,
30 ):
31 """
32 Initialize the retriever-based search engine.
34 Args:
35 retriever: Any LangChain BaseRetriever instance
36 max_results: Maximum number of results to return
37 name: Display name for this retriever (defaults to retriever class name)
38 **kwargs: Additional parameters passed to parent
39 """
40 super().__init__(
41 max_results=max_results,
42 settings_snapshot=settings_snapshot,
43 **kwargs,
44 )
45 self.retriever = retriever
46 self.name = name if name is not None else retriever.__class__.__name__
48 def run(
49 self, query: str, research_context: Dict[str, Any] | None = None
50 ) -> List[Dict[str, Any]]:
51 """
52 Execute search using the LangChain retriever.
54 Args:
55 query: Search query
56 research_context: Context from previous research to use.
58 Returns:
59 List of search results in LDR format
60 """
61 try:
62 # Use the retriever to get relevant documents
63 docs = self.retriever.invoke(query)
65 # Convert LangChain documents to LDR search result format
66 results = []
67 for i, doc in enumerate(docs[: self.max_results]):
68 result = self._convert_document_to_result(doc, i)
69 results.append(result)
71 logger.info(
72 f"Retriever '{self.name}' returned {len(results)} results for query: {query}"
73 )
74 return results
76 except Exception as e:
77 safe_msg = self._scrub_error(e)
78 logger.exception(
79 f"Error in retriever search ({type(e).__name__}): {safe_msg}"
80 )
81 return []
83 def _convert_document_to_result(
84 self, doc: Document, index: int
85 ) -> Dict[str, Any]:
86 """
87 Convert a LangChain Document to LDR search result format.
89 Args:
90 doc: LangChain Document
91 index: Result index
93 Returns:
94 Search result in LDR format
95 """
96 # Extract metadata
97 metadata = doc.metadata or {}
99 # Build the result
100 return {
101 # Required fields for LDR
102 "title": metadata.get("title", f"Document {index + 1}"),
103 "url": metadata.get(
104 "source",
105 metadata.get("url", f"retriever://{self.name}/doc_{index}"),
106 ),
107 "snippet": doc.page_content[:SNIPPET_LENGTH_LONG]
108 if doc.page_content
109 else "",
110 # Optional fields
111 "full_content": doc.page_content,
112 "author": metadata.get("author", ""),
113 "date": metadata.get("date", ""),
114 # Include all metadata for flexibility
115 "metadata": metadata,
116 # Score if available
117 "score": metadata.get("score", 1.0),
118 # Source information
119 "source": self.name,
120 "retriever_type": self.retriever.__class__.__name__,
121 }
123 def _get_previews(self, query: str) -> List[Dict[str, Any]]:
124 """
125 Get preview information from the retriever.
127 Args:
128 query: Search query
130 Returns:
131 List of preview dictionaries
132 """
133 try:
134 # Use the retriever to get relevant documents
135 docs = self.retriever.invoke(query)
137 # Convert to preview format
138 previews = []
139 for i, doc in enumerate(docs[: self.max_results]):
140 preview = self._convert_document_to_result(doc, i)
141 previews.append(preview)
143 logger.info(
144 f"Retriever '{self.name}' returned {len(previews)} previews for query: {query}"
145 )
146 return previews
148 except Exception as e:
149 safe_msg = self._scrub_error(e)
150 logger.exception(
151 f"Error getting previews from retriever ({type(e).__name__}): {safe_msg}"
152 )
153 return []
155 def _get_full_content(
156 self, relevant_items: List[Dict[str, Any]]
157 ) -> List[Dict[str, Any]]:
158 """
159 For retrievers, previews already contain full content.
161 Args:
162 relevant_items: List of relevant preview dictionaries
164 Returns:
165 Same list with full content (already included)
166 """
167 # For retrievers, the preview already contains the full content
168 # Just ensure the 'full_content' field is present
169 for item in relevant_items:
170 if "full_content" not in item and "snippet" in item:
171 item["full_content"] = item["snippet"]
172 return relevant_items
174 async def arun(self, query: str) -> List[Dict[str, Any]]:
175 """
176 Async version of search using the retriever.
178 Args:
179 query: Search query
181 Returns:
182 List of search results in LDR format
183 """
184 try:
185 # Use async retriever if available
186 if hasattr(self.retriever, "aget_relevant_documents"):
187 docs = await self.retriever.aget_relevant_documents(query)
188 else:
189 # Fall back to sync version
190 logger.debug(
191 f"Retriever '{self.name}' doesn't support async, using sync version"
192 )
193 return self.run(query)
195 # Convert documents to results
196 results = []
197 for i, doc in enumerate(docs[: self.max_results]):
198 result = self._convert_document_to_result(doc, i)
199 results.append(result)
201 logger.info(
202 f"Retriever '{self.name}' returned {len(results)} async results for query: {query}"
203 )
204 return results
206 except Exception as e:
207 safe_msg = self._scrub_error(e)
208 logger.exception(
209 f"Error in async retriever search ({type(e).__name__}): {safe_msg}"
210 )
211 return []