Skip to content

Commit 4fe39b5

Browse files
committed
Refactor Cypher query string for linear ranker
1 parent cbfae5e commit 4fe39b5

File tree

2 files changed

+7
-12
lines changed

2 files changed

+7
-12
lines changed

src/neo4j_graphrag/neo4j_queries.py

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -198,20 +198,15 @@ def _get_hybrid_query_linear(neo4j_version_is_5_23_or_above: bool, alpha: float)
198198
f"{NODE_VECTOR_INDEX_QUERY} "
199199
"WITH collect({node: node, score: score}) AS nodes, max(score) AS vector_index_max_score "
200200
"UNWIND nodes AS n "
201-
"RETURN n.node AS node, (n.score / vector_index_max_score) AS score, 'vector' AS source "
201+
"WITH n.node AS node, (n.score / vector_index_max_score) AS rawScore "
202+
"RETURN node, rawScore * $alpha AS score "
202203
"UNION "
203204
f"{FULL_TEXT_SEARCH_QUERY} "
204205
"WITH collect({node: node, score: score}) AS nodes, max(score) AS ft_index_max_score "
205206
"UNWIND nodes AS n "
206-
"RETURN n.node AS node, (n.score / ft_index_max_score) AS score, 'ft' AS source } "
207-
"WITH node, "
208-
"sum(CASE WHEN source = 'vector' THEN score * "
209-
+ "$alpha "
210-
+ " ELSE 0 END) + "
211-
+ "sum(CASE WHEN source = 'ft' THEN score * "
212-
+ "(1 - $alpha) "
213-
+ " ELSE 0 END) AS score "
214-
"ORDER BY score DESC LIMIT $top_k"
207+
"WITH n.node AS node, (n.score / ft_index_max_score) AS rawScore "
208+
"RETURN node, rawScore * (1 - $alpha) AS score } "
209+
"WITH node, sum(score) AS score ORDER BY score DESC LIMIT $top_k"
215210
)
216211
return call_prefix + query_body
217212

tests/unit/test_neo4j_queries.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -256,8 +256,8 @@ def test_get_query_tail_ordering_no_retrieval_query() -> None:
256256

257257
def test_get_hybrid_query_linear_with_alpha() -> None:
258258
query = _get_hybrid_query_linear(neo4j_version_is_5_23_or_above=True, alpha=0.7)
259-
vector_substr = "CASE WHEN source = 'vector' THEN score * $alpha"
260-
ft_substr = "CASE WHEN source = 'ft' THEN score * (1 - $alpha)"
259+
vector_substr = "rawScore * $alpha"
260+
ft_substr = "rawScore * (1 - $alpha)"
261261
assert vector_substr in query
262262
assert ft_substr in query
263263

0 commit comments

Comments
 (0)