local-deep-research 0.1.0__py3-none-any.whl

local_deep_research/web_search_engines/engines/search_engine_serpapi.py

@@ -0,0 +1,230 @@
+from langchain_community.utilities import SerpAPIWrapper
+from typing import Dict, List, Any, Optional
+import os
+from langchain_core.language_models import BaseLLM
+
+from local_deep_research.web_search_engines.search_engine_base import BaseSearchEngine
+from local_deep_research import config
+
+
+class SerpAPISearchEngine(BaseSearchEngine):
+    """Google search engine implementation using SerpAPI with two-phase approach"""
+
+    def __init__(self,
+                max_results: int = 10,
+                region: str = "us",
+                time_period: str = "y",
+                safe_search: bool = True,
+                search_language: str = "English",
+                api_key: Optional[str] = None,
+                language_code_mapping: Optional[Dict[str, str]] = None,
+                llm: Optional[BaseLLM] = None,
+                include_full_content: bool = False,
+                max_filtered_results: Optional[int] = None,
+                **kwargs):
+        """
+        Initialize the SerpAPI search engine.
+        
+        Args:
+            max_results: Maximum number of search results
+            region: Region code for search results
+            time_period: Time period for search results
+            safe_search: Whether to enable safe search
+            search_language: Language for search results
+            api_key: SerpAPI API key (can also be set in SERP_API_KEY env)
+            language_code_mapping: Mapping from language names to codes
+            llm: Language model for relevance filtering
+            include_full_content: Whether to include full webpage content in results
+            max_filtered_results: Maximum number of results to keep after filtering
+            **kwargs: Additional parameters (ignored but accepted for compatibility)
+        """
+        # Initialize the BaseSearchEngine with the LLM and max_filtered_results
+        super().__init__(llm=llm, max_filtered_results=max_filtered_results)
+        
+        self.max_results = max_results
+        self.include_full_content = include_full_content
+        
+        # Set up language code mapping
+        if language_code_mapping is None:
+            language_code_mapping = {
+                "english": "en",
+                "spanish": "es",
+                "chinese": "zh",
+                "hindi": "hi",
+                "french": "fr",
+                "arabic": "ar",
+                "bengali": "bn",
+                "portuguese": "pt",
+                "russian": "ru",
+            }
+        
+        # Get API key
+        serpapi_api_key = api_key or os.getenv("SERP_API_KEY")
+        if not serpapi_api_key:
+            raise ValueError("SERP_API_KEY not found. Please provide api_key or set the SERP_API_KEY environment variable.")
+        
+        # Get language code
+        language_code = language_code_mapping.get(search_language.lower(), "en")
+        
+        # Initialize SerpAPI wrapper
+        self.engine = SerpAPIWrapper(
+            serpapi_api_key=serpapi_api_key,
+            params={
+                "engine": "google",
+                "hl": language_code,
+                "gl": region,
+                "safe": "active" if safe_search else "off",
+                "tbs": f"qdr:{time_period}",
+                "num": max_results,
+            }
+        )
+        
+        # If full content is requested, initialize FullSearchResults
+        if include_full_content:
+            # Import FullSearchResults only if needed
+            try:
+                from local_deep_research.web_search_engines.engines.full_search  import FullSearchResults
+                self.full_search = FullSearchResults(
+                    llm=llm,
+                    web_search=self.engine,
+                    language=search_language,
+                    max_results=max_results,
+                    region=region,
+                    time=time_period,
+                    safesearch="Moderate" if safe_search else "Off"
+                )
+            except ImportError:
+                print("Warning: FullSearchResults not available. Full content retrieval disabled.")
+                self.include_full_content = False
+    
+    def _get_previews(self, query: str) -> List[Dict[str, Any]]:
+        """
+        Get preview information from SerpAPI.
+        
+        Args:
+            query: The search query
+            
+        Returns:
+            List of preview dictionaries
+        """
+        print("Getting search results from SerpAPI")
+        
+        try:
+            # Get search results from SerpAPI
+            organic_results = self.engine.results(query).get("organic_results", [])
+            
+            # Format results as previews
+            previews = []
+            for result in organic_results:
+                preview = {
+                    "id": result.get("position", len(previews)),  # Use position as ID
+                    "title": result.get("title", ""),
+                    "link": result.get("link", ""),
+                    "snippet": result.get("snippet", ""),
+                    "displayed_link": result.get("displayed_link", ""),
+                    "position": result.get("position")
+                }
+                
+                # Store full SerpAPI result for later
+                preview["_full_result"] = result
+                
+                previews.append(preview)
+            
+            # Store the previews for potential full content retrieval
+            self._search_results = previews
+            
+            return previews
+            
+        except Exception as e:
+            print(f"Error getting SerpAPI results: {e}")
+            return []
+    
+    def _get_full_content(self, relevant_items: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """
+        Get full content for the relevant search results.
+        If include_full_content is True and FullSearchResults is available,
+        retrieves full webpage content for the results.
+        
+        Args:
+            relevant_items: List of relevant preview dictionaries
+            
+        Returns:
+            List of result dictionaries with full content if requested
+        """
+        # Check if we should get full content
+        if hasattr(config, 'SEARCH_SNIPPETS_ONLY') and config.SEARCH_SNIPPETS_ONLY:
+            print("Snippet-only mode, skipping full content retrieval")
+            
+            # Return the relevant items with their full SerpAPI information
+            results = []
+            for item in relevant_items:
+                # Use the full result if available, otherwise use the preview
+                if "_full_result" in item:
+                    result = item["_full_result"]
+                    # Remove temporary field
+                    if "_full_result" in result:
+                        del result["_full_result"]
+                else:
+                    result = item
+                    
+                results.append(result)
+                
+            return results
+            
+        # If full content retrieval is enabled
+        if self.include_full_content and hasattr(self, 'full_search'):
+            print("Retrieving full webpage content")
+            
+            try:
+                # Extract only the links from relevant items
+                links = [item.get("link") for item in relevant_items if item.get("link")]
+                
+                # Use FullSearchResults to get full content
+                # This is a simplified approach - in a real implementation,
+                # you would need to fetch and process the URLs
+                results_with_content = self.full_search._get_full_content(relevant_items)
+                
+                return results_with_content
+                
+            except Exception as e:
+                print(f"Error retrieving full content: {e}")
+                # Fall back to returning the items without full content
+        
+        # Return items with their full SerpAPI information
+        results = []
+        for item in relevant_items:
+            # Use the full result if available, otherwise use the preview
+            if "_full_result" in item:
+                result = item["_full_result"].copy()
+                # Remove temporary field
+                if "_full_result" in result:
+                    del result["_full_result"]
+            else:
+                result = item.copy()
+                if "_full_result" in result:
+                    del result["_full_result"]
+                
+            results.append(result)
+            
+        return results
+    
+    def run(self, query: str) -> List[Dict[str, Any]]:
+        """
+        Execute a search using SerpAPI with the two-phase approach.
+        
+        Args:
+            query: The search query
+            
+        Returns:
+            List of search results
+        """
+        print("---Execute a search using SerpAPI (Google)---")
+        
+        # Use the implementation from the parent class which handles all phases
+        results = super().run(query)
+        
+        # Clean up
+        if hasattr(self, '_search_results'):
+            del self._search_results
+            
+        return results