Files

454 lines
18 KiB
Python
Raw Permalink Normal View History

2026-05-12 19:40:31 +09:00
# Copyright (c) "Neo4j"
# Neo4j Sweden AB [https://neo4j.com]
# #
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# #
# https://www.apache.org/licenses/LICENSE-2.0
# #
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Any
from unittest.mock import patch
import pytest
from neo4j_graphrag.exceptions import InvalidHybridSearchRankerError
from neo4j_graphrag.filters import FilterClassification
from neo4j_graphrag.neo4j_queries import (
_build_search_clause_vector_query,
_build_hybrid_search_clause_query,
_build_hybrid_search_clause_query_linear,
get_query_tail,
get_search_query,
_get_hybrid_query_linear,
)
from neo4j_graphrag.types import EntityType, SearchType
def test_vector_search_basic() -> None:
expected = (
"CALL db.index.vector.queryNodes($vector_index_name, $top_k * $effective_search_ratio, $query_vector) "
"YIELD node, score "
"WITH node, score LIMIT $top_k "
"RETURN node { .*, `None`: null } AS node, labels(node) AS nodeLabels, elementId(node) AS elementId, elementId(node) AS id, score"
)
result, params = get_search_query(SearchType.VECTOR)
assert result.strip() == expected.strip()
assert params == {}
def test_rel_vector_search_basic() -> None:
expected = (
"CALL db.index.vector.queryRelationships($vector_index_name, $top_k * $effective_search_ratio, $query_vector) "
"YIELD relationship, score "
"WITH relationship, score LIMIT $top_k "
"RETURN relationship { .*, `None`: null } AS relationship, type(relationship) as relationshipType, "
"elementId(relationship) AS elementId, elementId(relationship) AS id, score"
)
result, params = get_search_query(SearchType.VECTOR, EntityType.RELATIONSHIP)
assert result.strip() == expected.strip()
assert params == {}
def test_rel_vector_search_filters_err() -> None:
with pytest.raises(Exception) as exc_info:
get_search_query(
SearchType.VECTOR, EntityType.RELATIONSHIP, filters={"filter": None}
)
assert str(exc_info.value) == "Filters are not supported for relationship indexes"
def test_rel_vector_search_hybrid_err() -> None:
with pytest.raises(Exception) as exc_info:
get_search_query(SearchType.HYBRID, EntityType.RELATIONSHIP)
assert (
str(exc_info.value) == "Hybrid search is not supported for relationship indexes"
)
def test_hybrid_search_basic() -> None:
expected = (
"CALL { "
"CALL db.index.vector.queryNodes($vector_index_name, $top_k * $effective_search_ratio, $query_vector) "
"YIELD node, score "
"WITH node, score LIMIT $top_k "
"WITH collect({node:node, score:score}) AS nodes, max(score) AS vector_index_max_score "
"UNWIND nodes AS n "
"RETURN n.node AS node, (n.score / vector_index_max_score) 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 ft_index_max_score "
"UNWIND nodes AS n "
"RETURN n.node AS node, (n.score / ft_index_max_score) AS score "
"} "
"WITH node, max(score) AS score ORDER BY score DESC, elementId(node) LIMIT $top_k "
"RETURN node { .*, `None`: null } AS node, labels(node) AS nodeLabels, elementId(node) AS elementId, elementId(node) AS id, score"
)
result, _ = get_search_query(SearchType.HYBRID)
assert result.strip() == expected.strip()
def test_vector_search_with_properties() -> None:
properties = ["name", "age"]
expected = (
"CALL db.index.vector.queryNodes($vector_index_name, $top_k * $effective_search_ratio, $query_vector) "
"YIELD node, score "
"WITH node, score LIMIT $top_k "
"RETURN node {.name, .age} AS node, labels(node) AS nodeLabels, elementId(node) AS elementId, elementId(node) AS id, score"
)
result, _ = get_search_query(SearchType.VECTOR, return_properties=properties)
assert result.strip() == expected.strip()
def test_vector_search_with_retrieval_query() -> None:
retrieval_query = "MATCH (n) RETURN n LIMIT 10"
expected = (
"CALL db.index.vector.queryNodes($vector_index_name, $top_k * $effective_search_ratio, $query_vector) "
"YIELD node, score "
"WITH node, score LIMIT $top_k " + retrieval_query
)
result, _ = get_search_query(SearchType.VECTOR, retrieval_query=retrieval_query)
assert result.strip() == expected.strip()
@patch("neo4j_graphrag.neo4j_queries.get_metadata_filter", return_value=["True", {}])
def test_vector_search_with_filters(_mock: Any) -> None:
expected = (
"MATCH (node:`Label`) "
"WHERE node.`vector` IS NOT NULL "
"AND size(node.`vector`) = toInteger($embedding_dimension)"
" AND (True) "
"WITH node, "
"vector.similarity.cosine(node.`vector`, $query_vector) AS score "
"ORDER BY score DESC, elementId(node) LIMIT $top_k "
"RETURN node { .*, `vector`: null } AS node, labels(node) AS nodeLabels, elementId(node) AS elementId, elementId(node) AS id, score"
)
result, params = get_search_query(
SearchType.VECTOR,
node_label="Label",
embedding_node_property="vector",
embedding_dimension=1,
filters={"field": "value"},
)
assert result.strip() == expected.strip()
assert params == {"embedding_dimension": 1}
@patch(
"neo4j_graphrag.neo4j_queries.get_metadata_filter",
return_value=["True", {"param": "value"}],
)
def test_vector_search_with_params_from_filters(_mock: Any) -> None:
expected = (
"MATCH (node:`Label`) "
"WHERE node.`vector` IS NOT NULL "
"AND size(node.`vector`) = toInteger($embedding_dimension)"
" AND (True) "
"WITH node, "
"vector.similarity.cosine(node.`vector`, $query_vector) AS score "
"ORDER BY score DESC, elementId(node) LIMIT $top_k "
"RETURN node { .*, `vector`: null } AS node, labels(node) AS nodeLabels, elementId(node) AS elementId, elementId(node) AS id, score"
)
result, params = get_search_query(
SearchType.VECTOR,
node_label="Label",
embedding_node_property="vector",
embedding_dimension=1,
filters={"field": "value"},
)
assert result.strip() == expected.strip()
assert params == {"embedding_dimension": 1, "param": "value"}
def test_hybrid_search_with_retrieval_query() -> None:
retrieval_query = "MATCH (n) RETURN n LIMIT 10"
expected = (
"CALL { "
"CALL db.index.vector.queryNodes($vector_index_name, $top_k * $effective_search_ratio, $query_vector) "
"YIELD node, score "
"WITH node, score LIMIT $top_k "
"WITH collect({node:node, score:score}) AS nodes, max(score) AS vector_index_max_score "
"UNWIND nodes AS n "
"RETURN n.node AS node, (n.score / vector_index_max_score) 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 ft_index_max_score "
"UNWIND nodes AS n "
"RETURN n.node AS node, (n.score / ft_index_max_score) AS score "
"} "
"WITH node, max(score) AS score ORDER BY score DESC, elementId(node) LIMIT $top_k "
+ retrieval_query
)
result, _ = get_search_query(SearchType.HYBRID, retrieval_query=retrieval_query)
assert result.strip() == expected.strip()
def test_hybrid_search_with_properties() -> None:
properties = ["name", "age"]
expected = (
"CALL { "
"CALL db.index.vector.queryNodes($vector_index_name, $top_k * $effective_search_ratio, $query_vector) "
"YIELD node, score "
"WITH node, score LIMIT $top_k "
"WITH collect({node:node, score:score}) AS nodes, max(score) AS vector_index_max_score "
"UNWIND nodes AS n "
"RETURN n.node AS node, (n.score / vector_index_max_score) 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 ft_index_max_score "
"UNWIND nodes AS n "
"RETURN n.node AS node, (n.score / ft_index_max_score) AS score "
"} "
"WITH node, max(score) AS score ORDER BY score DESC, elementId(node) LIMIT $top_k "
"RETURN node {.name, .age} AS node, labels(node) AS nodeLabels, elementId(node) AS elementId, elementId(node) AS id, score"
)
result, _ = get_search_query(SearchType.HYBRID, return_properties=properties)
assert result.strip() == expected.strip()
def test_get_query_tail_with_retrieval_query() -> None:
retrieval_query = "MATCH (n) RETURN n LIMIT 10"
expected = retrieval_query
result = get_query_tail(retrieval_query=retrieval_query)
assert result.strip() == expected.strip()
def test_get_query_tail_with_properties() -> None:
properties = ["name", "age"]
expected = "RETURN node {.name, .age} AS node, labels(node) AS nodeLabels, elementId(node) AS elementId, elementId(node) AS id, score"
result = get_query_tail(return_properties=properties)
assert result.strip() == expected.strip()
def test_get_query_tail_with_fallback() -> None:
fallback = "HELLO"
expected = fallback
result = get_query_tail(fallback_return=fallback)
assert result.strip() == expected.strip()
def test_get_query_tail_ordering_all() -> None:
retrieval_query = "MATCH (n) RETURN n LIMIT 10"
properties = ["name", "age"]
fallback = "HELLO"
expected = retrieval_query
result = get_query_tail(
retrieval_query=retrieval_query,
return_properties=properties,
fallback_return=fallback,
)
assert result.strip() == expected.strip()
def test_get_query_tail_ordering_no_retrieval_query() -> None:
properties = ["name", "age"]
fallback = "HELLO"
expected = "RETURN node {.name, .age} AS node, labels(node) AS nodeLabels, elementId(node) AS elementId, elementId(node) AS id, score"
result = get_query_tail(
return_properties=properties,
fallback_return=fallback,
)
assert result.strip() == expected.strip()
def test_get_hybrid_query_linear_with_alpha() -> None:
query = _get_hybrid_query_linear(neo4j_version_is_5_23_or_above=True, alpha=0.7)
vector_substr = "rawScore * $alpha"
ft_substr = "rawScore * (1 - $alpha)"
assert vector_substr in query
assert ft_substr in query
def test_invalid_hybrid_search_ranker_error() -> None:
with pytest.raises(InvalidHybridSearchRankerError):
get_search_query(SearchType.HYBRID, ranker="invalid")
class TestBuildSearchClauseVectorQuery:
def test_no_filters(self) -> None:
query, params = _build_search_clause_vector_query(
index_name="my_index",
node_label="Document",
)
expected = (
"MATCH (node:`Document`) "
"SEARCH node IN (VECTOR INDEX `my_index` "
"FOR $query_vector "
"LIMIT $top_k * $effective_search_ratio) "
"SCORE AS score "
"WITH node, score ORDER BY score DESC LIMIT $top_k"
)
assert query == expected
assert params == {}
def test_no_filters_with_none_classification(self) -> None:
query, params = _build_search_clause_vector_query(
index_name="my_index",
node_label="Document",
filter_classification=None,
)
assert "WHERE" not in query
assert params == {}
def test_empty_filter_classification(self) -> None:
classification = FilterClassification(is_compatible=True)
query, params = _build_search_clause_vector_query(
index_name="my_index",
node_label="Document",
filter_classification=classification,
)
assert "WHERE" not in query
assert params == {}
def test_with_where_predicates(self) -> None:
classification = FilterClassification(
is_compatible=True,
cypher_where_clause="node.`year` = $_e_year",
params={"_e_year": 2024},
)
query, params = _build_search_clause_vector_query(
index_name="my_index",
node_label="Document",
filter_classification=classification,
)
expected = (
"MATCH (node:`Document`) "
"SEARCH node IN (VECTOR INDEX `my_index` "
"FOR $query_vector "
"WHERE node.`year` = $_e_year "
"LIMIT $top_k * $effective_search_ratio) "
"SCORE AS score "
"WITH node, score ORDER BY score DESC LIMIT $top_k"
)
assert query == expected
assert params == {"_e_year": 2024}
def test_with_multiple_predicates(self) -> None:
classification = FilterClassification(
is_compatible=True,
cypher_where_clause="node.`year` >= $_e_year AND node.`status` = $_e_status",
params={"_e_year": 2020, "_e_status": "active"},
)
query, params = _build_search_clause_vector_query(
index_name="vector_idx",
node_label="Article",
filter_classification=classification,
)
assert "WHERE node.`year` >= $_e_year AND node.`status` = $_e_status" in query
assert params == {"_e_year": 2020, "_e_status": "active"}
def test_backtick_escaping_in_label_and_index(self) -> None:
query, _ = _build_search_clause_vector_query(
index_name="my-special-index",
node_label="My Label",
)
assert "`my-special-index`" in query
assert "`My Label`" in query
def test_query_structure_with_filters(self) -> None:
"""Verify MATCH ... SEARCH ... FOR ... WHERE ... LIMIT ... SCORE order."""
classification = FilterClassification(
is_compatible=True,
cypher_where_clause="node.`x` > $_e_x",
params={"_e_x": 5},
)
query, _ = _build_search_clause_vector_query(
index_name="idx",
node_label="N",
filter_classification=classification,
)
# Verify ordering of clauses
match_pos = query.index("MATCH")
search_pos = query.index("SEARCH")
for_pos = query.index("FOR")
where_pos = query.index("WHERE")
limit_pos = query.index("LIMIT")
score_pos = query.index("SCORE")
assert match_pos < search_pos < for_pos < where_pos < limit_pos < score_pos
def test_filter_params_empty_dict(self) -> None:
classification = FilterClassification(
is_compatible=True,
cypher_where_clause="node.`a` = $_e_a",
params={},
)
_, params = _build_search_clause_vector_query(
index_name="idx",
node_label="N",
filter_classification=classification,
)
assert params == {}
class TestBuildHybridSearchClauseQuery:
def test_basic_structure(self) -> None:
query = _build_hybrid_search_clause_query(
vector_index_name="vec_idx",
fulltext_index_name="ft_idx",
node_label="Document",
)
assert "CALL () {" in query
assert "UNION" in query
assert "VECTOR INDEX `vec_idx`" in query
assert "db.index.fulltext.queryNodes" in query
assert "`Document`" in query
assert "max(score) AS score" in query
assert "ORDER BY score DESC" in query
def test_vector_uses_search_fulltext_uses_procedure(self) -> None:
query = _build_hybrid_search_clause_query(
vector_index_name="v",
fulltext_index_name="f",
node_label="N",
)
assert "SEARCH node IN (VECTOR INDEX" in query
assert "FULLTEXT INDEX" not in query
assert "db.index.fulltext.queryNodes" in query
assert "$query_vector" in query
assert "$query_text" in query
assert "vector_index_max_score" in query
assert "ft_index_max_score" in query
def test_backtick_escaping(self) -> None:
query = _build_hybrid_search_clause_query(
vector_index_name="my-vec",
fulltext_index_name="my-ft",
node_label="My Label",
)
assert "`my-vec`" in query
assert "`My Label`" in query
class TestBuildHybridSearchClauseQueryLinear:
def test_basic_structure(self) -> None:
query = _build_hybrid_search_clause_query_linear(
vector_index_name="vec_idx",
fulltext_index_name="ft_idx",
node_label="Document",
)
assert "CALL () {" in query
assert "UNION" in query
assert "$alpha" in query
assert "(1 - $alpha)" in query
assert "sum(score)" in query
assert "ORDER BY score DESC" in query
def test_vector_uses_search_fulltext_uses_procedure(self) -> None:
query = _build_hybrid_search_clause_query_linear(
vector_index_name="v",
fulltext_index_name="f",
node_label="N",
)
assert "SEARCH node IN (VECTOR INDEX" in query
assert "FULLTEXT INDEX" not in query
assert "db.index.fulltext.queryNodes" in query
assert "$query_vector" in query
assert "$query_text" in query