diff --git a/packages/leann-backend-hnsw/leann_backend_hnsw/hnsw_backend.py b/packages/leann-backend-hnsw/leann_backend_hnsw/hnsw_backend.py index e1afb36..d6b87f9 100644 --- a/packages/leann-backend-hnsw/leann_backend_hnsw/hnsw_backend.py +++ b/packages/leann-backend-hnsw/leann_backend_hnsw/hnsw_backend.py @@ -124,7 +124,9 @@ class HNSWSearcher(BaseSearcher): ) from . import faiss # type: ignore - self.distance_metric = self.meta.get("distance_metric", "mips").lower() + self.distance_metric = ( + self.meta.get("backend_kwargs", {}).get("distance_metric", "mips").lower() + ) metric_enum = get_metric_map().get(self.distance_metric) if metric_enum is None: raise ValueError(f"Unsupported distance_metric '{self.distance_metric}'.") @@ -200,6 +202,16 @@ class HNSWSearcher(BaseSearcher): params.efSearch = complexity params.beam_size = beam_width + # For OpenAI embeddings with cosine distance, disable relative distance check + # This prevents early termination when all scores are in a narrow range + embedding_model = self.meta.get("embedding_model", "").lower() + if self.distance_metric == "cosine" and any( + openai_model in embedding_model for openai_model in ["text-embedding", "openai"] + ): + params.check_relative_distance = False + else: + params.check_relative_distance = True + # PQ pruning: direct mapping to HNSW's pq_pruning_ratio params.pq_pruning_ratio = prune_ratio diff --git a/packages/leann-core/src/leann/embedding_server_manager.py b/packages/leann-core/src/leann/embedding_server_manager.py index 3c8a028..5a75ac7 100644 --- a/packages/leann-core/src/leann/embedding_server_manager.py +++ b/packages/leann-core/src/leann/embedding_server_manager.py @@ -293,6 +293,8 @@ class EmbeddingServerManager: command.extend(["--passages-file", str(passages_file)]) if embedding_mode != "sentence-transformers": command.extend(["--embedding-mode", embedding_mode]) + if kwargs.get("distance_metric"): + command.extend(["--distance-metric", kwargs["distance_metric"]]) return command diff --git a/packages/leann-core/src/leann/searcher_base.py b/packages/leann-core/src/leann/searcher_base.py index cc02fb2..02ec430 100644 --- a/packages/leann-core/src/leann/searcher_base.py +++ b/packages/leann-core/src/leann/searcher_base.py @@ -63,12 +63,19 @@ class BaseSearcher(LeannBackendSearcherInterface, ABC): if not self.embedding_model: raise ValueError("Cannot use recompute mode without 'embedding_model' in meta.json.") + # Get distance_metric from meta if not provided in kwargs + distance_metric = ( + kwargs.get("distance_metric") + or self.meta.get("backend_kwargs", {}).get("distance_metric") + or "mips" + ) + server_started, actual_port = self.embedding_server_manager.start_server( port=port, model_name=self.embedding_model, embedding_mode=self.embedding_mode, passages_file=passages_source_file, - distance_metric=kwargs.get("distance_metric"), + distance_metric=distance_metric, enable_warmup=kwargs.get("enable_warmup", False), ) if not server_started: