feat: make the score in VectorStoreResult consistent (#482)

ludwiktrammer · web-flow · commit ea7388325c28 · 2025-04-01T10:39:39.000+02:00
diff --git a/examples/document-search/chroma.py b/examples/document-search/chroma.py
@@ -84,7 +84,7 @@ async def main() -> None:
         index_name="jokes",
         default_options=VectorStoreOptions(
             k=10,
-            max_distance=0.22,
+            score_threshold=0.88,
         ),
         embedder=embedder,
     )
@@ -103,7 +103,7 @@ async def main() -> None:
     query = "I'm boiling my water and I need a joke"
     vector_store_kwargs = {
         "k": 2,
-        "max_distance": 0.6,
+        "score_threshold": 0.4,
     }
     results = await document_search.search(
         query,
diff --git a/examples/document-search/configurable.py b/examples/document-search/configurable.py
@@ -73,7 +73,7 @@ class to rephrase the query.
             "distance_method": "l2",
             "default_options": {
                 "k": 3,
-                "max_distance": 1.2,
+                "score_threshold": -1.2,
             },
             "embedder": {
                 "type": "ragbits.core.embeddings.litellm:LiteLLMEmbedder",
diff --git a/examples/document-search/otel.py b/examples/document-search/otel.py
@@ -121,7 +121,7 @@ async def main() -> None:
     query = "I'm boiling my water and I need a joke"
     vector_store_kwargs = {
         "k": 2,
-        "max_distance": None,
+        "score_threshold": None,
     }
     results = await document_search.search(
         query,
diff --git a/examples/document-search/qdrant.py b/examples/document-search/qdrant.py
@@ -94,7 +94,7 @@ async def main() -> None:
     query = "I'm boiling my water and I need a joke"
     vector_store_kwargs = {
         "k": 2,
-        "max_distance": 0.6,
+        "score_threshold": 0.6,
     }
     results = await document_search.search(
         query,
diff --git a/examples/evaluation/document-search/advanced/config/pipeline/vector_store/chroma.yaml b/examples/evaluation/document-search/advanced/config/pipeline/vector_store/chroma.yaml
@@ -6,7 +6,7 @@ config:
   distance_method: l2
   default_options:
     k: 3
-    max_distance: 1.2
+    score_threshold: -1.2
   embedder:
     type: ragbits.core.embeddings.litellm:LiteLLMEmbedder
     config:
diff --git a/examples/evaluation/document-search/advanced/config/pipeline/vector_store/chroma_optimization.yaml b/examples/evaluation/document-search/advanced/config/pipeline/vector_store/chroma_optimization.yaml
@@ -6,7 +6,7 @@ config:
   distance_method: l2
   default_options:
     k: 3
-    max_distance: 1.2
+    score_threshold: -1.2
   embedder:
     type: ragbits.core.embeddings.litellm:LiteLLMEmbedder
     config:
diff --git a/examples/evaluation/document-search/basic/evaluate.py b/examples/evaluation/document-search/basic/evaluate.py
@@ -35,7 +35,7 @@
                     "distance_method": "l2",
                     "default_options": {
                         "k": 3,
-                        "max_distance": 1.2,
+                        "score_threshold": -1.2,
                     },
                     "embedder": {
                         "type": "ragbits.core.embeddings.litellm:LiteLLMEmbedder",
diff --git a/packages/ragbits-core/CHANGELOG.md b/packages/ragbits-core/CHANGELOG.md
@@ -1,7 +1,7 @@
 # CHANGELOG
 
 ## Unreleased
-
+- Make the score in VectorStoreResult consistent (always bigger is better)
 - Add router option to LiteLLMEmbedder (#440)
 - Fix: make unflatten_dict symmetric to flatten_dict (#461)
 
@@ -11,7 +11,6 @@
 - Fix Qdrant vector store serialization (#419)
 
 ## 0.11.0 (2025-03-25)
-
 - Add HybridSearchVectorStore which can aggregate results from multiple VectorStores (#412)
 
 ## 0.10.2 (2025-03-21)
diff --git a/packages/ragbits-core/src/ragbits/core/vector_stores/_cli.py b/packages/ragbits-core/src/ragbits/core/vector_stores/_cli.py
@@ -101,7 +101,7 @@ async def run() -> None:
 def query(
     text: Annotated[str, typer.Argument(help="Text to query the vector store with")],
     k: Annotated[int, typer.Option(help="Number of entries to retrieve")] = 5,
-    max_distance: Annotated[float | None, typer.Option(help="Maximum distance to the query vector")] = None,
+    score_threshold: Annotated[float | None, typer.Option(help="Minimum score for result to be returned")] = None,
     columns: Annotated[
         str,
         typer.Option(
@@ -117,7 +117,7 @@ def query(
     async def run() -> None:
         if state.vector_store is None:
             raise ValueError("Vector store not initialized")
-        options = VectorStoreOptions(k=k, max_distance=max_distance)
+        options = VectorStoreOptions(k=k, score_threshold=score_threshold)
 
         entries = await state.vector_store.retrieve(
             text=text,
diff --git a/packages/ragbits-core/src/ragbits/core/vector_stores/base.py b/packages/ragbits-core/src/ragbits/core/vector_stores/base.py
@@ -44,6 +44,11 @@ class VectorStoreResult(BaseModel):
     """
     An object representing a query result from a vector store.
     Contains the entry, its vector, and the similarity score.
+
+    The bigger the score, the more similar the entry is to the query.
+    This holds even when using metrics that naturally follow the
+    opposite convention (e.g. Euclidean distance). In this case,
+    the score is simply the negative of the distance.
     """
 
     entry: VectorStoreEntry
@@ -54,10 +59,17 @@ class VectorStoreResult(BaseModel):
 class VectorStoreOptions(Options):
     """
     An object representing the options for the vector store.
+
+    Attributes:
+        k: The number of entries to return.
+        score_threshold: The minimum similarity score for an entry to be returned.
+            Note that this is based on score, which may be different from the raw
+            similarity metric used by the vector store (see `VectorStoreResult`
+            for more details).
     """
 
     k: int = 5
-    max_distance: float | None = None
+    score_threshold: float | None = None
 
 
 VectorStoreOptionsT = TypeVar("VectorStoreOptionsT", bound=VectorStoreOptions)
diff --git a/packages/ragbits-core/src/ragbits/core/vector_stores/chroma.py b/packages/ragbits-core/src/ragbits/core/vector_stores/chroma.py
@@ -19,6 +19,11 @@
     WhereQuery,
 )
 
+# The metrics for which smaller values are considered better.
+# This is used to determine whether to reverse the score.
+# (since Ragbits always follows the "biggest is better" convention)
+_SMALLER_IS_BETTER_METRICS = {"l2"}
+
 
 class ChromaVectorStore(VectorStoreWithExternalEmbedder[VectorStoreOptions]):
     """
@@ -136,6 +141,24 @@ async def store(self, entries: list[VectorStoreEntry]) -> None:
                 documents=documents,
             )
 
+    def _calculate_score(self, distance: float) -> float:
+        """
+        Calculates the score based on the distance.
+
+        Args:
+            distance: The distance.
+
+        Returns:
+            The score.
+        """
+        if self._distance_method in _SMALLER_IS_BETTER_METRICS:
+            # We need "bigger is better" convention, so we reverse the distance
+            return -1 * distance
+
+        # Chroma converts "bigger is better" metrics to "smaller is better" by subtracting the distance from 1
+        # So we need to reverse this operation
+        return 1 - distance
+
     async def retrieve(
         self,
         text: str,
@@ -155,6 +178,7 @@ async def retrieve(
             MetadataNotFoundError: If the metadata is not found.
         """
         merged_options = (self.default_options | options) if options else self.default_options
+
         with trace(
             text=text,
             options=merged_options.dict(),
@@ -178,7 +202,7 @@ async def retrieve(
             )
 
             ids = [id for batch in results.get("ids", []) for id in batch]
-            distances = [distance for batch in results.get("distances") or [] for distance in batch]
+            scores = [self._calculate_score(distance) for batch in results.get("distances") or [] for distance in batch]
             documents = [document for batch in results.get("documents") or [] for document in batch]
             embeddings = [embedding for batch in results.get("embeddings") or [] for embedding in batch]
 
@@ -191,7 +215,7 @@ async def retrieve(
 
             outputs.results = [
                 VectorStoreResult(
-                    score=distance,
+                    score=score,
                     vector=vector,
                     entry=VectorStoreEntry(
                         id=id,
@@ -200,10 +224,10 @@ async def retrieve(
                         metadata=metadata,
                     ),
                 )
-                for id, metadata, distance, document, image, vector in zip(
-                    ids, unflattened_metadatas, distances, documents, images, embeddings, strict=True
+                for id, metadata, score, document, image, vector in zip(
+                    ids, unflattened_metadatas, scores, documents, images, embeddings, strict=True
                 )
-                if merged_options.max_distance is None or distance <= merged_options.max_distance
+                if merged_options.score_threshold is None or score >= merged_options.score_threshold
             ]
 
             return outputs.results
diff --git a/packages/ragbits-core/src/ragbits/core/vector_stores/hybrid.py b/packages/ragbits-core/src/ragbits/core/vector_stores/hybrid.py
@@ -36,16 +36,7 @@ class OrderedHybridRetrivalStrategy(HybridRetrivalStrategy):
     A class that orders the results by score and deduplicates them by choosing the first occurrence of each entry.
     """
 
-    def __init__(self, reverse: bool = False) -> None:
-        """
-        Constructs a new OrderedHybridRetrivalStrategy instance.
-
-        Args:
-            reverse: if True orders the results in descending order by score, otherwise in ascending order.
-        """
-        self._reverse = reverse
-
-    def join(self, results: list[list[VectorStoreResult]]) -> list[VectorStoreResult]:
+    def join(self, results: list[list[VectorStoreResult]]) -> list[VectorStoreResult]:  # noqa: PLR6301
         """
         Joins the multiple lists of results into a single list.
 
@@ -56,7 +47,7 @@ def join(self, results: list[list[VectorStoreResult]]) -> list[VectorStoreResult
             The joined list of results.
         """
         all_results = [result for sublist in results for result in sublist]
-        all_results.sort(key=lambda result: result.score, reverse=self._reverse)
+        all_results.sort(key=lambda result: result.score, reverse=True)
 
         return list({result.entry.id: result for result in all_results}.values())
 
diff --git a/packages/ragbits-core/src/ragbits/core/vector_stores/in_memory.py b/packages/ragbits-core/src/ragbits/core/vector_stores/in_memory.py
@@ -88,12 +88,12 @@ async def retrieve(
             results: list[VectorStoreResult] = []
 
             for entry_id, vector in self._embeddings.items():
-                distance = float(np.linalg.norm(np.array(vector) - np.array(query_vector)))
-                result = VectorStoreResult(entry=self._entries[entry_id], vector=vector, score=distance)
-                if merged_options.max_distance is None or result.score <= merged_options.max_distance:
+                score = float(np.linalg.norm(np.array(vector) - np.array(query_vector))) * -1
+                result = VectorStoreResult(entry=self._entries[entry_id], vector=vector, score=score)
+                if merged_options.score_threshold is None or result.score >= merged_options.score_threshold:
                     results.append(result)
 
-            outputs.results = sorted(results, key=lambda r: r.score)[: merged_options.k]
+            outputs.results = sorted(results, key=lambda r: r.score, reverse=True)[: merged_options.k]
             return outputs.results
 
     @traceable
diff --git a/packages/ragbits-core/src/ragbits/core/vector_stores/pgvector.py b/packages/ragbits-core/src/ragbits/core/vector_stores/pgvector.py
@@ -1,6 +1,6 @@
 import json
 import re
-from typing import Any
+from typing import Any, NamedTuple
 from uuid import UUID
 
 import asyncpg
@@ -18,15 +18,26 @@
     WhereQuery,
 )
 
+
+class DistanceOp(NamedTuple):
+    """
+    Structure for keeping details of distance pgvector's operations.
+    """
+
+    function_name: str
+    operator: str
+    score_formula: str  # formula to calculate score based on distance
+
+
 DISTANCE_OPS = {
-    "cosine": ("vector_cosine_ops", "<=>"),
-    "l2": ("vector_l2_ops", "<->"),
-    "l1": ("vector_l1_ops", "<+>"),
-    "ip": ("vector_ip_ops", "<#>"),
-    "bit_hamming": ("bit_hamming_ops", "<~>"),
-    "bit_jaccard": ("bit_jaccard_ops", "<%>"),
-    "sparsevec_l2": ("sparsevec_l2_ops", "<->"),
-    "halfvec_l2": ("halfvec_l2_ops", "<->"),
+    "cosine": DistanceOp("vector_cosine_ops", "<=>", "1 - distance"),
+    "l2": DistanceOp("vector_l2_ops", "<->", "distance * -1"),
+    "l1": DistanceOp("vector_l1_ops", "<+>", "distance * -1"),
+    "ip": DistanceOp("vector_ip_ops", "<#>", "distance * -1"),
+    "bit_hamming": DistanceOp("bit_hamming_ops", "<~>", "distance * -1"),
+    "bit_jaccard": DistanceOp("bit_jaccard_ops", "<%>", "distance * -1"),
+    "sparsevec_l2": DistanceOp("sparsevec_l2_ops", "<->", "distance * -1"),
+    "halfvec_l2": DistanceOp("halfvec_l2_ops", "<->", "distance * -1"),
 }
 
 
@@ -115,22 +126,23 @@ def _create_retrieve_query(
         Returns:
             str: sql query.
         """
-        distance_operator = DISTANCE_OPS[self._distance_method][1]
+        distance_operator = DISTANCE_OPS[self._distance_method].operator
         if not query_options:
             query_options = self.default_options
 
+        # We select both distance and score because pgvector require ordering by distance (ascending).
+        # in order to use its KNN index. We calculate the score based on the distance.
+        score_formula = DISTANCE_OPS[self._distance_method].score_formula.replace(
+            "distance", f"(vector {distance_operator} $1)"
+        )
         # _table_name has been validated in the class constructor, and it is a valid table name.
-        query = f"SELECT *, vector {distance_operator} $1 as distance FROM {self._table_name}"  # noqa S608
+        query = f"SELECT *, vector {distance_operator} $1 as distance, {score_formula} as score FROM {self._table_name}"  # noqa S608
 
         values: list[Any] = [str(vector)]
 
-        if query_options.max_distance and self._distance_method == "ip":
-            query += """ WHERE distance BETWEEN $2 AND $3"""
-            values.extend([-1 * query_options.max_distance, query_options.max_distance])
-
-        elif query_options.max_distance:
-            query += " WHERE distance < $2"
-            values.extend([query_options.max_distance])
+        if query_options.score_threshold is not None:
+            query += " WHERE score >= $2"
+            values.extend([query_options.score_threshold])
 
         query += " ORDER BY distance"
 
@@ -181,7 +193,7 @@ async def create_table(self) -> None:
                     SELECT FROM information_schema.tables
                     WHERE table_name = $1
                 ); """
-            distance = DISTANCE_OPS[self._distance_method][0]
+            distance = DISTANCE_OPS[self._distance_method].function_name
             create_vector_extension = "CREATE EXTENSION IF NOT EXISTS vector;"
             # _table_name and has been validated in the class constructor, and it is a valid table name.
             # _vector_size has been validated in the class constructor, and it is a valid vector size.
@@ -354,7 +366,7 @@ async def retrieve(
                             metadata=json.loads(record["metadata"]),
                         ),
                         vector=json.loads(record["vector"]),
-                        score=record["distance"],
+                        score=record["score"],
                     )
                     for record in results
                 ]
diff --git a/packages/ragbits-core/src/ragbits/core/vector_stores/qdrant.py b/packages/ragbits-core/src/ragbits/core/vector_stores/qdrant.py
@@ -6,7 +6,14 @@
 import httpx
 import qdrant_client
 from qdrant_client import AsyncQdrantClient, models
-from qdrant_client.models import Distance, FieldCondition, Filter, MatchValue, VectorParams
+from qdrant_client.local.distances import DistanceOrder, distance_to_order
+from qdrant_client.models import (
+    Distance,
+    FieldCondition,
+    Filter,
+    MatchValue,
+    VectorParams,
+)
 from typing_extensions import Self
 
 from ragbits.core.audit import trace
@@ -170,7 +177,14 @@ async def retrieve(self, text: str, options: VectorStoreOptionsT | None = None)
             The retrieved entries.
         """
         merged_options = (self.default_options | options) if options else self.default_options
-        score_threshold = 1 - merged_options.max_distance if merged_options.max_distance else None
+
+        # Ragbits has a "larger is better" convention for all scores, so we need to reverse the score if the distance
+        # method is "smaller is better".
+        reverse_score = distance_to_order(self._distance_method) == DistanceOrder.SMALLER_IS_BETTER
+        score_multiplier = -1 if reverse_score else 1
+        score_threshold = (
+            None if merged_options.score_threshold is None else merged_options.score_threshold * score_multiplier
+        )
         with trace(
             text=text,
             options=merged_options,
@@ -197,7 +211,7 @@ async def retrieve(self, text: str, options: VectorStoreOptionsT | None = None)
                 outputs.results.append(
                     VectorStoreResult(
                         entry=entry,
-                        score=point.score,
+                        score=point.score * score_multiplier,
                         vector=cast(list[float], point.vector),
                     )
                 )
diff --git a/packages/ragbits-core/tests/integration/vector_stores/test_vector_store.py b/packages/ragbits-core/tests/integration/vector_stores/test_vector_store.py
diff --git a/packages/ragbits-core/tests/unit/vector_stores/test_chroma.py b/packages/ragbits-core/tests/unit/vector_stores/test_chroma.py
diff --git a/packages/ragbits-core/tests/unit/vector_stores/test_from_config.py b/packages/ragbits-core/tests/unit/vector_stores/test_from_config.py
diff --git a/packages/ragbits-core/tests/unit/vector_stores/test_hybrid.py b/packages/ragbits-core/tests/unit/vector_stores/test_hybrid.py
diff --git a/packages/ragbits-core/tests/unit/vector_stores/test_in_memory.py b/packages/ragbits-core/tests/unit/vector_stores/test_in_memory.py
diff --git a/packages/ragbits-core/tests/unit/vector_stores/test_pgvector.py b/packages/ragbits-core/tests/unit/vector_stores/test_pgvector.py
diff --git a/packages/ragbits-core/tests/unit/vector_stores/test_qdrant.py b/packages/ragbits-core/tests/unit/vector_stores/test_qdrant.py

Original file line number	Diff line number	Diff line change
`@@ -121,7 +121,7 @@ async def main() -> None:`
`121`	`121`	`query = "I'm boiling my water and I need a joke"`
`122`	`122`	`vector_store_kwargs = {`
`123`	`123`	`"k": 2,`
`124`		`- "max_distance": None,`
	`124`	`+ "score_threshold": None,`
`125`	`125`	`}`
`126`	`126`	`results = await document_search.search(`
`127`	`127`	`query,`
Original file line number	Diff line number	Diff line change
`@@ -94,7 +94,7 @@ async def main() -> None:`
`94`	`94`	`query = "I'm boiling my water and I need a joke"`
`95`	`95`	`vector_store_kwargs = {`
`96`	`96`	`"k": 2,`
`97`		`- "max_distance": 0.6,`
	`97`	`+ "score_threshold": 0.6,`
`98`	`98`	`}`
`99`	`99`	`results = await document_search.search(`
`100`	`100`	`query,`