Add thresholding for Semantic indexes

Author: greengori11a

doc/source/semantic_indexes.rst

@@ -50,6 +50,18 @@ Can be queried using :class:`~neomodel.semantic_filters.FulltextFilter`. Such as
 
 Where the result will be a list of length topk of nodes with the form (ProductNode, score).
 
+If you would like to filter the nodes based on a threshold (ie. nodes similarity >= threshold), you can use the following::
+
+    from neomodel.semantic_filters import FulltextFilter
+    result = Product.nodes.filter(
+        fulltext_filter=FulltextFilter(
+            topk=10,
+            fulltext_attribute_name="description",
+            query_string="product",
+            threshold=0.08)).all()
+
+Only nodes above threshold = 0.08 will be returned.
+
 The :class:`~neomodel.semantic_filters.FulltextFilter` can be used in conjunction with the normal filter types.
 
 .. attention:: 
@@ -103,10 +115,26 @@ The following node vector index property::
 Can be queried using :class:`~neomodel.semantic_filters.VectorFilter`. Such as::
 
     from neomodel.semantic_filters import VectorFilter
-    result = someNode.nodes.filter(vector_filter=VectorFilter(topk=3, vector_attribute_name="vector")).all()
+    result = someNode.nodes.filter(
+        vector_filter=VectorFilter(
+            topk=3, 
+            vector_attribute_name="vector", 
+            candidate_vector=[0.25, 0.25])).all()
 
 Where the result will be a list of length topk of tuples having the form (someNode, score). 
 
+If you would like to filter the nodes based on a threshold (ie. nodes similarity >= threshold), you can use the following::
+
+    from neomodel.semantic_filters import VectorFilter
+    result = someNode.nodes.filter(
+    vector_filter=VectorFilter(
+        topk=3, 
+        vector_attribute_name="vector", 
+        candidate_vector=[0.25, 0.25], 
+        threshold=0.85)).all()
+
+Only nodes above threshold = 0.85 will be returned.
+
 The :class:`~neomodel.semantic_filters.VectorFilter` can be used in conjunction with the normal filter types.
 
 .. attention:: 

neomodel/async_/match.py

@@ -577,13 +577,17 @@ def build_vector_query(self):
                 f"Attribute {vector_filter.vector_attribute_name} is not declared with a vector index."
             )
 
+        if type(vector_filter.threshold) not in [float, type(None)]:
+            raise ValueError(f"Vector Filter Threshold must be a float or None.")
+
         vector_filter.index_name = f"vector_index_{source_class.__label__}_{vector_filter.vector_attribute_name}"
         vector_filter.node_set_label = source_class.__label__.lower()
 
         self._ast.vector_index_query = vector_filter
         self._ast.return_clause = f"{vector_filter.node_set_label}, score"
         self._ast.result_class = source_class.__class__
 
