add normalize cosine distance flag

Author: rbs333

redisvl/index/index.py

@@ -16,7 +16,12 @@
     Union,
 )
 
-from redisvl.utils.utils import deprecated_argument, deprecated_function, sync_wrapper
+from redisvl.utils.utils import (
+    deprecated_argument,
+    deprecated_function,
+    norm_cosine_distance,
+    sync_wrapper,
+)
 
 if TYPE_CHECKING:
     from redis.commands.search.aggregation import AggregateResult
@@ -30,14 +35,21 @@
 
 from redisvl.exceptions import RedisModuleVersionError, RedisSearchError
 from redisvl.index.storage import BaseStorage, HashStorage, JsonStorage
-from redisvl.query import BaseQuery, CountQuery, FilterQuery
+from redisvl.query import (
+    BaseQuery,
+    CountQuery,
+    FilterQuery,
+    VectorQuery,
+    VectorRangeQuery,
+)
 from redisvl.query.filter import FilterExpression
 from redisvl.redis.connection import (
     RedisConnectionFactory,
     convert_index_info_to_schema,
 )
 from redisvl.redis.utils import convert_bytes
 from redisvl.schema import IndexSchema, StorageType
+from redisvl.schema.fields import VectorDistanceMetric
 from redisvl.utils.log import get_logger
 
 logger = get_logger(__name__)
@@ -50,7 +62,7 @@
 
 
 def process_results(
-    results: "Result", query: BaseQuery, storage_type: StorageType
+    results: "Result", query: BaseQuery, schema: IndexSchema
 ) -> List[Dict[str, Any]]:
     """Convert a list of search Result objects into a list of document
     dictionaries.
@@ -75,11 +87,18 @@ def process_results(
 
     # Determine if unpacking JSON is needed
     unpack_json = (
-        (storage_type == StorageType.JSON)
+        (schema.index.storage_type == StorageType.JSON)
         and isinstance(query, FilterQuery)
         and not query._return_fields  # type: ignore
     )
 
+    normalize_cosine_distance = (
+        (isinstance(query, VectorQuery) or isinstance(query, VectorRangeQuery))
+        and query._normalize_cosine_distance
+        and schema.fields[query._vector_field_name].attrs.distance_metric  # type: ignore
+        == VectorDistanceMetric.COSINE
+    )
+
     # Process records
     def _process(doc: "Document") -> Dict[str, Any]:
         doc_dict = doc.__dict__
@@ -93,6 +112,12 @@ def _process(doc: "Document") -> Dict[str, Any]:
                 return {"id": doc_dict.get("id"), **json_data}
             raise ValueError(f"Unable to parse json data from Redis {json_data}")
 
+        if normalize_cosine_distance:
+            # convert float back to string to be consistent
+            doc_dict[query.DISTANCE_ID] = str(  # type: ignore
+                norm_cosine_distance(float(doc_dict[query.DISTANCE_ID]))  # type: ignore
+            )
+
         # Remove 'payload' if present
         doc_dict.pop("payload", None)
 
@@ -665,9 +690,7 @@ def search(self, *args, **kwargs) -> "Result":
     def _query(self, query: BaseQuery) -> List[Dict[str, Any]]:
         """Execute a query and process results."""
         results = self.search(query.query, query_params=query.params)
-        return process_results(
-            results, query=query, storage_type=self.schema.index.storage_type
-        )
+        return process_results(results, query=query, schema=self.schema)
 
     def query(self, query: BaseQuery) -> List[Dict[str, Any]]:
         """Execute a query on the index.
@@ -1219,9 +1242,7 @@ async def search(self, *args, **kwargs) -> "Result":
     async def _query(self, query: BaseQuery) -> List[Dict[str, Any]]:
         """Asynchronously execute a query and process results."""
         results = await self.search(query.query, query_params=query.params)
-        return process_results(
-            results, query=query, storage_type=self.schema.index.storage_type
-        )
+        return process_results(results, query=query, schema=self.schema)
 
     async def query(self, query: BaseQuery) -> List[Dict[str, Any]]:
         """Asynchronously execute a query on the index.

redisvl/query/query.py

@@ -188,6 +188,7 @@ def __init__(
         dialect: int = 2,
         sort_by: Optional[str] = None,
         in_order: bool = False,
+        normalize_cosine_distance: bool = False,
     ):
         """A query for running a vector search along with an optional filter
         expression.
