feat(retrievers): add async retriever support via neo4j.AsyncDriver (#406) - #552
feat(retrievers): add async retriever support via neo4j.AsyncDriver (#406)#5521HazyOne707 wants to merge 1 commit into
Conversation
…eo4j#406) Adds AsyncRetriever base class and async variants of all core retrievers, allowing users to use neo4j.AsyncDriver without maintaining two driver instances or spawning threads. Changes: - types.py: Add AsyncNeo4jDriverModel (validates neo4j.AsyncDriver) and async model variants for all retrievers - retrievers/base.py: Add AsyncRetriever ABC with async search(), async get_search_results(), async _fetch_index_infos() - retrievers/async_vector.py: AsyncVectorRetriever, AsyncVectorCypherRetriever - retrievers/async_hybrid.py: AsyncHybridRetriever, AsyncHybridCypherRetriever - retrievers/async_text2cypher.py: AsyncText2CypherRetriever - retrievers/__init__.py: Export all async retriever classes Usage: driver = neo4j.AsyncGraphDatabase.driver(URI, auth=AUTH) retriever = await AsyncVectorRetriever(driver, 'my-index').async_init() results = await retriever.search(query_text='find something', top_k=5) Closes neo4j#406
|
Hi @stellasia — I've opened this PR to address #406. I also submitted the Neo4j CLA to cla@neo4j.com earlier today, so that should be processing. Happy to make any adjustments based on your feedback — particularly around the |
| """Fetch Neo4j schema if not provided. Must be awaited after construction.""" | ||
| if not self.neo4j_schema: | ||
| try: | ||
| self.neo4j_schema = get_schema(self.driver) |
There was a problem hiding this comment.
get_schema is sync. Either we should implement an sdync version or use asyncio.to_thread
| ) | ||
| self.result_formatter = validated_data.result_formatter | ||
| self._node_label = None | ||
| self._node_embedding_property = None |
There was a problem hiding this comment.
In some places we have _node_embedding_property while in some others we have _embedding_node_property
|
|
||
| use_search_clause = False | ||
| filter_cls: Optional[FilterClassification] = None | ||
| if supports_search_clause(self.driver, self.neo4j_database): |
There was a problem hiding this comment.
supports_search_clause silently always returns False with an AsyncDriver
|
|
||
| use_search_clause = False | ||
| filter_cls: Optional[FilterClassification] = None | ||
| if supports_search_clause(self.driver, self.neo4j_database): |
There was a problem hiding this comment.
Same issue here: supports_search_clause always returns False
| self._embedding_dimension = result["dimensions"] | ||
| self._filterable_properties = result.get("filterable_properties") or [] | ||
| except IndexError as e: | ||
| raise Exception(f"No index with name {self.index_name} found") from e |
There was a problem hiding this comment.
Does self.index_name exist on hybrid retrievers? I don't see it in AsyncHybridRetriever
| logger.debug("AsyncText2CypherRetriever prompt: %s", prompt) | ||
|
|
||
| try: | ||
| llm_result = self.llm.invoke(prompt) |
There was a problem hiding this comment.
llm_invoke is a sync method
| @@ -0,0 +1,193 @@ | |||
| # Copyright (c) "Neo4j" | |||
There was a problem hiding this comment.
How about tests for AsyncHybridRetriever, AsyncHybridCypherRetriever, and AsyncText2CypherRetriever?
| raise EmbeddingRequiredError( | ||
| "Embedding method required for text query." | ||
| ) | ||
| query_vector = self.embedder.embed_query(query_text) |
There was a problem hiding this comment.
embed_query() is a sync one.
2 participants
Summary
Closes #406
Adds async variants of all core retrievers so users can use
neo4j.AsyncDriverwithout maintaining two driver instances in parallel or spawning threads to avoid blocking the event loop.Usage
Design
async_init()is a separate awaitable method (since__init__cannot be async) that fetches index metadata from Neo4j — mirrors the sync retriever pattern where_fetch_index_infos()is called in__init__driver.execute_query()calls are replaced withawait driver.execute_query()AsyncText2CypherRetriever.async_init()fetches the Neo4j schema if not provided at construction timeChanges
src/neo4j_graphrag/types.pyAsyncNeo4jDriverModel— validatesneo4j.AsyncDriverAsyncVectorRetrieverModel,AsyncVectorCypherRetrieverModel,AsyncHybridRetrieverModel,AsyncHybridCypherRetrieverModel,AsyncText2CypherRetrieverModelsrc/neo4j_graphrag/retrievers/base.pyAsyncRetrieverABC withasync search(),async get_search_results(),async _fetch_index_infos()New files
retrievers/async_vector.py—AsyncVectorRetriever,AsyncVectorCypherRetrieverretrievers/async_hybrid.py—AsyncHybridRetriever,AsyncHybridCypherRetrieverretrievers/async_text2cypher.py—AsyncText2CypherRetrieverretrievers/__init__.pytests/unit/retrievers/test_async_vector.py