+
     def build_fulltext_query(self):
         """
         Query a free text indexed property on the node.
@@ -604,6 +608,9 @@ def build_fulltext_query(self):
                 f"Attribute {full_text_filter.fulltext_attribute_name} is not declared with a full text index."
             )
 
+        if type(full_text_filter.threshold) not in [float, type(None)]:
+            raise ValueError(f"Full Text Filter Threshold must be a float or None.")
+
         full_text_filter.index_name = f"fulltext_index_{source_class.__label__}_{full_text_filter.fulltext_attribute_name}"
         full_text_filter.node_set_label = source_class.__label__.lower()
 
@@ -1005,19 +1012,33 @@ def build_query(self) -> str:
         if self._ast.vector_index_query:
             query += f"""CALL () {{ 
                 CALL db.index.vector.queryNodes("{self._ast.vector_index_query.index_name}", {self._ast.vector_index_query.topk}, {self._ast.vector_index_query.vector}) 
-                YIELD node AS {self._ast.vector_index_query.node_set_label}, score 
+                YIELD node AS {self._ast.vector_index_query.node_set_label}, score """
+
+            if self._ast.vector_index_query.threshold:
+                query += f"""
+                WHERE score >= {self._ast.vector_index_query.threshold}
+                """
+
+            query += f"""
                 RETURN {self._ast.vector_index_query.node_set_label}, score 
-                }}"""
+            }}"""
 
             # This ensures that we bring the context of the new nodeSet and score along with us for metadata filtering
             query += f""" WITH {self._ast.vector_index_query.node_set_label}, score"""
 
         if self._ast.fulltext_index_query:
             query += f"""CALL () {{
                 CALL db.index.fulltext.queryNodes("{self._ast.fulltext_index_query.index_name}", "{self._ast.fulltext_index_query.query_string}")
-                YIELD node AS {self._ast.fulltext_index_query.node_set_label}, score
+                YIELD node AS {self._ast.fulltext_index_query.node_set_label}, score"""
+
+            if self._ast.fulltext_index_query.threshold:
+                query += f"""
+                WHERE score >= {self._ast.fulltext_index_query.threshold}
+                """
+
+            query += f"""
                 RETURN {self._ast.fulltext_index_query.node_set_label}, score LIMIT {self._ast.fulltext_index_query.topk}
-                }}
+            }}
                 """
             # This ensures that we bring the context of the new nodeSet and score along with us for metadata filtering
             query += f""" WITH {self._ast.fulltext_index_query.node_set_label}, score"""

neomodel/semantic_filters.py

@@ -1,3 +1,5 @@
+from typing import Union
+
 class VectorFilter(object):
     """
     Represents a CALL db.index.vector.query* neo functions call within the OGM
@@ -6,35 +8,42 @@ class VectorFilter(object):
     :type topk: int
     :param vector_attribute_name: The property name for vector indexed property on the searched object.
     :type vector_attribute_name: str
+    :param threshold: Threshold for vector similarity.
+    :type threshold: float
     :param candidate_vector: The vector you are finding the nearest topk neighbours for.
     :type candidate_vector: list[float]
 
     """
 
     def __init__(
-        self, topk: int, vector_attribute_name: str, candidate_vector: list[float]
+            self, topk: int, vector_attribute_name: str, candidate_vector: list[float], threshold: Union[float, None] =  None
     ):
         self.topk = topk
         self.vector_attribute_name = vector_attribute_name
+        self.threshold = threshold
         self.index_name = None
         self.node_set_label = None
         self.vector = candidate_vector
 
 class FulltextFilter(object):
     """
-    Represents a CALL db.index.fulltext.query* neo function call within the OGM.
+    Represents a CALL db.index.fulltext.query* neo functon call within the OGM.
     :param query_strng: The string you are finding the nearest
     :type query_string: str
     :param freetext_attribute_name: The property name for the free text indexed property.
     :type fulltext_attribute_name: str
+    :param threshold: Threshold for vector similarity.
+    :type threshold: float
     :param topk: Amount to nodes to return
     :type topk: int
 
     """
 
-    def __init__(self, query_string: str, fulltext_attribute_name: str, topk: int):
+    def __init__(self, query_string: str, fulltext_attribute_name: str, topk: int, threshold: Union[float, None] =  None
+    ):
+        self.topk = topk
         self.query_string = query_string
         self.fulltext_attribute_name = fulltext_attribute_name
+        self.threshold = threshold
         self.index_name = None
         self.node_set_label = None
-        self.topk = topk

neomodel/sync_/match.py

@@ -575,6 +575,9 @@ def build_vector_query(self):
                 f"Attribute {vector_filter.vector_attribute_name} is not declared with a vector index."
             )
 
+        if type(vector_filter.threshold) not in [float, type(None)]:
+            raise ValueError(f"Vector Filter Threshold must be a float or None.")
+
         vector_filter.index_name = f"vector_index_{source_class.__label__}_{vector_filter.vector_attribute_name}"
         vector_filter.node_set_label = source_class.__label__.lower()
 
