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_local_all.py

@@ -3,18 +3,19 @@ Search engine that searches across all local collections
 """
 
 import logging
-from typing import Dict, List, Any, Optional
+from typing import Any, Dict, List, Optional
 
 import toml
 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.config import LOCAL_COLLECTIONS_FILE
+from ...config.config_files import LOCAL_COLLECTIONS_FILE
+from ..search_engine_base import BaseSearchEngine
+from ..search_engine_factory import create_search_engine
 
 # Setup logging
 logger = logging.getLogger(__name__)
 
+
 class LocalAllSearchEngine(BaseSearchEngine):
     """
     Search engine that searches across all local document collections.
@@ -26,11 +27,11 @@ class LocalAllSearchEngine(BaseSearchEngine):
         llm: Optional[BaseLLM] = None,
         max_results: int = 10,
         max_filtered_results: Optional[int] = None,
-        **kwargs
+        **kwargs,
     ):
         """
         Initialize the local all-collections search engine.
-        
+
         Args:
             llm: Language model for relevance filtering
             max_results: Maximum number of search results
@@ -38,8 +39,10 @@ class LocalAllSearchEngine(BaseSearchEngine):
             **kwargs: Additional parameters passed to LocalSearchEngine instances
         """
         # 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
+        )
+
         # Find all local collection search engines
         self.local_engines = {}
         try:
@@ -48,78 +51,79 @@ class LocalAllSearchEngine(BaseSearchEngine):
             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
+                        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", "")
+                            "description": collection.get("description", ""),
                         }
                 except Exception as e:
-                    logger.error(f"Error creating search engine for collection '{collection_id}': {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
-        )
-        
+        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]]:
+        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
         """
@@ -131,7 +135,7 @@ class LocalAllSearchEngine(BaseSearchEngine):
                 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():
@@ -140,14 +144,16 @@ class LocalAllSearchEngine(BaseSearchEngine):
                 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}")
+                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