Fixes the hybrid retriever query

alexthomas93 · Aug 29, 2024 · 86c6567 · 86c6567
1 parent af27578
commit 86c6567
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/src/neo4j_genai/neo4j_queries.py b/src/neo4j_genai/neo4j_queries.py
@@ -53,7 +53,9 @@
 def _get_hybrid_query() -> str:
     return (
         f"CALL {{ {VECTOR_INDEX_QUERY} "
-        f"RETURN node, score "
+        f"WITH collect({{node:node, score:score}}) AS nodes, max(score) AS max "
+        f"UNWIND nodes AS n "
+        f"RETURN n.node AS node, (n.score / max) AS score "
         f"UNION "
         f"{FULL_TEXT_SEARCH_QUERY} "
         f"WITH collect({{node:node, score:score}}) AS nodes, max(score) AS max "

diff --git a/tests/unit/test_neo4j_queries.py b/tests/unit/test_neo4j_queries.py
@@ -35,7 +35,9 @@ def test_hybrid_search_basic() -> None:
         "CALL { "
         "CALL db.index.vector.queryNodes($vector_index_name, $top_k, $query_vector) "
         "YIELD node, score "
-        "RETURN node, score UNION "
+        "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 "
         "CALL db.index.fulltext.queryNodes($fulltext_index_name, $query_text, {limit: $top_k}) "
         "YIELD node, score "
         "WITH collect({node:node, score:score}) AS nodes, max(score) AS max "
@@ -125,7 +127,9 @@ def test_hybrid_search_with_retrieval_query() -> None:
         "CALL { "
         "CALL db.index.vector.queryNodes($vector_index_name, $top_k, $query_vector) "
         "YIELD node, score "
-        "RETURN node, score UNION "
+        "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 "
         "CALL db.index.fulltext.queryNodes($fulltext_index_name, $query_text, {limit: $top_k}) "
         "YIELD node, score "
         "WITH collect({node:node, score:score}) AS nodes, max(score) AS max "
@@ -145,7 +149,9 @@ def test_hybrid_search_with_properties() -> None:
         "CALL { "
         "CALL db.index.vector.queryNodes($vector_index_name, $top_k, $query_vector) "
         "YIELD node, score "
-        "RETURN node, score UNION "
+        "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 "
         "CALL db.index.fulltext.queryNodes($fulltext_index_name, $query_text, {limit: $top_k}) "
         "YIELD node, score "
         "WITH collect({node:node, score:score}) AS nodes, max(score) AS max "