@@ -602,6 +605,9 @@ def build_fulltext_query(self):
                 f"Attribute {full_text_filter.fulltext_attribute_name} is not declared with a full text index."
             )
 
+        if type(full_text_filter.threshold) not in [float, type(None)]:
+            raise ValueError(f"Full Text Filter Threshold must be a float or None.")
+
         full_text_filter.index_name = f"fulltext_index_{source_class.__label__}_{full_text_filter.fulltext_attribute_name}"
         full_text_filter.node_set_label = source_class.__label__.lower()
 
@@ -1003,19 +1009,33 @@ def build_query(self) -> str:
         if self._ast.vector_index_query:
             query += f"""CALL () {{ 
                 CALL db.index.vector.queryNodes("{self._ast.vector_index_query.index_name}", {self._ast.vector_index_query.topk}, {self._ast.vector_index_query.vector}) 
-                YIELD node AS {self._ast.vector_index_query.node_set_label}, score 
+                YIELD node AS {self._ast.vector_index_query.node_set_label}, score """
+
+            if self._ast.vector_index_query.threshold:
+                query += f"""
+                WHERE score >= {self._ast.vector_index_query.threshold}
+                """
+
+            query += f"""
                 RETURN {self._ast.vector_index_query.node_set_label}, score 
-                }}"""
+            }}"""
 
             # This ensures that we bring the context of the new nodeSet and score along with us for metadata filtering
             query += f""" WITH {self._ast.vector_index_query.node_set_label}, score"""
 
         if self._ast.fulltext_index_query:
             query += f"""CALL () {{
                 CALL db.index.fulltext.queryNodes("{self._ast.fulltext_index_query.index_name}", "{self._ast.fulltext_index_query.query_string}")
-                YIELD node AS {self._ast.fulltext_index_query.node_set_label}, score
+                YIELD node AS {self._ast.fulltext_index_query.node_set_label}, score"""
+
+            if self._ast.fulltext_index_query.threshold:
+                query += f"""
+                WHERE score >= {self._ast.fulltext_index_query.threshold}
+                """
+
+            query += f"""
                 RETURN {self._ast.fulltext_index_query.node_set_label}, score LIMIT {self._ast.fulltext_index_query.topk}
-                }}
+            }}
                 """
             # This ensures that we bring the context of the new nodeSet and score along with us for metadata filtering
             query += f""" WITH {self._ast.fulltext_index_query.node_set_label}, score"""

test/async_/test_fulltextfilter.py

@@ -122,6 +122,46 @@ class fulltextNodeBis(AsyncStructuredNode):
     assert all(isinstance(x[1], float) for x in result)
 
 
+@mark_async_test
+async def test_fulltextfilter_threshold():
+    """
+    Tests that the fulltext query is run, and only nodes above threshold returns.
+    """
+
+    if not await adb.version_is_higher_than("5.16"):
+        pytest.skip("Not supported before 5.16")
+
+    class fulltextNodeThresh(AsyncStructuredNode):
+        description = StringProperty(
+            fulltext_index=FulltextIndex(
+                analyzer="standard-no-stop-words", eventually_consistent=False
+            )
+        )
+        other = StringProperty()
+
+    await adb.install_labels(fulltextNodeThresh)
+
+    node1 = await fulltextNodeThresh(other="thing", description="Another thing").save()
+
+    node2 = await fulltextNodeThresh(
+        other="other thing", description="Another other thing"
+    ).save()
+
+    fulltextFilterThresh= fulltextNodeThresh.nodes.filter(
+        fulltext_filter=FulltextFilter(
+            topk=3, fulltext_attribute_name="description", query_string="thing", threshold=0.09
+        ),
+        other="thing",
+    )
+
+    result = await fulltextFilterThresh.all()
+
+    print(result)
+    assert len(result) == 1
+    assert all(isinstance(x[0], fulltextNodeThresh) for x in result)
+    assert result[0][0].other == "thing"
+    assert all(x[1] >= 0.09 for x in result)
+
 @mark_async_test
 async def test_dont_duplicate_fulltext_filter_node():
     """

