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

local_deep_research/web_search_engines/engines/search_engine_google_pse.py

@@ -1,38 +1,41 @@
-from typing import Dict, List, Any, Optional
+import logging
 import os
-import requests
-import time
 import random
-import logging
-from requests.exceptions import RequestException
-from urllib.parse import quote_plus
+import time
+from typing import Any, Dict, List, Optional
+
+import requests
 from langchain_core.language_models import BaseLLM
+from requests.exceptions import RequestException
 
-from local_deep_research.web_search_engines.search_engine_base import BaseSearchEngine
+from ..search_engine_base import BaseSearchEngine
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+
 class GooglePSESearchEngine(BaseSearchEngine):
     """Google Programmable Search Engine implementation"""
 
-    def __init__(self,
-                max_results: int = 10,
-                region: str = "us",
-                safe_search: bool = True,
-                search_language: str = "English",
-                api_key: Optional[str] = None,
-                search_engine_id: Optional[str] = None,
-                llm: Optional[BaseLLM] = None,
-                include_full_content: bool = False,
-                max_filtered_results: Optional[int] = None,
-                max_retries: int = 3,
-                retry_delay: float = 2.0,
-                **kwargs):
+    def __init__(
+        self,
+        max_results: int = 10,
+        region: str = "us",
+        safe_search: bool = True,
+        search_language: str = "English",
+        api_key: Optional[str] = None,
+        search_engine_id: Optional[str] = None,
+        llm: Optional[BaseLLM] = None,
+        include_full_content: bool = False,
+        max_filtered_results: Optional[int] = None,
+        max_retries: int = 3,
+        retry_delay: float = 2.0,
+        **kwargs,
+    ):
         """
         Initialize the Google Programmable Search Engine.
-        
+
         Args:
             max_results: Maximum number of search results
             region: Region code for search results
@@ -48,17 +51,19 @@ class GooglePSESearchEngine(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_full_content = include_full_content
-        
+
         # Retry configuration
         self.max_retries = max_retries
         self.retry_delay = retry_delay
-        
+
         # Rate limiting - keep track of last request time
         self.last_request_time = 0
         self.min_request_interval = 0.5  # Minimum time between requests in seconds
-        
+
         # Language code mapping
         language_code_mapping = {
             "english": "en",
@@ -70,82 +75,86 @@ class GooglePSESearchEngine(BaseSearchEngine):
             "korean": "ko",
             "portuguese": "pt",
             "russian": "ru",
-            "chinese": "zh-CN"
+            "chinese": "zh-CN",
         }
-        
+
         # Get language code
         search_language = search_language.lower()
         self.language = language_code_mapping.get(search_language, "en")
-        
+
         # Safe search setting
         self.safe = "active" if safe_search else "off"
-        
+
         # Region/Country setting
         self.region = region
-        
+
         # API key and Search Engine ID
         self.api_key = api_key or os.getenv("GOOGLE_PSE_API_KEY")
         self.search_engine_id = search_engine_id or os.getenv("GOOGLE_PSE_ENGINE_ID")
-        
+
         if not self.api_key:
-            raise ValueError("Google API key is required. Set it in the GOOGLE_PSE_API_KEY environment variable.")
+            raise ValueError(
+                "Google API key is required. Set it in the GOOGLE_PSE_API_KEY environment variable."
+            )
         if not self.search_engine_id:
-            raise ValueError("Google Search Engine ID is required. Set it in the GOOGLE_PSE_ENGINE_ID environment variable.")
-            
+            raise ValueError(
+                "Google Search Engine ID is required. Set it in the GOOGLE_PSE_ENGINE_ID environment variable."
+            )
+
         # Validate connection and credentials
         self._validate_connection()
-    
+
     def _validate_connection(self):
         """Test the connection to ensure API key and Search Engine ID are valid"""
         try:
             # Make a minimal test query
             response = self._make_request("test")
-            
+
             # Check if we got a valid response
             if response.get("error"):
                 error_msg = response["error"].get("message", "Unknown error")
                 raise ValueError(f"Google PSE API error: {error_msg}")
-                
+
             # If we get here, the connection is valid
             logger.info("Google PSE connection validated successfully")
             return True
-            
+
         except Exception as e:
             # Log the error and re-raise
             logger.error(f"Error validating Google PSE connection: {str(e)}")
             raise
-    
+
     def _respect_rate_limit(self):
         """Ensure we don't exceed rate limits by adding appropriate delay between requests"""
         current_time = time.time()
         elapsed = current_time - self.last_request_time
-        
+
         # If we've made a request recently, wait until the minimum interval has passed
         if elapsed < self.min_request_interval:
             sleep_time = self.min_request_interval - elapsed
-            logger.debug(f"Rate limiting: sleeping for {sleep_time:.2f}s")
+            logger.debug("Rate limiting: sleeping for %.2f s", sleep_time)
             time.sleep(sleep_time)
-        
+
         # Update the last request time
         self.last_request_time = time.time()
 
     def _make_request(self, query: str, start_index: int = 1) -> Dict:
         """
         Make a request to the Google PSE API with retry logic and rate limiting
-        
+
         Args:
             query: Search query string
             start_index: Starting index for pagination
-            
+
         Returns:
             JSON response from the API
-            
+
         Raises:
             RequestException: If all retry attempts fail
         """
         # Base URL for the API
         url = "https://www.googleapis.com/customsearch/v1"
