Coverage for src/local_deep_research/web_search_engines/engines/search_engine_retriever.py: 100%
60 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-22 21:20 +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"
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 # Retriever documents may be chunks of the source URL, not the
112 # complete resource that other cache consumers request.
113 "full_content": doc.page_content,
114 "full_content_is_complete": False,
115 "author": metadata.get("author", ""),
116 "date": metadata.get("date", ""),
117 # Include all metadata for flexibility
118 "metadata": metadata,
119 # Score if available
120 "score": metadata.get("score", 1.0),
121 # Source information
122 "source": self.name,
123 "retriever_type": self.retriever.__class__.__name__,
124 }
126 def _get_previews(self, query: str) -> List[Dict[str, Any]]:
127 """
128 Get preview information from the retriever.
130 Args:
131 query: Search query
133 Returns:
134 List of preview dictionaries
135 """
136 try:
137 # Use the retriever to get relevant documents
138 docs = self.retriever.invoke(query)
140 # Convert to preview format
141 previews = []
142 for i, doc in enumerate(docs[: self.max_results]):
143 preview = self._convert_document_to_result(doc, i)
144 previews.append(preview)
146 logger.info(
147 f"Retriever '{self.name}' returned {len(previews)} previews"
148 )
149 return previews
151 except Exception as e:
152 safe_msg = self._scrub_error(e)
153 logger.exception(
154 f"Error getting previews from retriever ({type(e).__name__}): {safe_msg}"
155 )
156 return []
158 def _get_full_content(
159 self, relevant_items: List[Dict[str, Any]]
160 ) -> List[Dict[str, Any]]:
161 """
162 For retrievers, previews already contain full content.
164 Args:
165 relevant_items: List of relevant preview dictionaries
167 Returns:
168 Same list with full content (already included)
169 """
170 # For retrievers, the preview already contains the full content
171 # Just ensure the 'full_content' field is present
172 for item in relevant_items:
173 if "full_content" not in item and "snippet" in item:
174 item["full_content"] = item["snippet"]
175 return relevant_items
177 async def arun(self, query: str) -> List[Dict[str, Any]]:
178 """
179 Async version of search using the retriever.
181 Args:
182 query: Search query
184 Returns:
185 List of search results in LDR format
186 """
187 try:
188 # Use async retriever if available
189 if hasattr(self.retriever, "aget_relevant_documents"):
190 docs = await self.retriever.aget_relevant_documents(query)
191 else:
192 # Fall back to sync version
193 logger.debug(
194 f"Retriever '{self.name}' doesn't support async, using sync version"
195 )
196 return self.run(query)
198 # Convert documents to results
199 results = []
200 for i, doc in enumerate(docs[: self.max_results]):
201 result = self._convert_document_to_result(doc, i)
202 results.append(result)
204 logger.info(
205 f"Retriever '{self.name}' returned {len(results)} async results"
206 )
207 return results
209 except Exception as e:
210 safe_msg = self._scrub_error(e)
211 logger.exception(
212 f"Error in async retriever search ({type(e).__name__}): {safe_msg}"
213 )
214 return []