test/async_/test_vectorfilter.py

@@ -47,6 +47,38 @@ class someNode(AsyncStructuredNode):
     assert all(isinstance(x[0], someNode) for x in result)
     assert all(isinstance(x[1], float) for x in result)
 
+@mark_async_test
+async def test_vectorfilter_thresholding():
+    """
+    Tests that the vector query is run, and only node above threshold returns.
+    """
+    # Vector Indexes only exist from 5.13 onwards
+    if not await adb.version_is_higher_than("5.13"):
+        pytest.skip("Vector Index not Generally Available in Neo4j.")
+
+    class someNodeThresh(AsyncStructuredNode):
+        name = StringProperty()
+        vector = ArrayProperty(
+            base_property=FloatProperty(), vector_index=VectorIndex(2, "cosine")
+        )
+
+    await adb.install_labels(someNodeThresh)
+
+    john = await someNodeThresh(name="John", vector=[float(0.5), float(0.5)]).save()
+    fred = await someNodeThresh(name="Fred", vector=[float(1.0), float(0.0)]).save()
+
+    vectorsearchFilterThreshold = someNodeThresh.nodes.filter(
+        vector_filter=VectorFilter(
+            topk=3, vector_attribute_name="vector", candidate_vector=[0.25, 0], threshold=0.8
+        ),
+        name="John",
+    )
+    result = await vectorsearchFilterThreshold.all()
+
+    assert len(result) == 1
+    assert all(isinstance(x[0], someNodeThresh) for x in result)
+    assert result[0][0].name == "John"
+    assert all(x[1] >= 0.8 for x in result)
 
 @mark_async_test
 async def test_vectorfilter_with_node_propertyfilter():

test/sync_/test_fulltextfilter.py

@@ -122,6 +122,50 @@ class fulltextNodeBis(StructuredNode):
     assert all(isinstance(x[1], float) for x in result)
 
 
+@mark_sync_test
+def test_fulltextfilter_threshold():
+    """
+    Tests that the fulltext query is run, and only nodes above threshold returns.
+    """
+
+    if not db.version_is_higher_than("5.16"):
+        pytest.skip("Not supported before 5.16")
+
+    class fulltextNodeThresh(StructuredNode):
+        description = StringProperty(
+            fulltext_index=FulltextIndex(
+                analyzer="standard-no-stop-words", eventually_consistent=False
+            )
+        )
+        other = StringProperty()
+
+    db.install_labels(fulltextNodeThresh)
+
+    node1 = fulltextNodeThresh(other="thing", description="Another thing").save()
+
+    node2 = fulltextNodeThresh(
+        other="other thing", description="Another other thing"
+    ).save()
+
+    fulltextFilterThresh = fulltextNodeThresh.nodes.filter(
+        fulltext_filter=FulltextFilter(
+            topk=3,
+            fulltext_attribute_name="description",
+            query_string="thing",
+            threshold=0.09,
+        ),
+        other="thing",
+    )
+
+    result = fulltextFilterThresh.all()
+
+    print(result)
+    assert len(result) == 1
+    assert all(isinstance(x[0], fulltextNodeThresh) for x in result)
+    assert result[0][0].other == "thing"
+    assert all(x[1] >= 0.09 for x in result)
+
+
 @mark_sync_test
 def test_dont_duplicate_fulltext_filter_node():
     """

test/sync_/test_indexing.py

@@ -3,7 +3,13 @@
 import pytest
 from pytest import raises
 
-from neomodel import IntegerProperty, StringProperty, StructuredNode, UniqueProperty, db
+from neomodel import (
+    IntegerProperty,
+    StringProperty,
+    StructuredNode,
+    UniqueProperty,
+    db,
+)
 from neomodel.exceptions import ConstraintValidationFailed