-        
+
         # Parameters for the request
         params = {
             "key": self.api_key,
@@ -155,127 +164,151 @@ class GooglePSESearchEngine(BaseSearchEngine):
             "start": start_index,
             "safe": self.safe,
             "lr": f"lang_{self.language}",
-            "gl": self.region
+            "gl": self.region,
         }
-        
+
         # Implement retry logic with exponential backoff
         attempt = 0
         last_exception = None
-        
+
         while attempt < self.max_retries:
             try:
                 # Respect rate limits
                 self._respect_rate_limit()
-                
+
                 # Add jitter to retries after the first attempt
                 if attempt > 0:
                     jitter = random.uniform(0.5, 1.5)
                     sleep_time = self.retry_delay * (2 ** (attempt - 1)) * jitter
-                    logger.info(f"Retry attempt {attempt+1}/{self.max_retries} for query '{query}'. Waiting {sleep_time:.2f}s")
+                    logger.info(
+                        "Retry attempt %s / %s for query '%s'. Waiting %s s",
+                        attempt + 1,
+                        self.max_retries,
+                        query,
+                        f"{sleep_time:.2f}",
+                    )
                     time.sleep(sleep_time)
-                
+
                 # Make the request
-                logger.debug(f"Making request to Google PSE API: {query} (start_index={start_index})")
+                logger.debug(
+                    "Making request to Google PSE API: %s (start_index=%s)",
+                    query,
+                    start_index,
+                )
                 response = requests.get(url, params=params, timeout=10)
-                
+
                 # Check for HTTP errors
                 response.raise_for_status()
-                
+
                 # Return the JSON response
                 return response.json()
-                
+
             except RequestException as e:
-                logger.warning(f"Request error on attempt {attempt+1}/{self.max_retries}: {str(e)}")
+                logger.warning(
+                    "Request error on attempt %s / %s: %s",
+                    attempt + 1,
+                    self.max_retries,
+                    str(e),
+                )
                 last_exception = e
             except Exception as e:
-                logger.warning(f"Error on attempt {attempt+1}/{self.max_retries}: {str(e)}")
+                logger.warning(
+                    "Error on attempt %s / %s: %s",
+                    attempt + 1,
+                    self.max_retries,
+                    str(e),
+                )
                 last_exception = e
-            
+
             attempt += 1
-        
+
         # If we get here, all retries failed
         error_msg = f"Failed to get response from Google PSE API after {self.max_retries} attempts"
         logger.error(error_msg)
-        
+
         if last_exception:
             raise RequestException(f"{error_msg}: {str(last_exception)}")
         else:
             raise RequestException(error_msg)
-        
+
     def _get_previews(self, query: str) -> List[Dict[str, Any]]:
         """Get search result previews/snippets"""
         results = []
-        
+
         # Google PSE API returns a maximum of 10 results per request
         # We may need to make multiple requests to get the desired number
         start_index = 1
         total_results = 0
-        
+
         while total_results < self.max_results:
             try:
                 response = self._make_request(query, start_index)
-                
+
                 # Break if no items
                 if "items" not in response:
                     break
-                    
+
                 items = response.get("items", [])
-                
+
                 # Process each result
                 for item in items:
                     title = item.get("title", "")
                     snippet = item.get("snippet", "")
                     url = item.get("link", "")
-                    
+
                     # Skip results without URL
                     if not url:
                         continue
-                        
-                    results.append({
-                        "title": title,
-                        "snippet": snippet,
-                        "url": url,
-                        "source": "Google Programmable Search"
-                    })
-                    
+
+                    results.append(
+                        {
+                            "title": title,
+                            "snippet": snippet,
+                            "url": url,
+                            "source": "Google Programmable Search",
+                        }
+                    )
+
                     total_results += 1
                     if total_results >= self.max_results:
                         break
-                        
+
                 # Check if there are more results
                 if not items or total_results >= self.max_results:
                     break
-                    
+
                 # Update start index for next request
                 start_index += len(items)
-                
+
                 # Add a small delay between multiple requests to be respectful of the API
                 if total_results < self.max_results:
                     time.sleep(self.min_request_interval)
-                
+
             except Exception as e:
-                logger.error(f"Error getting search results: {str(e)}")
+                logger.error("Error getting search results: %s", str(e))
                 break
-                
-        logger.info(f"Retrieved {len(results)} search results for query: '{query}'")
+
+        logger.info("Retrieved %s search results for query: '%s'", len(results), query)
         return results
-        
-    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 search results"""
         # Use the BaseSearchEngine implementation
         return super()._get_full_content(relevant_items)
-        
+
     def run(self, query: str) -> List[Dict[str, Any]]:
         """Run the search engine to get results for a query"""
         # Get search result previews/snippets
         search_results = self._get_previews(query)
-        
+
         # Filter for relevance if we have an LLM and max_filtered_results
         if self.llm and self.max_filtered_results:
             search_results = self._filter_for_relevance(query, search_results)
-            
+
         # Get full content if needed
         if self.include_full_content:
             search_results = self._get_full_content(search_results)
-            
-        return search_results 
+
+        return search_results