local-deep-research 0.1.26__py3-none-any.whl → 0.2.0__py3-none-any.whl

local_deep_research/web_search_engines/engines/search_engine_wikipedia.py

@@ -1,10 +1,11 @@
+import logging
+from typing import Any, Dict, List, Optional
+
 import wikipedia
-from typing import Dict, List, Any, Optional
 from langchain_core.language_models import BaseLLM
-import logging
 
-from local_deep_research.web_search_engines.search_engine_base import BaseSearchEngine
-from local_deep_research import config
+from ...config import search_config
+from ..search_engine_base import BaseSearchEngine
 
 # Setup logging
 logger = logging.getLogger(__name__)
@@ -12,18 +13,20 @@ logger = logging.getLogger(__name__)
 
 class WikipediaSearchEngine(BaseSearchEngine):
     """Wikipedia search engine implementation with two-phase approach"""
-    
-    def __init__(self, 
-                max_results: int = 10, 
-                language: str = "en", 
-                include_content: bool = True,
-                sentences: int = 5,
-                llm: Optional[BaseLLM] = None,
-                max_filtered_results: Optional[int] = None,
-                **kwargs):
+
+    def __init__(
+        self,
+        max_results: int = 10,
+        language: str = "en",
+        include_content: bool = True,
+        sentences: int = 5,
+        llm: Optional[BaseLLM] = None,
+        max_filtered_results: Optional[int] = None,
+        **kwargs,
+    ):
         """
         Initialize the Wikipedia search engine.
-        
+
         Args:
             max_results: Maximum number of search results
             language: Language code for Wikipedia (e.g., 'en', 'fr', 'es')
@@ -34,38 +37,42 @@ class WikipediaSearchEngine(BaseSearchEngine):
             **kwargs: Additional parameters (ignored but accepted for compatibility)
         """
         # Initialize the BaseSearchEngine with LLM, max_filtered_results, and max_results
-        super().__init__(llm=llm, max_filtered_results=max_filtered_results, max_results=max_results)
+        super().__init__(
+            llm=llm, max_filtered_results=max_filtered_results, max_results=max_results
+        )
         self.include_content = include_content
         self.sentences = sentences
-        
+
         # Set the Wikipedia language
         wikipedia.set_lang(language)
-    
+
     def _get_previews(self, query: str) -> List[Dict[str, Any]]:
         """
         Get preview information (titles and summaries) for Wikipedia pages.
-        
+
         Args:
             query: The search query
-            
+
         Returns:
             List of preview dictionaries
         """
-        logger.info(f"Getting Wikipedia page previews for query: {query}")
-        
+        logger.info("Getting Wikipedia page previews for query: %s", query)
+
         try:
             # Get search results (just titles)
             search_results = wikipedia.search(query, results=self.max_results)
-            
-            logger.info(f"Found {len(search_results)} Wikipedia results: {search_results}")
-            
+
+            logger.info(
+                f"Found {len(search_results)} Wikipedia results: {search_results}"
+            )
+
             if not search_results:
-                logger.info(f"No Wikipedia results found for query: {query}")
+                logger.info("No Wikipedia results found for query: %s", query)
                 return []
-            
+
             # Create a cache for full pages (will be populated on-demand)
             self._page_cache = {}
-            
+
             # Generate previews with summaries
             previews = []
             for title in search_results:
@@ -73,83 +80,104 @@ class WikipediaSearchEngine(BaseSearchEngine):
                     # Get just the summary, with auto_suggest=False to be more precise
                     summary = None
                     try:
-                        summary = wikipedia.summary(title, sentences=self.sentences, auto_suggest=False)
+                        summary = wikipedia.summary(
+                            title, sentences=self.sentences, auto_suggest=False
+                        )
                     except wikipedia.exceptions.DisambiguationError as e:
                         # If disambiguation error, try the first option
                         if e.options and len(e.options) > 0:
-                            logger.info(f"Disambiguation for '{title}', trying first option: {e.options[0]}")
+                            logger.info(
+                                f"Disambiguation for '{title}', trying first option: {e.options[0]}"
+                            )
                             try:
-                                summary = wikipedia.summary(e.options[0], sentences=self.sentences, auto_suggest=False)
+                                summary = wikipedia.summary(
+                                    e.options[0],
+                                    sentences=self.sentences,
+                                    auto_suggest=False,
+                                )
                                 title = e.options[0]  # Use the new title
                             except Exception as inner_e:
-                                logger.error(f"Error with disambiguation option: {inner_e}")
+                                logger.error(
+                                    f"Error with disambiguation option: {inner_e}"
+                                )
                                 continue
                         else:
-                            logger.warning(f"Disambiguation with no options for '{title}'")
+                            logger.warning(
+                                f"Disambiguation with no options for '{title}'"
+                            )
                             continue
-                    
+
                     if summary:
                         preview = {
                             "id": title,  # Use title as ID
                             "title": title,
                             "snippet": summary,
-                            "link": f"https://en.wikipedia.org/wiki/{title.replace(' ', '_')}"
+                            "link": f"https://en.wikipedia.org/wiki/{title.replace(' ', '_')}",
                         }
-                        
+
                         previews.append(preview)
-                    
-                except (wikipedia.exceptions.PageError,
-                       wikipedia.exceptions.WikipediaException) as e:
+
+                except (
+                    wikipedia.exceptions.PageError,
+                    wikipedia.exceptions.WikipediaException,
+                ) as e:
                     # Skip pages with errors
                     logger.warning(f"Error getting summary for '{title}': {e}")
                     continue
                 except Exception as e:
                     logger.error(f"Unexpected error for '{title}': {e}")
                     continue
