fix test

langchain-ai · Dec 7, 2024 · 9266045 · 9266045
1 parent 5c06afb
commit 9266045
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/libs/neo4j/tests/unit_tests/vectorstores/test_neo4j.py b/libs/neo4j/tests/unit_tests/vectorstores/test_neo4j.py
@@ -203,8 +203,9 @@ def test_converting_to_yaml() -> None:
 def test_get_search_index_query_hybrid_node_neo4j_5_23_above() -> None:
     expected_query = (
         "CALL () { "
-        "CALL db.index.vector.queryNodes($index, $k, $embedding) "
+        "CALL db.index.vector.queryNodes($index, $k * $ef, $embedding) "
         "YIELD node, score "
+        "WITH node, score LIMIT $k "
         "WITH collect({node:node, score:score}) AS nodes, max(score) AS max "
         "UNWIND nodes AS n "
         "RETURN n.node AS node, (n.score / max) AS score UNION "
@@ -225,8 +226,9 @@ def test_get_search_index_query_hybrid_node_neo4j_5_23_above() -> None:
 def test_get_search_index_query_hybrid_node_neo4j_5_23_below() -> None:
     expected_query = (
         "CALL { "
-        "CALL db.index.vector.queryNodes($index, $k, $embedding) "
+        "CALL db.index.vector.queryNodes($index, $k * $ef, $embedding) "
         "YIELD node, score "
+        "WITH node, score LIMIT $k "
         "WITH collect({node:node, score:score}) AS nodes, max(score) AS max "
         "UNWIND nodes AS n "
         "RETURN n.node AS node, (n.score / max) AS score UNION "