feat: different search_args and docstrings

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

@@ -175,7 +175,7 @@ class EmbeddingServerManager:
         self.backend_module_name = backend_module_name
         self.server_process: Optional[subprocess.Popen] = None
         self.server_port: Optional[int] = None
-        atexit.register(self.stop_server)
+        # atexit.register(self.stop_server)
 
     def start_server(self, port: int, model_name: str, **kwargs) -> bool:
         """

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

@@ -1,42 +1,55 @@
 from abc import ABC, abstractmethod
 import numpy as np
-from typing import Dict, Any
+from typing import Dict, Any, Literal
 
 class LeannBackendBuilderInterface(ABC):
-    """用于构建索引的后端接口"""
+    """Backend interface for building indexes"""
     
     @abstractmethod 
     def build(self, data: np.ndarray, index_path: str, **kwargs) -> None:
-        """构建索引
+        """Build index
         
         Args:
-            data: 向量数据 (N, D)
-            index_path: 索引保存路径
-            **kwargs: 后端特定的构建参数
+            data: Vector data (N, D)
+            index_path: Path to save index
+            **kwargs: Backend-specific build parameters
         """
         pass
 
 class LeannBackendSearcherInterface(ABC):
-    """用于搜索的后端接口"""
+    """Backend interface for searching"""
     
     @abstractmethod
     def __init__(self, index_path: str, **kwargs):
-        """初始化搜索器
+        """Initialize searcher
         
         Args:
-            index_path: 索引文件路径
-            **kwargs: 后端特定的加载参数
+            index_path: Path to index file
+            **kwargs: Backend-specific loading parameters
         """
         pass
     
     @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 nearest neighbors
         
         Args:
-            query: 查询向量 (1, D) 或 (B, D)
-            top_k: 返回的最近邻数量
-            **kwargs: 搜索参数
+            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", "local", or "proportional"
+            zmq_port: ZMQ port for embedding server communication
+            **kwargs: Backend-specific parameters
             
         Returns:
             {"labels": [...], "distances": [...]}
@@ -44,16 +57,16 @@ class LeannBackendSearcherInterface(ABC):
         pass
 
 class LeannBackendFactoryInterface(ABC):
-    """后端工厂接口"""
+    """Backend factory interface"""
     
     @staticmethod
     @abstractmethod
     def builder(**kwargs) -> LeannBackendBuilderInterface:
-        """创建 Builder 实例"""
+        """Create Builder instance"""
         pass
     
     @staticmethod
     @abstractmethod  
     def searcher(index_path: str, **kwargs) -> LeannBackendSearcherInterface:
-        """创建 Searcher 实例"""
+        """Create Searcher instance"""
         pass

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()
-