Coverage for src/local_deep_research/web_search_engines/engines/search_engine_wikipedia.py: 100%
131 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-20 01:24 +0000
1import json
2from typing import Any, Dict, List, Optional
4import requests
5import wikipedia
6from langchain_core.language_models import BaseLLM
8from ..search_engine_base import BaseSearchEngine, Exposure, Sensitivity
9from ...security import sanitize_for_log
10from ...security.secure_logging import logger
13# Wikipedia / MediaWiki returns a non-JSON HTML body with HTTP 429 when a
14# client is rate-limited. The `wikipedia` PyPI library calls
15# ``Response.json()`` unconditionally, so the symptom we observe is a
16# ``JSONDecodeError`` on every subsequent ``wikipedia.summary()`` call.
17# Catch that family explicitly so we can short-circuit instead of
18# emitting a full traceback per title (the original behaviour spammed
19# the log with 14+ stack traces for a single rate-limited query).
20_TRANSIENT_DECODE_ERRORS: tuple = (
21 json.JSONDecodeError,
22 requests.exceptions.JSONDecodeError,
23)
26class WikipediaSearchEngine(BaseSearchEngine):
27 """Wikipedia search engine implementation with two-phase approach"""
29 # Mark as public search engine
30 is_public = True
31 egress_sensitivity = Sensitivity.NON_SENSITIVE
32 egress_exposure = Exposure.EXPOSING
33 is_lexical = True
34 needs_llm_relevance_filter = True
36 def __init__(
37 self,
38 max_results: int = 10,
39 language: str = "en",
40 include_content: bool = True,
41 sentences: int = 5,
42 llm: Optional[BaseLLM] = None,
43 max_filtered_results: Optional[int] = None,
44 settings_snapshot: Optional[Dict[str, Any]] = None,
45 **kwargs,
46 ):
47 """
48 Initialize the Wikipedia search engine.
50 Args:
51 max_results: Maximum number of search results
52 language: Language code for Wikipedia (e.g., 'en', 'fr', 'es')
53 include_content: Whether to include full page content in results
54 sentences: Number of sentences to include in summary
55 llm: Language model for relevance filtering
56 max_filtered_results: Maximum number of results to keep after filtering
57 **kwargs: Additional parameters (ignored but accepted for compatibility)
58 """
59 # Initialize the BaseSearchEngine with LLM, max_filtered_results, and max_results
60 super().__init__(
61 llm=llm,
62 max_filtered_results=max_filtered_results,
63 max_results=max_results,
64 settings_snapshot=settings_snapshot,
65 )
66 self.include_content = include_content
67 self.sentences = sentences
69 # Set the Wikipedia language
70 wikipedia.set_lang(language)
72 def _get_previews(self, query: str) -> List[Dict[str, Any]]:
73 """
74 Get preview information (titles and summaries) for Wikipedia pages.
76 Args:
77 query: The search query
79 Returns:
80 List of preview dictionaries
81 """
82 logger.info(f"Getting Wikipedia page previews for query: {query}")
84 try:
85 # Apply rate limiting before search request
86 self._last_wait_time = self.rate_tracker.apply_rate_limit(
87 self.engine_type
88 )
90 # Get search results (just titles)
91 search_results = wikipedia.search(query, results=self.max_results)
92 safe_search_results = [
93 sanitize_for_log(title) for title in search_results
94 ]
96 logger.info(
97 f"Found {len(search_results)} Wikipedia results: "
98 f"{safe_search_results}"
99 )
101 if not search_results:
102 logger.info(f"No Wikipedia results found for query: {query}")
103 return []
105 # Generate previews with summaries.
106 # NOTE: This loop is intentionally sequential. Do NOT parallelize with
107 # ThreadPoolExecutor because:
108 # 1. The `wikipedia` PyPI library is not thread-safe — it uses global
109 # mutable state (API_URL, RATE_LIMIT_LAST_CALL) and an unlocked cache.
110 # Concurrent threads would corrupt the library's built-in rate limiting.
111 # 2. self._last_wait_time is a shared instance attribute with no lock —
112 # concurrent writes would feed incorrect data to record_outcome().
113 # 3. Downstream _filter_for_relevance uses positional indices — random
114 # completion order would cause the LLM to select wrong articles.
115 previews = []
116 for title in search_results:
117 safe_title = sanitize_for_log(title)
118 try:
119 # Get just the summary, with auto_suggest=False to be more precise
120 summary = None
121 try:
122 # Apply rate limiting before summary request
123 self._last_wait_time = (
124 self.rate_tracker.apply_rate_limit(self.engine_type)
125 )
127 summary = wikipedia.summary(
128 title, sentences=self.sentences, auto_suggest=False
129 )
130 except wikipedia.exceptions.DisambiguationError as e:
131 # If disambiguation error, try the first option
132 if e.options and len(e.options) > 0:
133 # Page titles are external content; sanitize
134 # before interpolation to prevent log injection.
135 first_option = e.options[0]
136 safe_first_option = sanitize_for_log(first_option)
137 logger.info(
138 f"Disambiguation for '{safe_title}', trying "
139 f"first option: {safe_first_option}"
140 )
141 try:
142 summary = wikipedia.summary(
143 first_option,
144 sentences=self.sentences,
145 auto_suggest=False,
146 )
147 title = first_option # Use the new title
148 except Exception as inner_e:
149 safe_msg = self._scrub_error(inner_e)
150 logger.exception(
151 f"Error with disambiguation option ({type(inner_e).__name__}): {safe_msg}"
152 )
153 continue
154 else:
155 logger.warning(
156 f"Disambiguation with no options for '{safe_title}'"
157 )
158 continue
160 if summary:
161 preview = {
162 "id": title, # Use title as ID
163 "title": title,
164 "snippet": summary,
165 "link": f"https://en.wikipedia.org/wiki/{title.replace(' ', '_')}",
166 "source": "Wikipedia",
167 }
169 previews.append(preview)
171 except (
172 wikipedia.exceptions.PageError,
173 wikipedia.exceptions.WikipediaException,
174 ):
175 # Skip pages with errors
176 logger.warning(f"Error getting summary for '{safe_title}'")
177 continue
178 except _TRANSIENT_DECODE_ERRORS:
179 # MediaWiki almost certainly returned a 429 (or other
180 # non-JSON page) — every remaining title in this batch
181 # will hit the same throttle. Bail out with one warning
182 # instead of a per-title traceback storm.
183 logger.warning(
184 "Wikipedia rate-limited (non-JSON response while "
185 "fetching '{}'); returning {} previews collected so far",
186 safe_title,
187 len(previews),
188 )
189 break
190 except Exception as e:
191 safe_msg = self._scrub_error(e)
192 logger.exception(
193 f"Unexpected error for '{safe_title}' "
194 f"({type(e).__name__}): {safe_msg}"
195 )
196 continue
198 logger.info(
199 f"Successfully created {len(previews)} previews from Wikipedia"
200 )
201 return previews
203 except _TRANSIENT_DECODE_ERRORS:
204 # Same 429-style failure on the outer wikipedia.search() call —
205 # log once at warning level and return an empty list.
206 logger.warning(
207 "Wikipedia rate-limited on search for query '{}'; "
208 "returning no previews",
209 query,
210 )
211 return []
212 except Exception as e:
213 safe_msg = self._scrub_error(e)
214 logger.exception(
215 f"Error getting Wikipedia previews ({type(e).__name__}): {safe_msg}"
216 )
217 return []
219 def _get_full_content(
220 self, relevant_items: List[Dict[str, Any]]
221 ) -> List[Dict[str, Any]]:
222 """
223 Get full content for the relevant Wikipedia pages.
225 Args:
226 relevant_items: List of relevant preview dictionaries
228 Returns:
229 List of result dictionaries with full content
230 """
231 logger.info(
232 f"Getting full content for {len(relevant_items)} relevant Wikipedia pages"
233 )
235 results = []
236 for item in relevant_items:
237 title = item.get("id") # Title stored as ID
239 if not title:
240 results.append(item)
241 continue
243 safe_title = sanitize_for_log(str(title))
244 try:
245 # Apply rate limiting before page request
246 self._last_wait_time = self.rate_tracker.apply_rate_limit(
247 self.engine_type
248 )
250 # Get the full page
251 page = wikipedia.page(title, auto_suggest=False)
253 # Create a full result with all information
254 result = {
255 "title": page.title,
256 "link": page.url,
257 "snippet": item.get("snippet", ""), # Keep existing snippet
258 "source": "Wikipedia",
259 }
261 # Add additional information
262 result["content"] = page.content
263 result["full_content"] = page.content
264 result["categories"] = page.categories
265 result["references"] = page.references
266 result["links"] = page.links
267 result["images"] = page.images
268 result["sections"] = page.sections
270 results.append(result)
272 except (
273 wikipedia.exceptions.DisambiguationError,
274 wikipedia.exceptions.PageError,
275 wikipedia.exceptions.WikipediaException,
276 ):
277 # If error, use the preview
278 logger.warning(f"Error getting full content for '{safe_title}'")
279 results.append(item)
280 except Exception as e:
281 safe_msg = self._scrub_error(e)
282 logger.exception(
283 f"Unexpected error getting full content for "
284 f"'{safe_title}' ({type(e).__name__}): {safe_msg}"
285 )
286 results.append(item)
288 return results
290 def get_summary(self, title: str, sentences: Optional[int] = None) -> str:
291 """
292 Get a summary of a specific Wikipedia page.
294 Args:
295 title: Title of the Wikipedia page
296 sentences: Number of sentences to include (defaults to self.sentences)
298 Returns:
299 Summary of the page
300 """
301 sentences = sentences or self.sentences
302 try:
303 return str(
304 wikipedia.summary(
305 title, sentences=sentences, auto_suggest=False
306 )
307 )
308 except wikipedia.exceptions.DisambiguationError as e:
309 if e.options and len(e.options) > 0:
310 return str(
311 wikipedia.summary(
312 e.options[0], sentences=sentences, auto_suggest=False
313 )
314 )
315 raise
317 def get_page(self, title: str) -> Dict[str, Any]:
318 """
319 Get detailed information about a specific Wikipedia page.
321 Args:
322 title: Title of the Wikipedia page
324 Returns:
325 Dictionary with page information
326 """
327 include_content = self.include_content
329 try:
330 page = wikipedia.page(title, auto_suggest=False)
332 result = {
333 "title": page.title,
334 "link": page.url,
335 "snippet": self.get_summary(title, self.sentences),
336 "source": "Wikipedia",
337 }
339 # Add additional information if requested
340 if include_content:
341 result["content"] = page.content
342 result["full_content"] = page.content
343 result["categories"] = page.categories
344 result["references"] = page.references
345 result["links"] = page.links
346 result["images"] = page.images
347 result["sections"] = page.sections
349 return result
350 except wikipedia.exceptions.DisambiguationError as e:
351 if e.options and len(e.options) > 0:
352 return self.get_page(e.options[0])
353 raise
355 def set_language(self, language: str) -> None:
356 """
357 Change the Wikipedia language.
359 Args:
360 language: Language code (e.g., 'en', 'fr', 'es')
361 """
362 wikipedia.set_lang(language)