Skip to content

Commit 0eb888f

Browse files
willtaiwilltai
authored
Refactor id variables in neo4j queries (#207)
* Refactor id variables in neo4j queries * Add warning in get_search_query * Add duplicate elementId(node) as nodeId * Rename to elementId * Rename node_id and rel_id to element id * Revert start_node_id and end_node_id * Rename to node_element_id and rel_element_id * Update docstring
1 parent 18e8e2a commit 0eb888f

File tree

6 files changed

+40
-26
lines changed

6 files changed

+40
-26
lines changed

‎src/neo4j_graphrag/experimental/components/types.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ class Neo4jNode(BaseModel):
4949
"""Represents a Neo4j node.
5050
5151
Attributes:
52-
id (str): The ID of the node.
52+
id (str): The element ID of the node.
5353
label (str): The label of the node.
5454
properties (dict[str, Any]): A dictionary of properties attached to the node.
5555
embedding_properties (Optional[dict[str, list[float]]]): A list of embedding properties attached to the node.

‎src/neo4j_graphrag/indexes.py

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -278,7 +278,7 @@ def upsert_vector(
278278
279279
Args:
280280
driver (neo4j.Driver): Neo4j Python driver instance.
281-
node_id (int): The id of the node.
281+
node_id (int): The element id of the node.
282282
embedding_property (str): The name of the property to store the vector in.
283283
vector (list[float]): The vector to store.
284284
neo4j_database (Optional[str]): The name of the Neo4j database. If not provided, this defaults to "neo4j" in the database (`see reference to documentation <https://neo4j.com/docs/operations-manual/current/database-administration/#manage-databases-default>`_).
@@ -288,7 +288,7 @@ def upsert_vector(
288288
"""
289289
try:
290290
parameters = {
291-
"id": node_id,
291+
"node_element_id": node_id,
292292
"embedding_property": embedding_property,
293293
"vector": vector,
294294
}
@@ -334,7 +334,7 @@ def upsert_vector_on_relationship(
334334
335335
Args:
336336
driver (neo4j.Driver): Neo4j Python driver instance.
337-
rel_id (int): The id of the relationship.
337+
rel_id (int): The element id of the relationship.
338338
embedding_property (str): The name of the property to store the vector in.
339339
vector (list[float]): The vector to store.
340340
neo4j_database (Optional[str]): The name of the Neo4j database. If not provided, this defaults to "neo4j" in the database (`see reference to documentation <https://neo4j.com/docs/operations-manual/current/database-administration/#manage-databases-default>`_).
@@ -344,7 +344,7 @@ def upsert_vector_on_relationship(
344344
"""
345345
try:
346346
parameters = {
347-
"id": rel_id,
347+
"rel_element_id": rel_id,
348348
"embedding_property": embedding_property,
349349
"vector": vector,
350350
}
@@ -391,7 +391,7 @@ async def async_upsert_vector(
391391
392392
Args:
393393
driver (neo4j.AsyncDriver): Neo4j Python asynchronous driver instance.
394-
node_id (int): The id of the node.
394+
node_id (int): The element id of the node.
395395
embedding_property (str): The name of the property to store the vector in.
396396
vector (list[float]): The vector to store.
397397
neo4j_database (Optional[str]): The name of the Neo4j database. If not provided, this defaults to "neo4j" in the database (`see reference to documentation <https://neo4j.com/docs/operations-manual/current/database-administration/#manage-databases-default>`_).
@@ -401,7 +401,7 @@ async def async_upsert_vector(
401401
"""
402402
try:
403403
parameters = {
404-
"id": node_id,
404+
"node_id": node_id,
405405
"embedding_property": embedding_property,
406406
"vector": vector,
407407
}
@@ -448,7 +448,7 @@ async def async_upsert_vector_on_relationship(
448448
449449
Args:
450450
driver (neo4j.AsyncDriver): Neo4j Python asynchronous driver instance.
451-
rel_id (int): The id of the relationship.
451+
rel_id (int): The element id of the relationship.
452452
embedding_property (str): The name of the property to store the vector in.
453453
vector (list[float]): The vector to store.
454454
neo4j_database (Optional[str]): The name of the Neo4j database. If not provided, this defaults to "neo4j" in the database (`see reference to documentation <https://neo4j.com/docs/operations-manual/current/database-administration/#manage-databases-default>`_).
@@ -458,7 +458,7 @@ async def async_upsert_vector_on_relationship(
458458
"""
459459
try:
460460
parameters = {
461-
"id": rel_id,
461+
"rel_id": rel_id,
462462
"embedding_property": embedding_property,
463463
"vector": vector,
464464
}

‎src/neo4j_graphrag/neo4j_queries.py

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
# limitations under the License.
1515
from __future__ import annotations
1616

17+
import warnings
1718
from typing import Any, Optional
1819

1920
from neo4j_graphrag.filters import get_metadata_filter
@@ -100,15 +101,15 @@
100101

101102
UPSERT_VECTOR_ON_NODE_QUERY = (
102103
"MATCH (n) "
103-
"WHERE elementId(n) = $id "
104+
"WHERE elementId(n) = $node_element_id "
104105
"WITH n "
105106
"CALL db.create.setNodeVectorProperty(n, $embedding_property, $vector) "
106107
"RETURN n"
107108
)
108109

109110
UPSERT_VECTOR_ON_RELATIONSHIP_QUERY = (
110111
"MATCH ()-[r]->() "
111-
"WHERE elementId(r) = $id "
112+
"WHERE elementId(r) = $rel_element_id "
112113
"WITH r "
113114
"CALL db.create.setRelationshipVectorProperty(r, $embedding_property, $vector) "
114115
"RETURN r"
@@ -201,6 +202,11 @@ def get_search_query(
201202
tuple[str, dict[str, Any]]: query and parameters
202203
203204
"""
205+
warnings.warn(
206+
"The default returned 'id' field in the search results will be removed. Please switch to using 'elementId' instead.",
207+
DeprecationWarning,
208+
stacklevel=2,
209+
)
204210
if search_type == SearchType.HYBRID:
205211
if filters:
206212
raise Exception("Filters are not supported with Hybrid Search")
@@ -227,7 +233,7 @@ def get_search_query(
227233
query_tail = get_query_tail(
228234
retrieval_query,
229235
return_properties,
230-
fallback_return=f"RETURN node {{ .*, `{embedding_node_property}`: null }} AS node, labels(node) AS nodeLabels, elementId(node) AS id, score",
236+
fallback_return=f"RETURN node {{ .*, `{embedding_node_property}`: null }} AS node, labels(node) AS nodeLabels, elementId(node) AS elementId, elementId(node) AS id, score",
231237
)
232238
return f"{query} {query_tail}", params
233239

@@ -253,5 +259,5 @@ def get_query_tail(
253259
return retrieval_query
254260
if return_properties:
255261
return_properties_cypher = ", ".join([f".{prop}" for prop in return_properties])
256-
return f"RETURN node {{{return_properties_cypher}}} AS node, labels(node) AS nodeLabels, elementId(node) AS id, score"
262+
return f"RETURN node {{{return_properties_cypher}}} AS node, labels(node) AS nodeLabels, elementId(node) AS elementId, elementId(node) AS id, score"
257263
return fallback_return if fallback_return else ""

‎tests/unit/retrievers/external/test_weaviate.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -225,7 +225,7 @@ def test_match_query_with_return_properties() -> None:
225225
"WITH match_param[0] AS match_id_value, match_param[1] AS score "
226226
"MATCH (node) "
227227
"WHERE node[$id_property] = match_id_value "
228-
"RETURN node {.name, .age} AS node, labels(node) AS nodeLabels, elementId(node) AS id, score"
228+
"RETURN node {.name, .age} AS node, labels(node) AS nodeLabels, elementId(node) AS elementId, elementId(node) AS id, score"
229229
)
230230
assert match_query.strip() == expected.strip()
231231

‎tests/unit/test_indexes.py

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -219,15 +219,19 @@ def test_upsert_vector_happy_path(driver: MagicMock) -> None:
219219

220220
upsert_query = (
221221
"MATCH (n) "
222-
"WHERE elementId(n) = $id "
222+
"WHERE elementId(n) = $node_element_id "
223223
"WITH n "
224224
"CALL db.create.setNodeVectorProperty(n, $embedding_property, $vector) "
225225
"RETURN n"
226226
)
227227

228228
driver.execute_query.assert_called_once_with(
229229
upsert_query,
230-
{"id": id, "embedding_property": embedding_property, "vector": vector},
230+
{
231+
"node_element_id": id,
232+
"embedding_property": embedding_property,
233+
"vector": vector,
234+
},
231235
database_=None,
232236
)
233237

@@ -241,15 +245,19 @@ def test_upsert_vector_on_relationship_happy_path(driver: MagicMock) -> None:
241245

242246
upsert_query = (
243247
"MATCH ()-[r]->() "
244-
"WHERE elementId(r) = $id "
248+
"WHERE elementId(r) = $rel_element_id "
245249
"WITH r "
246250
"CALL db.create.setRelationshipVectorProperty(r, $embedding_property, $vector) "
247251
"RETURN r"
248252
)
249253

250254
driver.execute_query.assert_called_once_with(
251255
upsert_query,
252-
{"id": id, "embedding_property": embedding_property, "vector": vector},
256+
{
257+
"rel_element_id": id,
258+
"embedding_property": embedding_property,
259+
"vector": vector,
260+
},
253261
database_=None,
254262
)
255263

‎tests/unit/test_neo4j_queries.py

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ def test_vector_search_basic() -> None:
2323
expected = (
2424
"CALL db.index.vector.queryNodes($vector_index_name, $top_k, $query_vector) "
2525
"YIELD node, score "
26-
"RETURN node { .*, `None`: null } AS node, labels(node) AS nodeLabels, elementId(node) AS id, score"
26+
"RETURN node { .*, `None`: null } AS node, labels(node) AS nodeLabels, elementId(node) AS elementId, elementId(node) AS id, score"
2727
)
2828
result, params = get_search_query(SearchType.VECTOR)
2929
assert result.strip() == expected.strip()
@@ -45,7 +45,7 @@ def test_hybrid_search_basic() -> None:
4545
"RETURN n.node AS node, (n.score / ft_index_max_score) AS score "
4646
"} "
4747
"WITH node, max(score) AS score ORDER BY score DESC LIMIT $top_k "
48-
"RETURN node { .*, `None`: null } AS node, labels(node) AS nodeLabels, elementId(node) AS id, score"
48+
"RETURN node { .*, `None`: null } AS node, labels(node) AS nodeLabels, elementId(node) AS elementId, elementId(node) AS id, score"
4949
)
5050
result, _ = get_search_query(SearchType.HYBRID)
5151
assert result.strip() == expected.strip()
@@ -56,7 +56,7 @@ def test_vector_search_with_properties() -> None:
5656
expected = (
5757
"CALL db.index.vector.queryNodes($vector_index_name, $top_k, $query_vector) "
5858
"YIELD node, score "
59-
"RETURN node {.name, .age} AS node, labels(node) AS nodeLabels, elementId(node) AS id, score"
59+
"RETURN node {.name, .age} AS node, labels(node) AS nodeLabels, elementId(node) AS elementId, elementId(node) AS id, score"
6060
)
6161
result, _ = get_search_query(SearchType.VECTOR, return_properties=properties)
6262
assert result.strip() == expected.strip()
@@ -82,7 +82,7 @@ def test_vector_search_with_filters(_mock: Any) -> None:
8282
"WITH node, "
8383
"vector.similarity.cosine(node.`vector`, $query_vector) AS score "
8484
"ORDER BY score DESC LIMIT $top_k "
85-
"RETURN node { .*, `vector`: null } AS node, labels(node) AS nodeLabels, elementId(node) AS id, score"
85+
"RETURN node { .*, `vector`: null } AS node, labels(node) AS nodeLabels, elementId(node) AS elementId, elementId(node) AS id, score"
8686
)
8787
result, params = get_search_query(
8888
SearchType.VECTOR,
@@ -108,7 +108,7 @@ def test_vector_search_with_params_from_filters(_mock: Any) -> None:
108108
"WITH node, "
109109
"vector.similarity.cosine(node.`vector`, $query_vector) AS score "
110110
"ORDER BY score DESC LIMIT $top_k "
111-
"RETURN node { .*, `vector`: null } AS node, labels(node) AS nodeLabels, elementId(node) AS id, score"
111+
"RETURN node { .*, `vector`: null } AS node, labels(node) AS nodeLabels, elementId(node) AS elementId, elementId(node) AS id, score"
112112
)
113113
result, params = get_search_query(
114114
SearchType.VECTOR,
@@ -159,7 +159,7 @@ def test_hybrid_search_with_properties() -> None:
159159
"RETURN n.node AS node, (n.score / ft_index_max_score) AS score "
160160
"} "
161161
"WITH node, max(score) AS score ORDER BY score DESC LIMIT $top_k "
162-
"RETURN node {.name, .age} AS node, labels(node) AS nodeLabels, elementId(node) AS id, score"
162+
"RETURN node {.name, .age} AS node, labels(node) AS nodeLabels, elementId(node) AS elementId, elementId(node) AS id, score"
163163
)
164164
result, _ = get_search_query(SearchType.HYBRID, return_properties=properties)
165165
assert result.strip() == expected.strip()
@@ -174,7 +174,7 @@ def test_get_query_tail_with_retrieval_query() -> None:
174174

175175
def test_get_query_tail_with_properties() -> None:
176176
properties = ["name", "age"]
177-
expected = "RETURN node {.name, .age} AS node, labels(node) AS nodeLabels, elementId(node) AS id, score"
177+
expected = "RETURN node {.name, .age} AS node, labels(node) AS nodeLabels, elementId(node) AS elementId, elementId(node) AS id, score"
178178
result = get_query_tail(return_properties=properties)
179179
assert result.strip() == expected.strip()
180180

@@ -204,7 +204,7 @@ def test_get_query_tail_ordering_no_retrieval_query() -> None:
204204
properties = ["name", "age"]
205205
fallback = "HELLO"
206206

207-
expected = "RETURN node {.name, .age} AS node, labels(node) AS nodeLabels, elementId(node) AS id, score"
207+
expected = "RETURN node {.name, .age} AS node, labels(node) AS nodeLabels, elementId(node) AS elementId, elementId(node) AS id, score"
208208
result = get_query_tail(
209209
return_properties=properties,
210210
fallback_return=fallback,

0 commit comments

Comments
 (0)