@@ -213,6 +214,9 @@ def __init__(
             in_order (bool): Requires the terms in the field to have
                 the same order as the terms in the query filter, regardless of
                 the offsets between them. Defaults to False.
+            normalize_cosine_distance (bool): by default Redis returns cosine distance as a value
+                between 0 and 2 where 0 is the best match. If set to True, the cosine distance will be
+                converted to cosine similarity with a value between 0 and 1 where 1 is the best match.
 
         Raises:
             TypeError: If filter_expression is not of type redisvl.query.FilterExpression
@@ -224,6 +228,7 @@ def __init__(
         self._vector_field_name = vector_field_name
         self._dtype = dtype
         self._num_results = num_results
+        self._normalize_cosine_distance = normalize_cosine_distance
         self.set_filter(filter_expression)
         query_string = self._build_query_string()
 
@@ -284,6 +289,7 @@ def __init__(
         dialect: int = 2,
         sort_by: Optional[str] = None,
         in_order: bool = False,
+        normalize_cosine_distance: bool = False,
     ):
         """A query for running a filtered vector search based on semantic
         distance threshold.
@@ -312,6 +318,9 @@ def __init__(
             in_order (bool): Requires the terms in the field to have
                 the same order as the terms in the query filter, regardless of
                 the offsets between them. Defaults to False.
+            normalize_cosine_distance (bool): by default Redis returns cosine distance as a value
+                between 0 and 2 where 0 is the best match. If set to True, the cosine distance will be
+                converted to cosine similarity with a value between 0 and 1 where 1 is the best match.
 
         Raises:
             TypeError: If filter_expression is not of type redisvl.query.FilterExpression
@@ -324,6 +333,7 @@ def __init__(
         self._vector_field_name = vector_field_name
         self._dtype = dtype
         self._num_results = num_results
+        self._normalize_cosine_distance = normalize_cosine_distance
         self.set_distance_threshold(distance_threshold)
         self.set_filter(filter_expression)
         query_string = self._build_query_string()

redisvl/utils/utils.py

@@ -191,3 +191,10 @@ def wrapper():
             return
 
     return wrapper
+
+
+def norm_cosine_distance(value: float) -> float:
+    """
+    Normalize the cosine distance to a similarity score between 0 and 1.
+    """
+    return (2 - value) / 2

tests/integration/test_query.py

@@ -52,6 +52,24 @@ def sorted_vector_query():
     )
 
 
+@pytest.fixture
+def normalized_vector_query():
+    return VectorQuery(
+        vector=[0.1, 0.1, 0.5],
+        vector_field_name="user_embedding",
+        normalize_cosine_distance=True,
+        return_score=True,
+        return_fields=[
+            "user",
+            "credit_score",
+            "age",
+            "job",
+            "location",
+            "last_updated",
+        ],
+    )
+
+
 @pytest.fixture
 def filter_query():
     return FilterQuery(
@@ -83,6 +101,18 @@ def sorted_filter_query():
     )
 
 
+@pytest.fixture
+def normalized_range_query():
+    return RangeQuery(
+        vector=[0.1, 0.1, 0.5],
+        vector_field_name="user_embedding",
+        normalize_cosine_distance=True,
+        return_score=True,
+        return_fields=["user", "credit_score", "age", "job", "location"],
+        distance_threshold=0.2,
+    )
+
+
 @pytest.fixture
 def range_query():
     return RangeQuery(
@@ -154,6 +184,56 @@ def hash_preprocess(item: dict) -> dict:
     index.delete(drop=True)
 
 
+@pytest.fixture
+def L2_index(sample_data, redis_url):
+    # construct a search index from the schema
+    index = SearchIndex.from_dict(
+        {
+            "index": {
+                "name": "L2_index",
+                "prefix": "L2_index",
+                "storage_type": "hash",
+            },
+            "fields": [
+                {"name": "credit_score", "type": "tag"},
+                {"name": "job", "type": "text"},
+                {"name": "age", "type": "numeric"},
+                {"name": "last_updated", "type": "numeric"},
+                {"name": "location", "type": "geo"},
+                {
+                    "name": "user_embedding",
+                    "type": "vector",
+                    "attrs": {
+                        "dims": 3,
+                        "distance_metric": "L2",
+                        "algorithm": "flat",
+                        "datatype": "float32",
+                    },
+                },
+            ],
+        },
+        redis_url=redis_url,
+    )
+
+    # create the index (no data yet)
+    index.create(overwrite=True)
+
+    # Prepare and load the data
+    def hash_preprocess(item: dict) -> dict:
+        return {
+            **item,
+            "user_embedding": array_to_buffer(item["user_embedding"], "float32"),
+        }
+
+    index.load(sample_data, preprocess=hash_preprocess)
+
+    # run the test
+    yield index
+
+    # clean up
+    index.delete(drop=True)
+
+
 def test_search_and_query(index):
     # *=>[KNN 7 @user_embedding $vector AS vector_distance]
     v = VectorQuery(
@@ -531,3 +611,26 @@ def test_query_with_chunk_number_zero():
     assert (
         str(filter_conditions) == expected_query_str
     ), "Query with chunk_number zero is incorrect"
+
+
+def test_query_normalize_cosine_distance(index, normalized_vector_query):
+
+    res = index.query(normalized_vector_query)
+
+    for r in res:
+        assert 0 <= float(r["vector_distance"]) <= 1
+
+
+def test_query_normalize_cosine_distance_ip_distance(L2_index, normalized_vector_query):
+
+    res = L2_index.query(normalized_vector_query)
+
+    assert any(float(r["vector_distance"]) > 1 for r in res)
+
+
+def test_range_query_normalize_cosine_distance(index, normalized_range_query):
+
+    res = index.query(normalized_range_query)
+
+    for r in res:
+        assert 0 <= float(r["vector_distance"]) <= 1

tests/unit/test_utils.py

@@ -13,9 +13,16 @@
     assert_no_warnings,
     deprecated_argument,
     deprecated_function,
+    norm_cosine_distance,
 )
 
 
+def test_norm_cosine_distance():
+    input = 2
+    expected = 0
+    assert norm_cosine_distance(input) == expected
+
+
 def test_even_number_of_elements():
     """Test with an even number of elements"""
     values = ["key1", "value1", "key2", "value2"]