Skip to content

Commit 7ae97dd

Browse files
stellasiastellasia
authored
Use routing_=READ for all read queries (#217)
* Use routing_=neo4j.RoutingControl.READ for all READ queries * Update CHANGELOG
1 parent 99315c4 commit 7ae97dd

File tree

17 files changed

+70
-12
lines changed

17 files changed

+70
-12
lines changed

‎CHANGELOG.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,8 @@
1212

1313
### Changed
1414
- Updated all examples to use `neo4j_database` parameter instead of an undocumented neo4j driver constructor.
15+
- All `READ` queries are now routed to a reader replica (for clusters). This impacts all retrievers, the `Neo4jChunkReader` and `SinglePropertyExactMatchResolver` components.
16+
1517

1618
## 1.2.0
1719

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,7 @@ async def run(
9797
result, _, _ = self.driver.execute_query(
9898
query,
9999
database_=self.neo4j_database,
100+
routing_=neo4j.RoutingControl.READ,
100101
)
101102
chunks = []
102103
for record in result:

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,8 @@ async def run(self) -> ResolutionStats:
9494
match_query += self.filter_query
9595
stat_query = f"{match_query} RETURN count(entity) as c"
9696
records, _, _ = self.driver.execute_query(
97-
stat_query, database_=self.neo4j_database
97+
stat_query,
98+
database_=self.neo4j_database,
9899
)
99100
number_of_nodes_to_resolve = records[0].get("c")
100101
if number_of_nodes_to_resolve == 0:

‎src/neo4j_graphrag/retrievers/base.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,9 @@ def __init__(self, driver: neo4j.Driver, neo4j_database: Optional[str] = None):
8989

9090
def _get_version(self) -> tuple[tuple[int, ...], bool]:
9191
records, _, _ = self.driver.execute_query(
92-
"CALL dbms.components()", database_=self.neo4j_database
92+
"CALL dbms.components()",
93+
database_=self.neo4j_database,
94+
routing_=neo4j.RoutingControl.READ,
9395
)
9496
version = records[0]["versions"][0]
9597
# drop everything after the '-' first
@@ -145,7 +147,10 @@ def _fetch_index_infos(self, vector_index_name: str) -> None:
145147
"options.indexConfig.`vector.dimensions` as dimensions"
146148
)
147149
query_result = self.driver.execute_query(
148-
query, {"index_name": vector_index_name}, database_=self.neo4j_database
150+
query,
151+
{"index_name": vector_index_name},
152+
database_=self.neo4j_database,
153+
routing_=neo4j.RoutingControl.READ,
149154
)
150155
try:
151156
result = query_result.records[0]

‎src/neo4j_graphrag/retrievers/external/pinecone/pinecone.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -234,7 +234,10 @@ def get_search_results(
234234
logger.debug("Pinecone Store Cypher query: %s", search_query)
235235

236236
records, _, _ = self.driver.execute_query(
237-
search_query, parameters, database_=self.neo4j_database
237+
search_query,
238+
parameters,
239+
database_=self.neo4j_database,
240+
routing_=neo4j.RoutingControl.READ,
238241
)
239242

240243
return RawSearchResult(records=records)

‎src/neo4j_graphrag/retrievers/external/qdrant/qdrant.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -228,7 +228,10 @@ def get_search_results(
228228
logger.debug("Qdrant Store Cypher query: %s", search_query)
229229

230230
records, _, _ = self.driver.execute_query(
231-
search_query, parameters, database_=self.neo4j_database
231+
search_query,
232+
parameters,
233+
database_=self.neo4j_database,
234+
routing_=neo4j.RoutingControl.READ,
232235
)
233236

234237
return RawSearchResult(records=records)

‎src/neo4j_graphrag/retrievers/external/weaviate/weaviate.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -245,7 +245,10 @@ def get_search_results(
245245
logger.debug("Weaviate Store Cypher query: %s", search_query)
246246

247247
records, _, _ = self.driver.execute_query(
248-
search_query, parameters, database_=self.neo4j_database
248+
search_query,
249+
parameters,
250+
database_=self.neo4j_database,
251+
routing_=neo4j.RoutingControl.READ,
249252
)
250253

251254
return RawSearchResult(records=records)

‎src/neo4j_graphrag/retrievers/hybrid.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -201,7 +201,10 @@ def get_search_results(
201201
logger.debug("HybridRetriever Cypher query: %s", search_query)
202202

203203
records, _, _ = self.driver.execute_query(
204-
search_query, parameters, database_=self.neo4j_database
204+
search_query,
205+
parameters,
206+
database_=self.neo4j_database,
207+
routing_=neo4j.RoutingControl.READ,
205208
)
206209
return RawSearchResult(
207210
records=records,
@@ -358,7 +361,10 @@ def get_search_results(
358361
logger.debug("HybridRetriever Cypher query: %s", search_query)
359362

360363
records, _, _ = self.driver.execute_query(
361-
search_query, parameters, database_=self.neo4j_database
364+
search_query,
365+
parameters,
366+
database_=self.neo4j_database,
367+
routing_=neo4j.RoutingControl.READ,
362368
)
363369
return RawSearchResult(
364370
records=records,

‎src/neo4j_graphrag/retrievers/text2cypher.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -167,7 +167,9 @@ def get_search_results(
167167
t2c_query = llm_result.content
168168
logger.debug("Text2CypherRetriever Cypher query: %s", t2c_query)
169169
records, _, _ = self.driver.execute_query(
170-
query_=t2c_query, database_=self.neo4j_database
170+
query_=t2c_query,
171+
database_=self.neo4j_database,
172+
routing_=neo4j.RoutingControl.READ,
171173
)
172174
except CypherSyntaxError as e:
173175
raise Text2CypherRetrievalError(

‎src/neo4j_graphrag/retrievers/vector.py

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -207,7 +207,10 @@ def get_search_results(
207207
logger.debug("VectorRetriever Cypher query: %s", search_query)
208208

209209
records, _, _ = self.driver.execute_query(
210-
search_query, parameters, database_=self.neo4j_database
210+
search_query,
211+
parameters,
212+
database_=self.neo4j_database,
213+
routing_=neo4j.RoutingControl.READ,
211214
)
212215
return RawSearchResult(records=records)
213216

@@ -363,7 +366,10 @@ def get_search_results(
363366
logger.debug("VectorCypherRetriever Cypher query: %s", search_query)
364367

365368
records, _, _ = self.driver.execute_query(
366-
search_query, parameters, database_=self.neo4j_database
369+
search_query,
370+
parameters,
371+
database_=self.neo4j_database,
372+
routing_=neo4j.RoutingControl.READ,
367373
)
368374
return RawSearchResult(
369375
records=records,

‎tests/unit/experimental/components/test_neo4j_reader.py

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@ async def test_neo4j_chunk_reader(driver: Mock) -> None:
3535
driver.execute_query.assert_called_once_with(
3636
"MATCH (c:`Chunk`) RETURN c { .*, embedding: null } as chunk ORDER BY c.index",
3737
database_="mydb",
38+
routing_=neo4j.RoutingControl.READ,
3839
)
3940

4041
assert isinstance(res, TextChunks)
@@ -75,6 +76,7 @@ async def test_neo4j_chunk_reader_custom_lg_config(driver: Mock) -> None:
7576
driver.execute_query.assert_called_once_with(
7677
"MATCH (c:`Page`) RETURN c { .*, embedding: null } as chunk ORDER BY c.k",
7778
database_=None,
79+
routing_=neo4j.RoutingControl.READ,
7880
)
7981

8082
assert isinstance(res, TextChunks)
@@ -110,6 +112,7 @@ async def test_neo4j_chunk_reader_fetch_embedding(driver: Mock) -> None:
110112
driver.execute_query.assert_called_once_with(
111113
"MATCH (c:`Chunk`) RETURN c { .* } as chunk ORDER BY c.index",
112114
database_=None,
115+
routing_=neo4j.RoutingControl.READ,
113116
)
114117

115118
assert isinstance(res, TextChunks)

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,7 @@ def test_pinecone_retriever_search_happy_path(
9999
"id_property": "sync_id",
100100
},
101101
database_=None,
102+
routing_=neo4j.RoutingControl.READ,
102103
)
103104

104105
assert records == RetrieverResult(
@@ -171,6 +172,7 @@ def test_pinecone_retriever_search_return_properties(
171172
"id_property": "sync_id",
172173
},
173174
database_=None,
175+
routing_=neo4j.RoutingControl.READ,
174176
)
175177

176178
assert records == RetrieverResult(
@@ -230,6 +232,7 @@ def test_pinecone_retriever_search_retrieval_query(
230232
"id_property": "sync_id",
231233
},
232234
database_=None,
235+
routing_=neo4j.RoutingControl.READ,
233236
)
234237

235238
assert records == RetrieverResult(

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,7 @@ def test_qdrant_retriever_search_happy_path(
7474
"id_property": "sync_id",
7575
},
7676
database_=None,
77+
routing_=neo4j.RoutingControl.READ,
7778
)
7879

7980
assert records == RetrieverResult(
@@ -152,6 +153,7 @@ def test_qdrant_retriever_search_return_properties(
152153
"id_property": "sync_id",
153154
},
154155
database_=None,
156+
routing_=neo4j.RoutingControl.READ,
155157
)
156158

157159
assert records == RetrieverResult(
@@ -217,6 +219,7 @@ def test_qdrant_retriever_search_retrieval_query(
217219
"id_property": "sync_id",
218220
},
219221
database_=None,
222+
routing_=neo4j.RoutingControl.READ,
220223
)
221224

222225
assert records == RetrieverResult(

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,7 @@ def test_text_search_remote_vector_store_happy_path(driver: MagicMock) -> None:
8080
"id_property": "sync_id",
8181
},
8282
database_=None,
83+
routing_=neo4j.RoutingControl.READ,
8384
)
8485
assert records == RetrieverResult(
8586
items=[
@@ -146,6 +147,7 @@ def test_text_search_remote_vector_store_return_properties(driver: MagicMock) ->
146147
"id_property": "sync_id",
147148
},
148149
database_=None,
150+
routing_=neo4j.RoutingControl.READ,
149151
)
150152
assert records == RetrieverResult(
151153
items=[
@@ -193,6 +195,7 @@ def test_text_search_remote_vector_store_retrieval_query(driver: MagicMock) -> N
193195
"id_property": "sync_id",
194196
},
195197
database_=None,
198+
routing_=neo4j.RoutingControl.READ,
196199
)
197200

198201
assert records == RetrieverResult(

‎tests/unit/retrievers/test_hybrid.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515

1616
from unittest.mock import MagicMock, patch
1717

18+
import neo4j
1819
import pytest
1920
from neo4j_graphrag.exceptions import (
2021
EmbeddingRequiredError,
@@ -204,6 +205,7 @@ def test_hybrid_search_text_happy_path(
204205
"query_vector": embed_query_vector,
205206
},
206207
database_=None,
208+
routing_=neo4j.RoutingControl.READ,
207209
)
208210
embedder.embed_query.assert_called_once_with(query_text)
209211
assert records == RetrieverResult(
@@ -262,6 +264,7 @@ def test_hybrid_search_favors_query_vector_over_embedding_vector(
262264
"query_vector": query_vector,
263265
},
264266
database_=database,
267+
routing_=neo4j.RoutingControl.READ,
265268
)
266269
embedder.embed_query.assert_not_called()
267270

@@ -344,6 +347,7 @@ def test_hybrid_retriever_return_properties(
344347
"query_vector": embed_query_vector,
345348
},
346349
database_=None,
350+
routing_=neo4j.RoutingControl.READ,
347351
)
348352
assert records == RetrieverResult(
349353
items=[
@@ -410,6 +414,7 @@ def test_hybrid_cypher_retrieval_query_with_params(
410414
"param": "dummy-param",
411415
},
412416
database_=None,
417+
routing_=neo4j.RoutingControl.READ,
413418
)
414419

415420
assert records == RetrieverResult(

‎tests/unit/retrievers/test_text2cypher.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
# limitations under the License.
1515
from unittest.mock import MagicMock, patch
1616

17+
import neo4j
1718
import pytest
1819
from neo4j.exceptions import CypherSyntaxError, Neo4jError
1920
from neo4j_graphrag.exceptions import (
@@ -139,7 +140,9 @@ def test_t2c_retriever_happy_path(
139140
retriever.search(query_text=query_text)
140141
llm.invoke.assert_called_once_with(prompt)
141142
driver.execute_query.assert_called_once_with(
142-
query_=t2c_query, database_=neo4j_database
143+
query_=t2c_query,
144+
database_=neo4j_database,
145+
routing_=neo4j.RoutingControl.READ,
143146
)
144147

145148

‎tests/unit/retrievers/test_vector.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,7 @@ def test_similarity_search_vector_happy_path(
137137
"query_vector": query_vector,
138138
},
139139
database_=database,
140+
routing_=neo4j.RoutingControl.READ,
140141
)
141142
assert records == RetrieverResult(
142143
items=[
@@ -182,6 +183,7 @@ def test_similarity_search_text_happy_path(
182183
"query_vector": embed_query_vector,
183184
},
184185
database_=None,
186+
routing_=neo4j.RoutingControl.READ,
185187
)
186188
assert records == RetrieverResult(
187189
items=[
@@ -234,6 +236,7 @@ def test_similarity_search_text_return_properties(
234236
"query_vector": embed_query_vector,
235237
},
236238
database_=None,
239+
routing_=neo4j.RoutingControl.READ,
237240
)
238241
assert records == RetrieverResult(
239242
items=[
@@ -397,6 +400,7 @@ def test_retrieval_query_happy_path(
397400
"query_vector": embed_query_vector,
398401
},
399402
database_=database,
403+
routing_=neo4j.RoutingControl.READ,
400404
)
401405
assert records == RetrieverResult(
402406
items=[
@@ -458,6 +462,7 @@ def test_retrieval_query_with_result_format_function(
458462
"query_vector": embed_query_vector,
459463
},
460464
database_=None,
465+
routing_=neo4j.RoutingControl.READ,
461466
)
462467
assert records == RetrieverResult(
463468
items=[
@@ -520,6 +525,7 @@ def test_retrieval_query_with_params(
520525
"param": "dummy-param",
521526
},
522527
database_=None,
528+
routing_=neo4j.RoutingControl.READ,
523529
)
524530

525531
assert records == RetrieverResult(

0 commit comments

Comments
 (0)