feat: different search_args and docstrings

packages/leann-core/src/leann/searcher_base.py

@@ -1,9 +1,8 @@
-
 import json
 import pickle
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Dict, Any, List
+from typing import Dict, Any, Literal
 
 import numpy as np
 
@@ -40,7 +39,9 @@ class BaseSearcher(LeannBackendSearcherInterface, ABC):
 
         self.embedding_model = self.meta.get("embedding_model")
         if not self.embedding_model:
-            print("WARNING: embedding_model not found in meta.json. Recompute will fail.")
+            print(
+                "WARNING: embedding_model not found in meta.json. Recompute will fail."
+            )
 
         self.label_map = self._load_label_map()
 
@@ -54,7 +55,7 @@ class BaseSearcher(LeannBackendSearcherInterface, ABC):
         meta_path = self.index_dir / f"{self.index_path.name}.meta.json"
         if not meta_path.exists():
             raise FileNotFoundError(f"Leann metadata file not found at {meta_path}")
-        with open(meta_path, 'r', encoding='utf-8') as f:
+        with open(meta_path, "r", encoding="utf-8") as f:
             return json.load(f)
 
     def _load_label_map(self) -> Dict[int, str]:
@@ -62,16 +63,20 @@ class BaseSearcher(LeannBackendSearcherInterface, ABC):
         label_map_file = self.index_dir / "leann.labels.map"
         if not label_map_file.exists():
             raise FileNotFoundError(f"Label map file not found: {label_map_file}")
-        with open(label_map_file, 'rb') as f:
+        with open(label_map_file, "rb") as f:
             return pickle.load(f)
 
-    def _ensure_server_running(self, passages_source_file: str, port: int, **kwargs) -> None:
+    def _ensure_server_running(
+        self, passages_source_file: str, port: int, **kwargs
+    ) -> None:
         """
         Ensures the embedding server is running if recompute is needed.
         This is a helper for subclasses.
         """
         if not self.embedding_model:
-            raise ValueError("Cannot use recompute mode without 'embedding_model' in meta.json.")
+            raise ValueError(
+                "Cannot use recompute mode without 'embedding_model' in meta.json."
+            )
 
         server_started = self.embedding_server_manager.start_server(
             port=port,
@@ -85,15 +90,38 @@ class BaseSearcher(LeannBackendSearcherInterface, ABC):
             raise RuntimeError(f"Failed to start embedding server on port {port}")
 
     @abstractmethod
-    def search(self, query: np.ndarray, top_k: int, **kwargs) -> Dict[str, Any]:
+    def search(
+        self,
+        query: np.ndarray,
+        top_k: int,
+        complexity: int = 64,
+        beam_width: int = 1,
+        prune_ratio: float = 0.0,
+        recompute_embeddings: bool = False,
+        pruning_strategy: Literal["global", "local", "proportional"] = "global",
+        zmq_port: int = 5557,
+        **kwargs,
+    ) -> Dict[str, Any]:
         """
         Search for the top_k nearest neighbors of the query vector.
-        Must be implemented by subclasses.
+
+        Args:
+            query: Query vectors (B, D) where B is batch size, D is dimension
+            top_k: Number of nearest neighbors to return
+            complexity: Search complexity/candidate list size, higher = more accurate but slower
+            beam_width: Number of parallel search paths/IO requests per iteration
+            prune_ratio: Ratio of neighbors to prune via approximate distance (0.0-1.0)
+            recompute_embeddings: Whether to fetch fresh embeddings from server vs use stored PQ codes
+            pruning_strategy: PQ candidate selection strategy - "global" (default), "local", or "proportional"
+            zmq_port: ZMQ port for embedding server communication
+            **kwargs: Backend-specific parameters (e.g., batch_size, dedup_node_dis, etc.)
+
+        Returns:
+            Dict with 'labels' (list of lists) and 'distances' (ndarray)
         """
         pass
 
     def __del__(self):
         """Ensures the embedding server is stopped when the searcher is destroyed."""
-        if hasattr(self, 'embedding_server_manager'):
+        if hasattr(self, "embedding_server_manager"):
             self.embedding_server_manager.stop_server()
-