-            
+
             logger.info(f"Successfully created {len(previews)} previews from Wikipedia")
             return previews
-            
+
         except Exception as e:
             logger.error(f"Error getting Wikipedia previews: {e}")
             return []
-    
-    def _get_full_content(self, relevant_items: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+
+    def _get_full_content(
+        self, relevant_items: List[Dict[str, Any]]
+    ) -> List[Dict[str, Any]]:
         """
         Get full content for the relevant Wikipedia pages.
-        
+
         Args:
             relevant_items: List of relevant preview dictionaries
-            
+
         Returns:
             List of result dictionaries with full content
         """
         # Check if we should add full content
-        if hasattr(config, 'SEARCH_SNIPPETS_ONLY') and config.SEARCH_SNIPPETS_ONLY:
+        if (
+            hasattr(search_config, "SEARCH_SNIPPETS_ONLY")
+            and search_config.SEARCH_SNIPPETS_ONLY
+        ):
             logger.info("Snippet-only mode, skipping full content retrieval")
             return relevant_items
-            
-        logger.info(f"Getting full content for {len(relevant_items)} relevant Wikipedia pages")
-        
+
+        logger.info(
+            f"Getting full content for {len(relevant_items)} relevant Wikipedia pages"
+        )
+
         results = []
         for item in relevant_items:
             title = item.get("id")  # Title stored as ID
-            
+
             if not title:
                 results.append(item)
                 continue
-                
+
             try:
                 # Get the full page
                 page = wikipedia.page(title, auto_suggest=False)
-                
+
                 # Create a full result with all information
                 result = {
                     "title": page.title,
                     "link": page.url,
-                    "snippet": item.get("snippet", "")  # Keep existing snippet
+                    "snippet": item.get("snippet", ""),  # Keep existing snippet
                 }
-                
+
                 # Add additional information
                 result["content"] = page.content
                 result["full_content"] = page.content
@@ -158,29 +186,33 @@ class WikipediaSearchEngine(BaseSearchEngine):
                 result["links"] = page.links
                 result["images"] = page.images
                 result["sections"] = page.sections
-                
+
                 results.append(result)
-                
-            except (wikipedia.exceptions.DisambiguationError, 
-                   wikipedia.exceptions.PageError,
-                   wikipedia.exceptions.WikipediaException) as e:
+
+            except (
+                wikipedia.exceptions.DisambiguationError,
+                wikipedia.exceptions.PageError,
+                wikipedia.exceptions.WikipediaException,
+            ) as e:
                 # If error, use the preview
                 logger.warning(f"Error getting full content for '{title}': {e}")
                 results.append(item)
             except Exception as e:
-                logger.error(f"Unexpected error getting full content for '{title}': {e}")
+                logger.error(
+                    f"Unexpected error getting full content for '{title}': {e}"
+                )
                 results.append(item)
-        
+
         return results
-    
+
     def get_summary(self, title: str, sentences: Optional[int] = None) -> str:
         """
         Get a summary of a specific Wikipedia page.
-        
+
         Args:
             title: Title of the Wikipedia page
             sentences: Number of sentences to include (defaults to self.sentences)
-            
+
         Returns:
             Summary of the page
         """
@@ -189,31 +221,37 @@ class WikipediaSearchEngine(BaseSearchEngine):
             return wikipedia.summary(title, sentences=sentences, auto_suggest=False)
         except wikipedia.exceptions.DisambiguationError as e:
             if e.options and len(e.options) > 0:
-                return wikipedia.summary(e.options[0], sentences=sentences, auto_suggest=False)
+                return wikipedia.summary(
+                    e.options[0], sentences=sentences, auto_suggest=False
+                )
             raise
-    
+
     def get_page(self, title: str) -> Dict[str, Any]:
         """
         Get detailed information about a specific Wikipedia page.
-        
+
         Args:
             title: Title of the Wikipedia page
-            
+
         Returns:
             Dictionary with page information
         """
-        # Check if we should include full content
-        include_content = not (hasattr(config, 'SEARCH_SNIPPETS_ONLY') and config.SEARCH_SNIPPETS_ONLY)
-        
+        # Initialize include_content with our instance value
+        include_content = self.include_content
+
+        # Check if we should override with config setting
+        if hasattr(search_config, "SEARCH_SNIPPETS_ONLY"):
+            include_content = not search_config.SEARCH_SNIPPETS_ONLY
+
         try:
             page = wikipedia.page(title, auto_suggest=False)
-            
+
             result = {
                 "title": page.title,
                 "link": page.url,
-                "snippet": self.get_summary(title, self.sentences)
+                "snippet": self.get_summary(title, self.sentences),
             }
-            
+
             # Add additional information if requested
             if include_content:
                 result["content"] = page.content
@@ -223,18 +261,18 @@ class WikipediaSearchEngine(BaseSearchEngine):
                 result["links"] = page.links
                 result["images"] = page.images
                 result["sections"] = page.sections
-            
+
             return result
         except wikipedia.exceptions.DisambiguationError as e:
             if e.options and len(e.options) > 0:
                 return self.get_page(e.options[0])
             raise
-    
+
     def set_language(self, language: str) -> None:
         """
         Change the Wikipedia language.
-        
+
         Args:
             language: Language code (e.g., 'en', 'fr', 'es')
         """
-        wikipedia.set_lang(language)
+        wikipedia.set_lang(language)