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

local_deep_research/web_search_engines/engines/search_engine_local_all.py

@@ -0,0 +1,153 @@
+"""
+Search engine that searches across all local collections
+"""
+
+import logging
+from typing import Dict, List, Any, Optional
+from langchain_core.language_models import BaseLLM
+
+from local_deep_research.web_search_engines.search_engine_base import BaseSearchEngine
+from local_deep_research.web_search_engines.search_engine_factory import create_search_engine
+from local_deep_research import config
+
+# Setup logging
+logger = logging.getLogger(__name__)
+
+class LocalAllSearchEngine(BaseSearchEngine):
+    """
+    Search engine that searches across all local document collections.
+    Acts as a meta search engine specifically for local collections.
+    """
+    
+    def __init__(
+        self,
+        llm: Optional[BaseLLM] = None,
+        max_results: int = 10,
+        max_filtered_results: Optional[int] = None,
+        **kwargs
+    ):
+        """
+        Initialize the local all-collections search engine.
+        
+        Args:
+            llm: Language model for relevance filtering
+            max_results: Maximum number of search results
+            max_filtered_results: Maximum results after filtering
+            **kwargs: Additional parameters passed to LocalSearchEngine instances
+        """
+        # Initialize the base search engine
+        super().__init__(llm=llm, max_filtered_results=max_filtered_results)
+        
+        self.max_results = max_results
+        
+        # Find all local collection search engines
+        self.local_engines = {}
+        try:
+            from local_collections import LOCAL_COLLECTIONS
+            
+            for collection_id, collection in LOCAL_COLLECTIONS.items():
+                if not collection.get("enabled", True):
+                    continue
+                    
+                # Create a search engine for this collection
+                try:
+                    engine = create_search_engine(
+                        collection_id,
+                        llm=llm,
+                        max_filtered_results=max_filtered_results
+                    )
+                    
+                    if engine:
+                        self.local_engines[collection_id] = {
+                            "engine": engine,
+                            "name": collection.get("name", collection_id),
+                            "description": collection.get("description", "")
+                        }
+                except Exception as e:
+                    logger.error(f"Error creating search engine for collection '{collection_id}': {e}")
+        except ImportError:
+            logger.warning("No local collections configuration found")
+    
+    def _get_previews(self, query: str) -> List[Dict[str, Any]]:
+        """
+        Get preview information for documents from all local collections.
+        
+        Args:
+            query: The search query
+            
+        Returns:
+            List of preview dictionaries
+        """
+        logger.info(f"Searching across all local collections for query: {query}")
+        
+        all_previews = []
+        
+        # Get previews from each local search engine
+        for collection_id, engine_info in self.local_engines.items():
+            engine = engine_info["engine"]
+            try:
+                # Get previews from this engine
+                previews = engine._get_previews(query)
+                
+                # Add collection info to each preview
+                for preview in previews:
+                    preview["collection_id"] = collection_id
+                    preview["collection_name"] = engine_info["name"]
+                    preview["collection_description"] = engine_info["description"]
+                
+                all_previews.extend(previews)
+            except Exception as e:
+                logger.error(f"Error searching collection '{collection_id}': {e}")
+        
+        if not all_previews:
+            logger.info(f"No local documents found for query: {query}")
+            return []
+        
+        # Sort by similarity score if available
+        all_previews.sort(
+            key=lambda x: float(x.get("similarity", 0)), 
+            reverse=True
+        )
+        
+        # Limit to max_results
+        return all_previews[:self.max_results]
+    
+    def _get_full_content(self, relevant_items: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """
+        Get full content for the relevant documents.
+        Delegates to the appropriate collection's search engine.
+        
+        Args:
+            relevant_items: List of relevant preview dictionaries
+            
+        Returns:
+            List of result dictionaries with full content
+        """
+        # Group items by collection
+        items_by_collection = {}
+        for item in relevant_items:
+            collection_id = item.get("collection_id")
+            if collection_id and collection_id in self.local_engines:
+                if collection_id not in items_by_collection:
+                    items_by_collection[collection_id] = []
+                items_by_collection[collection_id].append(item)
+        
+        # Process each collection's items with its own engine
+        all_results = []
+        for collection_id, items in items_by_collection.items():
+            engine = self.local_engines[collection_id]["engine"]
+            try:
+                results = engine._get_full_content(items)
+                all_results.extend(results)
+            except Exception as e:
+                logger.error(f"Error getting full content from collection '{collection_id}': {e}")
+                # Fall back to returning the items without full content
+                all_results.extend(items)
+        
+        # Add any items that weren't processed
+        processed_ids = set(item["id"] for item in all_results)
+        for item in relevant_items:
+            if item["id"] not in processed_ids:
+                all_results.append(item)
+        
+        return all_results