From 2bfd8d2afc0bff659c618e263dd49ca2db2aaeb0 Mon Sep 17 00:00:00 2001 From: 1HazyOne707 Date: Sat, 27 Jun 2026 12:00:40 -0700 Subject: [PATCH] feat(retrievers): add async retriever support via neo4j.AsyncDriver (#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 #406 --- src/neo4j_graphrag/retrievers/__init__.py | 8 + src/neo4j_graphrag/retrievers/async_hybrid.py | 386 ++++++++++++++ .../retrievers/async_text2cypher.py | 195 ++++++++ src/neo4j_graphrag/retrievers/async_vector.py | 470 ++++++++++++++++++ src/neo4j_graphrag/retrievers/base.py | 64 +++ src/neo4j_graphrag/types.py | 59 +++ tests/unit/retrievers/test_async_vector.py | 193 +++++++ 7 files changed, 1375 insertions(+) create mode 100644 src/neo4j_graphrag/retrievers/async_hybrid.py create mode 100644 src/neo4j_graphrag/retrievers/async_text2cypher.py create mode 100644 src/neo4j_graphrag/retrievers/async_vector.py create mode 100644 tests/unit/retrievers/test_async_vector.py diff --git a/src/neo4j_graphrag/retrievers/__init__.py b/src/neo4j_graphrag/retrievers/__init__.py index 681cd0e8c..a4c128feb 100644 --- a/src/neo4j_graphrag/retrievers/__init__.py +++ b/src/neo4j_graphrag/retrievers/__init__.py @@ -13,6 +13,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +from .async_hybrid import AsyncHybridCypherRetriever, AsyncHybridRetriever +from .async_text2cypher import AsyncText2CypherRetriever +from .async_vector import AsyncVectorCypherRetriever, AsyncVectorRetriever from .hybrid import HybridCypherRetriever, HybridRetriever from .text2cypher import Text2CypherRetriever from .tools_retriever import ToolsRetriever @@ -25,6 +28,11 @@ "HybridCypherRetriever", "Text2CypherRetriever", "ToolsRetriever", + "AsyncVectorRetriever", + "AsyncVectorCypherRetriever", + "AsyncHybridRetriever", + "AsyncHybridCypherRetriever", + "AsyncText2CypherRetriever", ] diff --git a/src/neo4j_graphrag/retrievers/async_hybrid.py b/src/neo4j_graphrag/retrievers/async_hybrid.py new file mode 100644 index 000000000..84e5cc6a6 --- /dev/null +++ b/src/neo4j_graphrag/retrievers/async_hybrid.py @@ -0,0 +1,386 @@ +# 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 __future__ import annotations + +import logging +from typing import Any, Callable, Optional, Union + +import neo4j +from pydantic import ValidationError + +from neo4j_graphrag.embeddings.base import Embedder +from neo4j_graphrag.exceptions import ( + EmbeddingRequiredError, + RetrieverInitializationError, + SearchValidationError, + SearchQueryParseError, +) +from neo4j_graphrag.neo4j_queries import ( + _build_hybrid_search_clause_query, + _build_hybrid_search_clause_query_linear, + get_query_tail, + get_search_query, +) +from neo4j_graphrag.retrievers.base import AsyncRetriever +from neo4j_graphrag.utils.version_utils import supports_search_clause +from neo4j_graphrag.types import ( + AsyncNeo4jDriverModel, + AsyncHybridRetrieverModel, + AsyncHybridCypherRetrieverModel, + EmbedderModel, + HybridCypherSearchModel, + HybridSearchModel, + HybridSearchRanker, + RawSearchResult, + RetrieverResultItem, + SearchType, +) +from neo4j_graphrag.utils.logging import prettify + +logger = logging.getLogger(__name__) + + +class AsyncHybridRetriever(AsyncRetriever): + """ + Provides async retrieval using combination of vector search and fulltext search. + + Args: + driver (neo4j.AsyncDriver): The Neo4j async Python driver. + vector_index_name (str): Vector index name. + fulltext_index_name (str): Fulltext index name. + embedder (Optional[Embedder]): Embedder object to embed query text. + return_properties (Optional[list[str]]): List of node properties to return. + result_formatter (Optional[Callable[[neo4j.Record], RetrieverResultItem]]): Custom formatter. + neo4j_database (Optional[str]): The name of the Neo4j database. + + Raises: + RetrieverInitializationError: If validation of the input arguments fail. + """ + + def __init__( + self, + driver: neo4j.AsyncDriver, + vector_index_name: str, + fulltext_index_name: str, + embedder: Optional[Embedder] = None, + return_properties: Optional[list[str]] = None, + result_formatter: Optional[ + Callable[[neo4j.Record], RetrieverResultItem] + ] = None, + neo4j_database: Optional[str] = None, + ) -> None: + try: + driver_model = AsyncNeo4jDriverModel(driver=driver) + embedder_model = EmbedderModel(embedder=embedder) if embedder else None + validated_data = AsyncHybridRetrieverModel( + driver_model=driver_model, + vector_index_name=vector_index_name, + fulltext_index_name=fulltext_index_name, + embedder_model=embedder_model, + return_properties=return_properties, + result_formatter=result_formatter, + neo4j_database=neo4j_database, + ) + except ValidationError as e: + raise RetrieverInitializationError(e.errors()) from e + + super().__init__( + validated_data.driver_model.driver, validated_data.neo4j_database + ) + self.vector_index_name = validated_data.vector_index_name + self.fulltext_index_name = validated_data.fulltext_index_name + self.return_properties = validated_data.return_properties + self.embedder = ( + validated_data.embedder_model.embedder + if validated_data.embedder_model + else None + ) + self.result_formatter = validated_data.result_formatter + self._node_label = None + self._embedding_node_property = None + self._embedding_dimension = None + + async def async_init(self) -> "AsyncHybridRetriever": + """Fetch index metadata. Must be awaited after construction.""" + await self._fetch_index_infos(self.vector_index_name) + return self + + def default_record_formatter(self, record: neo4j.Record) -> RetrieverResultItem: + metadata = {"score": record.get("score")} + node = record.get("node") + return RetrieverResultItem(content=str(node), metadata=metadata) + + async def get_search_results( + self, + query_text: str, + query_vector: Optional[list[float]] = None, + top_k: int = 5, + effective_search_ratio: int = 1, + ranker: Union[str, HybridSearchRanker] = HybridSearchRanker.NAIVE, + alpha: Optional[float] = None, + ) -> RawSearchResult: + """Async hybrid search. See HybridRetriever.get_search_results for full docs.""" + try: + validated_data = HybridSearchModel( + query_vector=query_vector, + query_text=query_text, + top_k=top_k, + effective_search_ratio=effective_search_ratio, + ranker=ranker, + alpha=alpha, + ) + except ValidationError as e: + raise SearchValidationError(e.errors()) from e + + parameters = validated_data.model_dump(exclude_none=True) + parameters["vector_index_name"] = self.vector_index_name + parameters["fulltext_index_name"] = self.fulltext_index_name + + if query_text and not query_vector: + if not self.embedder: + raise EmbeddingRequiredError( + "Embedding method required for text query." + ) + query_vector = self.embedder.embed_query(query_text) + parameters["query_vector"] = query_vector + + use_search_clause = False + if supports_search_clause(self.driver, self.neo4j_database): + if self._node_label: + use_search_clause = True + + if use_search_clause: + if validated_data.ranker == HybridSearchRanker.LINEAR and validated_data.alpha: + search_query_base = _build_hybrid_search_clause_query_linear( + vector_index_name=self.vector_index_name, + fulltext_index_name=self.fulltext_index_name, + node_label=self._node_label or "", + ) + else: + search_query_base = _build_hybrid_search_clause_query( + vector_index_name=self.vector_index_name, + fulltext_index_name=self.fulltext_index_name, + node_label=self._node_label or "", + ) + query_tail = get_query_tail( + return_properties=self.return_properties, + fallback_return=( + f"RETURN node {{ .*, `{self._embedding_node_property}`: null }} AS node, " + "labels(node) AS nodeLabels, elementId(node) AS elementId, " + "elementId(node) AS id, score" + ) + if self._embedding_node_property + else ( + "RETURN node, labels(node) AS nodeLabels, " + "elementId(node) AS elementId, elementId(node) AS id, score" + ), + ) + search_query = f"{search_query_base} {query_tail}" + else: + search_query, _ = get_search_query( + search_type=SearchType.HYBRID, + return_properties=self.return_properties, + embedding_node_property=self._embedding_node_property, + neo4j_version_is_5_23_or_above=getattr(self, "neo4j_version_is_5_23_or_above", True), + ranker=validated_data.ranker, + alpha=validated_data.alpha, + ) + + if "ranker" in parameters: + del parameters["ranker"] + + logger.debug("AsyncHybridRetriever Cypher parameters: %s", prettify(parameters)) + logger.debug("AsyncHybridRetriever Cypher query: %s", search_query) + + try: + result = await self.driver.execute_query( + search_query, + parameters, + database_=self.neo4j_database, + routing_=neo4j.RoutingControl.READ, + ) + records = result.records + except neo4j.exceptions.ClientError as e: + if "org.apache.lucene.queryparser.classic.ParseException" in str(e): + raise SearchQueryParseError( + f"Invalid Lucene query generated from query_text: {query_text}" + ) from e + raise + return RawSearchResult( + records=records, + metadata={"query_vector": query_vector}, + ) + + +class AsyncHybridCypherRetriever(AsyncRetriever): + """ + Provides async retrieval using hybrid search augmented by a Cypher query. + + Args: + driver (neo4j.AsyncDriver): The Neo4j async Python driver. + vector_index_name (str): Vector index name. + fulltext_index_name (str): Fulltext index name. + retrieval_query (str): Cypher query appended to the hybrid search. + embedder (Optional[Embedder]): Embedder object to embed query text. + result_formatter (Optional[Callable[[neo4j.Record], RetrieverResultItem]]): Custom formatter. + neo4j_database (Optional[str]): The name of the Neo4j database. + + Raises: + RetrieverInitializationError: If validation of the input arguments fail. + """ + + def __init__( + self, + driver: neo4j.AsyncDriver, + vector_index_name: str, + fulltext_index_name: str, + retrieval_query: str, + embedder: Optional[Embedder] = None, + result_formatter: Optional[ + Callable[[neo4j.Record], RetrieverResultItem] + ] = None, + neo4j_database: Optional[str] = None, + ) -> None: + try: + driver_model = AsyncNeo4jDriverModel(driver=driver) + embedder_model = EmbedderModel(embedder=embedder) if embedder else None + validated_data = AsyncHybridCypherRetrieverModel( + driver_model=driver_model, + vector_index_name=vector_index_name, + fulltext_index_name=fulltext_index_name, + retrieval_query=retrieval_query, + embedder_model=embedder_model, + result_formatter=result_formatter, + neo4j_database=neo4j_database, + ) + except ValidationError as e: + raise RetrieverInitializationError(e.errors()) from e + + super().__init__( + validated_data.driver_model.driver, validated_data.neo4j_database + ) + self.vector_index_name = validated_data.vector_index_name + self.fulltext_index_name = validated_data.fulltext_index_name + self.retrieval_query = validated_data.retrieval_query + self.embedder = ( + validated_data.embedder_model.embedder + if validated_data.embedder_model + else None + ) + self.result_formatter = validated_data.result_formatter + self._node_label = None + self._embedding_node_property = None + self._embedding_dimension = None + + async def async_init(self) -> "AsyncHybridCypherRetriever": + """Fetch index metadata. Must be awaited after construction.""" + await self._fetch_index_infos(self.vector_index_name) + return self + + async def get_search_results( + self, + query_text: str, + query_vector: Optional[list[float]] = None, + top_k: int = 5, + effective_search_ratio: int = 1, + query_params: Optional[dict[str, Any]] = None, + ranker: Union[str, HybridSearchRanker] = HybridSearchRanker.NAIVE, + alpha: Optional[float] = None, + ) -> RawSearchResult: + """Async hybrid+cypher search. See HybridCypherRetriever.get_search_results for full docs.""" + try: + validated_data = HybridCypherSearchModel( + query_vector=query_vector, + query_text=query_text, + top_k=top_k, + effective_search_ratio=effective_search_ratio, + ranker=ranker, + alpha=alpha, + query_params=query_params, + ) + except ValidationError as e: + raise SearchValidationError(e.errors()) from e + + parameters = validated_data.model_dump(exclude_none=True) + parameters["vector_index_name"] = self.vector_index_name + parameters["fulltext_index_name"] = self.fulltext_index_name + + if query_text and not query_vector: + if not self.embedder: + raise EmbeddingRequiredError( + "Embedding method required for text query." + ) + query_vector = self.embedder.embed_query(query_text) + parameters["query_vector"] = query_vector + + if query_params: + for key, value in query_params.items(): + if key not in parameters: + parameters[key] = value + del parameters["query_params"] + + use_search_clause = False + if supports_search_clause(self.driver, self.neo4j_database): + if self._node_label: + use_search_clause = True + + if use_search_clause: + if validated_data.ranker == HybridSearchRanker.LINEAR and validated_data.alpha: + search_query_base = _build_hybrid_search_clause_query_linear( + vector_index_name=self.vector_index_name, + fulltext_index_name=self.fulltext_index_name, + node_label=self._node_label or "", + ) + else: + search_query_base = _build_hybrid_search_clause_query( + vector_index_name=self.vector_index_name, + fulltext_index_name=self.fulltext_index_name, + node_label=self._node_label or "", + ) + query_tail = get_query_tail(retrieval_query=self.retrieval_query) + search_query = f"{search_query_base} {query_tail}" + else: + search_query, _ = get_search_query( + search_type=SearchType.HYBRID, + retrieval_query=self.retrieval_query, + neo4j_version_is_5_23_or_above=getattr(self, "neo4j_version_is_5_23_or_above", True), + ranker=validated_data.ranker, + alpha=validated_data.alpha, + ) + + if "ranker" in parameters: + del parameters["ranker"] + + logger.debug("AsyncHybridCypherRetriever Cypher parameters: %s", prettify(parameters)) + logger.debug("AsyncHybridCypherRetriever Cypher query: %s", search_query) + + try: + result = await self.driver.execute_query( + search_query, + parameters, + database_=self.neo4j_database, + routing_=neo4j.RoutingControl.READ, + ) + records = result.records + except neo4j.exceptions.ClientError as e: + if "org.apache.lucene.queryparser.classic.ParseException" in str(e): + raise SearchQueryParseError( + f"Invalid Lucene query generated from query_text: {query_text}" + ) from e + raise + return RawSearchResult( + records=records, + metadata={"query_vector": query_vector}, + ) diff --git a/src/neo4j_graphrag/retrievers/async_text2cypher.py b/src/neo4j_graphrag/retrievers/async_text2cypher.py new file mode 100644 index 000000000..ef859df53 --- /dev/null +++ b/src/neo4j_graphrag/retrievers/async_text2cypher.py @@ -0,0 +1,195 @@ +# 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 __future__ import annotations + +import logging +from typing import Any, Callable, Dict, Optional + +import neo4j +from neo4j.exceptions import CypherSyntaxError, DriverError, Neo4jError +from pydantic import ValidationError + +from neo4j_graphrag.exceptions import ( + RetrieverInitializationError, + SchemaFetchError, + SearchValidationError, + Text2CypherRetrievalError, +) +from neo4j_graphrag.generation.prompts import Text2CypherTemplate +from neo4j_graphrag.llm import LLMInterface +from neo4j_graphrag.retrievers.base import AsyncRetriever +from neo4j_graphrag.retrievers.text2cypher import extract_cypher, READ_ONLY_QUERY_TYPE +from neo4j_graphrag.schema import get_schema +from neo4j_graphrag.types import ( + AsyncNeo4jDriverModel, + AsyncText2CypherRetrieverModel, + LLMModel, + Neo4jSchemaModel, + RawSearchResult, + RetrieverResultItem, + Text2CypherSearchModel, +) + +logger = logging.getLogger(__name__) + + +class AsyncText2CypherRetriever(AsyncRetriever): + """ + Allows async retrieval of records from a Neo4j database using natural language. + Converts a user's natural language query to a Cypher query using an LLM, + then retrieves records from a Neo4j database using the generated Cypher query. + + Args: + driver (neo4j.AsyncDriver): The Neo4j async Python driver. + llm (LLMInterface): LLM object to generate the Cypher query. + neo4j_schema (Optional[str]): Neo4j schema used to generate the Cypher query. + examples (Optional[list[str]]): Optional user input/query pairs for the LLM to use as examples. + result_formatter (Optional[Callable[[neo4j.Record], RetrieverResultItem]]): Custom formatter. + custom_prompt (Optional[str]): Optional custom prompt template. + neo4j_database (Optional[str]): The name of the Neo4j database. + + Raises: + RetrieverInitializationError: If validation of the input arguments fail. + """ + + def __init__( + self, + driver: neo4j.AsyncDriver, + llm: LLMInterface, + neo4j_schema: Optional[str] = None, + examples: Optional[list[str]] = None, + result_formatter: Optional[ + Callable[[neo4j.Record], RetrieverResultItem] + ] = None, + custom_prompt: Optional[str] = None, + neo4j_database: Optional[str] = None, + ) -> None: + try: + driver_model = AsyncNeo4jDriverModel(driver=driver) + llm_model = LLMModel(llm=llm) + neo4j_schema_model = ( + Neo4jSchemaModel(neo4j_schema=neo4j_schema) if neo4j_schema else None + ) + validated_data = AsyncText2CypherRetrieverModel( + driver_model=driver_model, + llm_model=llm_model, + neo4j_schema_model=neo4j_schema_model, + examples=examples, + result_formatter=result_formatter, + custom_prompt=custom_prompt, + neo4j_database=neo4j_database, + ) + except ValidationError as e: + raise RetrieverInitializationError(e.errors()) from e + + super().__init__( + validated_data.driver_model.driver, validated_data.neo4j_database + ) + self.llm = validated_data.llm_model.llm + self.examples = validated_data.examples + self.result_formatter = validated_data.result_formatter + self.custom_prompt = validated_data.custom_prompt + self.neo4j_schema = neo4j_schema or "" + + async def async_init(self) -> "AsyncText2CypherRetriever": + """Fetch Neo4j schema if not provided. Must be awaited after construction.""" + if not self.neo4j_schema: + try: + self.neo4j_schema = get_schema(self.driver) + except (Neo4jError, DriverError) as e: + error_message = getattr(e, "message", str(e)) + raise SchemaFetchError( + f"Failed to fetch schema for AsyncText2CypherRetriever: {error_message}" + ) from e + return self + + async def get_search_results( + self, + query_text: str, + prompt_params: Optional[Dict[str, Any]] = None, + ) -> RawSearchResult: + """Async version of Text2CypherRetriever.get_search_results. + + Converts query_text to a Cypher query using an LLM, then retrieves + records from Neo4j using the generated Cypher query. + + Args: + query_text (str): The natural language query. + prompt_params (Optional[Dict[str, Any]]): Additional prompt parameters. + + Raises: + SearchValidationError: If validation of the input arguments fail. + Text2CypherRetrievalError: If the LLM fails to generate a correct Cypher query. + + Returns: + RawSearchResult: The results of the search query. + """ + try: + validated_data = Text2CypherSearchModel(query_text=query_text) + except ValidationError as e: + raise SearchValidationError(e.errors()) from e + + prompt_template = Text2CypherTemplate(template=self.custom_prompt) + + if prompt_params is not None: + examples_to_use = prompt_params.pop("examples", None) or ( + "\n".join(self.examples) if self.examples else "" + ) + schema_to_use = prompt_params.pop("schema", None) or self.neo4j_schema + else: + examples_to_use = "\n".join(self.examples) if self.examples else "" + schema_to_use = self.neo4j_schema + prompt_params = dict() + + prompt = prompt_template.format( + schema=schema_to_use, + examples=examples_to_use, + query_text=validated_data.query_text, + **prompt_params, + ) + + logger.debug("AsyncText2CypherRetriever prompt: %s", prompt) + + try: + llm_result = self.llm.invoke(prompt) + t2c_query = extract_cypher(llm_result.content) + logger.debug("AsyncText2CypherRetriever Cypher query: %s", t2c_query) + + explain_result = await self.driver.execute_query( + query_=f"EXPLAIN {t2c_query}", + database_=self.neo4j_database, + routing_=neo4j.RoutingControl.READ, + ) + explain_summary = explain_result.summary + if explain_summary.query_type != READ_ONLY_QUERY_TYPE: + raise Text2CypherRetrievalError( + "Refusing to execute non-read-only Cypher " + f"(query_type={explain_summary.query_type!r}): {t2c_query}" + ) + result = await self.driver.execute_query( + query_=t2c_query, + database_=self.neo4j_database, + routing_=neo4j.RoutingControl.READ, + ) + records = result.records + except CypherSyntaxError as e: + raise Text2CypherRetrievalError( + f"Failed to get search result: {e.message}" + ) from e + + return RawSearchResult( + records=records, + metadata={"cypher": t2c_query}, + ) diff --git a/src/neo4j_graphrag/retrievers/async_vector.py b/src/neo4j_graphrag/retrievers/async_vector.py new file mode 100644 index 000000000..bb6dd514e --- /dev/null +++ b/src/neo4j_graphrag/retrievers/async_vector.py @@ -0,0 +1,470 @@ +# 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 __future__ import annotations + +import logging +from typing import Any, Callable, Optional + +import neo4j +from pydantic import ValidationError + +from neo4j_graphrag.embeddings.base import Embedder +from neo4j_graphrag.exceptions import ( + EmbeddingRequiredError, + RetrieverInitializationError, + SearchValidationError, +) +from neo4j_graphrag.filters import ( + FilterClassification, + classify_filter_for_search, + extract_filter_field_names, +) +from neo4j_graphrag.neo4j_queries import ( + _build_search_clause_vector_query, + get_query_tail, + get_search_query, +) +from neo4j_graphrag.retrievers.base import AsyncRetriever +from neo4j_graphrag.utils.version_utils import supports_search_clause +from neo4j_graphrag.types import ( + AsyncNeo4jDriverModel, + AsyncVectorRetrieverModel, + AsyncVectorCypherRetrieverModel, + EmbedderModel, + RawSearchResult, + RetrieverResultItem, + SearchType, + VectorCypherSearchModel, + VectorSearchModel, +) +from neo4j_graphrag.utils.logging import prettify + +logger = logging.getLogger(__name__) + + +class AsyncVectorRetriever(AsyncRetriever): + """ + Provides async retrieval method using vector search over embeddings. + + Example: + + .. code-block:: python + + import neo4j + from neo4j_graphrag.retrievers import AsyncVectorRetriever + + driver = neo4j.AsyncGraphDatabase.driver(URI, auth=AUTH) + retriever = AsyncVectorRetriever(driver, "vector-index-name", custom_embedder) + await retriever.search(query_text="Find me a book about Fremen", top_k=5) + + Args: + driver (neo4j.AsyncDriver): The Neo4j async Python driver. + index_name (str): Vector index name. + embedder (Optional[Embedder]): Embedder object to embed query text. + return_properties (Optional[list[str]]): List of node properties to return. + result_formatter (Optional[Callable[[neo4j.Record], RetrieverResultItem]]): Custom formatter. + neo4j_database (Optional[str]): The name of the Neo4j database. + + Raises: + RetrieverInitializationError: If validation of the input arguments fail. + """ + + def __init__( + self, + driver: neo4j.AsyncDriver, + index_name: str, + embedder: Optional[Embedder] = None, + return_properties: Optional[list[str]] = None, + result_formatter: Optional[ + Callable[[neo4j.Record], RetrieverResultItem] + ] = None, + neo4j_database: Optional[str] = None, + ) -> None: + try: + driver_model = AsyncNeo4jDriverModel(driver=driver) + embedder_model = EmbedderModel(embedder=embedder) if embedder else None + validated_data = AsyncVectorRetrieverModel( + driver_model=driver_model, + index_name=index_name, + embedder_model=embedder_model, + return_properties=return_properties, + result_formatter=result_formatter, + neo4j_database=neo4j_database, + ) + except ValidationError as e: + raise RetrieverInitializationError(e.errors()) from e + + super().__init__( + validated_data.driver_model.driver, validated_data.neo4j_database + ) + self.index_name = validated_data.index_name + self.return_properties = validated_data.return_properties + self.embedder = ( + validated_data.embedder_model.embedder + if validated_data.embedder_model + else None + ) + self.result_formatter = validated_data.result_formatter + self._node_label = None + self._embedding_node_property = None + self._embedding_dimension = None + self._filterable_properties: list[str] = [] + + async def async_init(self) -> "AsyncVectorRetriever": + """Fetch index metadata. Must be awaited after construction.""" + await self._fetch_index_infos(self.index_name) + return self + + def default_record_formatter(self, record: neo4j.Record) -> RetrieverResultItem: + metadata = { + "score": record.get("score"), + "nodeLabels": record.get("nodeLabels"), + "id": record.get("id"), + } + node = record.get("node") + return RetrieverResultItem( + content=str(node), + metadata=metadata, + ) + + async def get_search_results( + self, + query_vector: Optional[list[float]] = None, + query_text: Optional[str] = None, + top_k: int = 5, + effective_search_ratio: int = 1, + filters: Optional[dict[str, Any]] = None, + ) -> RawSearchResult: + """Async vector search. See VectorRetriever.get_search_results for full docs.""" + try: + validated_data = VectorSearchModel( + query_vector=query_vector, + query_text=query_text, + top_k=top_k, + effective_search_ratio=effective_search_ratio, + filters=filters, + ) + except ValidationError as e: + raise SearchValidationError(e.errors()) from e + + parameters = validated_data.model_dump(exclude_none=True) + parameters["vector_index_name"] = self.index_name + if filters: + del parameters["filters"] + + if query_text: + if not self.embedder: + raise EmbeddingRequiredError( + "Embedding method required for text query." + ) + query_vector = self.embedder.embed_query(query_text) + parameters["query_vector"] = query_vector + del parameters["query_text"] + + use_search_clause = False + filter_cls: Optional[FilterClassification] = None + if supports_search_clause(self.driver, self.neo4j_database): + if filters: + filter_cls = classify_filter_for_search(filters, node_alias="node") + missing = extract_filter_field_names(filters) - set( + self._filterable_properties + ) + if not filter_cls.is_compatible: + logger.warning( + "Filters are not compatible with SEARCH clause " + "in-index filtering; falling back to procedure-based " + "vector search with brute-force filtering." + ) + elif missing: + logger.warning( + "Filter properties %s not declared as filterable on " + "index '%s'; falling back to procedure-based vector " + "search. Recreate the index with filterable_properties " + "to use in-index filtering.", + sorted(missing), + self.index_name, + ) + elif self._node_label: + use_search_clause = True + else: + if self._node_label: + use_search_clause = True + + if use_search_clause: + search_query_base, search_params = _build_search_clause_vector_query( + index_name=self.index_name, + node_label=self._node_label or "", + filter_classification=filter_cls, + ) + query_tail = get_query_tail( + return_properties=self.return_properties, + fallback_return=( + f"RETURN node {{ .*, `{self._embedding_node_property}`: null }} AS node, " + "labels(node) AS nodeLabels, " + "elementId(node) AS elementId, " + "elementId(node) AS id, " + "score" + ) + if self._embedding_node_property + else ( + "RETURN node, labels(node) AS nodeLabels, " + "elementId(node) AS elementId, " + "elementId(node) AS id, " + "score" + ), + ) + search_query = f"{search_query_base} {query_tail}" + else: + search_query, search_params = get_search_query( + search_type=SearchType.VECTOR, + return_properties=self.return_properties, + node_label=self._node_label, + embedding_node_property=self._embedding_node_property, + embedding_dimension=self._embedding_dimension, + filters=filters, + ) + parameters.update(search_params) + + logger.debug("AsyncVectorRetriever Cypher parameters: %s", prettify(parameters)) + logger.debug("AsyncVectorRetriever Cypher query: %s", search_query) + + try: + result = await self.driver.execute_query( + search_query, + parameters, + database_=self.neo4j_database, + routing_=neo4j.RoutingControl.READ, + ) + records = result.records + except neo4j.exceptions.ClientError as e: + if use_search_clause and "PropertyNotFound" in str(e): + logger.warning( + "SEARCH clause failed; falling back to procedure-based vector search. Error: %s", + e, + ) + search_query, search_params = get_search_query( + search_type=SearchType.VECTOR, + return_properties=self.return_properties, + node_label=self._node_label, + embedding_node_property=self._embedding_node_property, + embedding_dimension=self._embedding_dimension, + filters=filters, + ) + parameters.update(search_params) + result = await self.driver.execute_query( + search_query, + parameters, + database_=self.neo4j_database, + routing_=neo4j.RoutingControl.READ, + ) + records = result.records + else: + raise + return RawSearchResult( + records=records, + metadata={"query_vector": query_vector}, + ) + + +class AsyncVectorCypherRetriever(AsyncRetriever): + """ + Provides async retrieval using vector similarity augmented by a Cypher query. + + Args: + driver (neo4j.AsyncDriver): The Neo4j async Python driver. + index_name (str): Vector index name. + retrieval_query (str): Cypher query appended to the vector search. + embedder (Optional[Embedder]): Embedder object to embed query text. + result_formatter (Optional[Callable[[neo4j.Record], RetrieverResultItem]]): Custom formatter. + neo4j_database (Optional[str]): The name of the Neo4j database. + + Raises: + RetrieverInitializationError: If validation of the input arguments fail. + """ + + def __init__( + self, + driver: neo4j.AsyncDriver, + index_name: str, + retrieval_query: str, + embedder: Optional[Embedder] = None, + result_formatter: Optional[ + Callable[[neo4j.Record], RetrieverResultItem] + ] = None, + neo4j_database: Optional[str] = None, + ) -> None: + try: + driver_model = AsyncNeo4jDriverModel(driver=driver) + embedder_model = EmbedderModel(embedder=embedder) if embedder else None + validated_data = AsyncVectorCypherRetrieverModel( + driver_model=driver_model, + index_name=index_name, + retrieval_query=retrieval_query, + embedder_model=embedder_model, + result_formatter=result_formatter, + neo4j_database=neo4j_database, + ) + except ValidationError as e: + raise RetrieverInitializationError(e.errors()) from e + + super().__init__( + validated_data.driver_model.driver, validated_data.neo4j_database + ) + self.index_name = validated_data.index_name + self.retrieval_query = validated_data.retrieval_query + self.embedder = ( + validated_data.embedder_model.embedder + if validated_data.embedder_model + else None + ) + self.result_formatter = validated_data.result_formatter + self._node_label = None + self._node_embedding_property = None + self._embedding_dimension = None + self._filterable_properties: list[str] = [] + + async def async_init(self) -> "AsyncVectorCypherRetriever": + """Fetch index metadata. Must be awaited after construction.""" + await self._fetch_index_infos(self.index_name) + return self + + async def get_search_results( + self, + query_vector: Optional[list[float]] = None, + query_text: Optional[str] = None, + top_k: int = 5, + effective_search_ratio: int = 1, + query_params: Optional[dict[str, Any]] = None, + filters: Optional[dict[str, Any]] = None, + ) -> RawSearchResult: + """Async vector+cypher search. See VectorCypherRetriever.get_search_results for full docs.""" + try: + validated_data = VectorCypherSearchModel( + query_vector=query_vector, + query_text=query_text, + top_k=top_k, + effective_search_ratio=effective_search_ratio, + query_params=query_params, + filters=filters, + ) + except ValidationError as e: + raise SearchValidationError(e.errors()) from e + + parameters = validated_data.model_dump(exclude_none=True) + parameters["vector_index_name"] = self.index_name + if filters: + del parameters["filters"] + + if query_text: + if not self.embedder: + raise EmbeddingRequiredError( + "Embedding method required for text query." + ) + query_vector = self.embedder.embed_query(query_text) + parameters["query_vector"] = query_vector + del parameters["query_text"] + + if query_params: + for key, value in query_params.items(): + if key not in parameters: + parameters[key] = value + del parameters["query_params"] + + use_search_clause = False + filter_cls: Optional[FilterClassification] = None + if supports_search_clause(self.driver, self.neo4j_database): + if filters: + filter_cls = classify_filter_for_search(filters, node_alias="node") + missing = extract_filter_field_names(filters) - set( + self._filterable_properties + ) + if not filter_cls.is_compatible: + logger.warning( + "Filters are not compatible with SEARCH clause " + "in-index filtering; falling back to procedure-based " + "vector search with brute-force filtering." + ) + elif missing: + logger.warning( + "Filter properties %s not declared as filterable on " + "index '%s'; falling back to procedure-based vector " + "search.", + sorted(missing), + self.index_name, + ) + elif self._node_label: + use_search_clause = True + else: + if self._node_label: + use_search_clause = True + + if use_search_clause: + search_query_base, search_params = _build_search_clause_vector_query( + index_name=self.index_name, + node_label=self._node_label or "", + filter_classification=filter_cls, + ) + query_tail = get_query_tail(retrieval_query=self.retrieval_query) + search_query = f"{search_query_base} {query_tail}" + else: + search_query, search_params = get_search_query( + search_type=SearchType.VECTOR, + retrieval_query=self.retrieval_query, + node_label=self._node_label, + embedding_node_property=self._node_embedding_property, + embedding_dimension=self._embedding_dimension, + filters=filters, + ) + parameters.update(search_params) + + logger.debug("AsyncVectorCypherRetriever Cypher parameters: %s", prettify(parameters)) + logger.debug("AsyncVectorCypherRetriever Cypher query: %s", search_query) + + try: + result = await self.driver.execute_query( + search_query, + parameters, + database_=self.neo4j_database, + routing_=neo4j.RoutingControl.READ, + ) + records = result.records + except neo4j.exceptions.ClientError as e: + if use_search_clause and "PropertyNotFound" in str(e): + logger.warning( + "SEARCH clause failed; falling back to procedure-based vector search. Error: %s", + e, + ) + search_query, search_params = get_search_query( + search_type=SearchType.VECTOR, + retrieval_query=self.retrieval_query, + node_label=self._node_label, + embedding_node_property=self._node_embedding_property, + embedding_dimension=self._embedding_dimension, + filters=filters, + ) + parameters.update(search_params) + result = await self.driver.execute_query( + search_query, + parameters, + database_=self.neo4j_database, + routing_=neo4j.RoutingControl.READ, + ) + records = result.records + else: + raise + return RawSearchResult( + records=records, + metadata={"query_vector": query_vector}, + ) diff --git a/src/neo4j_graphrag/retrievers/base.py b/src/neo4j_graphrag/retrievers/base.py index 4fbe712e5..f343448ff 100644 --- a/src/neo4j_graphrag/retrievers/base.py +++ b/src/neo4j_graphrag/retrievers/base.py @@ -443,6 +443,70 @@ def execute_func(**kwargs: Any) -> Any: ) +class AsyncRetriever(ABC): + """ + Abstract base class for async Neo4j retrievers using neo4j.AsyncDriver. + """ + + index_name: str + VERIFY_NEO4J_VERSION = True + + def __init__(self, driver: "neo4j.AsyncDriver", neo4j_database: Optional[str] = None): + from neo4j_graphrag.utils import driver_config + self.driver = driver_config.override_user_agent(driver) + self.neo4j_database = neo4j_database + + async def _fetch_index_infos(self, vector_index_name: str) -> None: + """Fetch node label and embedding property from the index definition (async).""" + query = ( + "SHOW VECTOR INDEXES " + "YIELD name, labelsOrTypes, properties, options " + "WHERE name = $index_name " + "RETURN labelsOrTypes as labels, properties, " + "options.indexConfig.`vector.dimensions` as dimensions, " + "options.indexConfig.`vector.filterable_properties` as filterable_properties" + ) + query_result = await self.driver.execute_query( + query, + {"index_name": vector_index_name}, + database_=self.neo4j_database, + routing_=neo4j.RoutingControl.READ, + ) + try: + result = query_result.records[0] + self._node_label = result["labels"][0] + self._embedding_node_property = result["properties"][0] + 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 + + async def search(self, *args: Any, **kwargs: Any) -> "RetrieverResult": + """Async search method. Calls get_search_results and formats results.""" + raw_result = await self.get_search_results(*args, **kwargs) + formatter = self.get_result_formatter() + search_items = [formatter(record) for record in raw_result.records] + metadata = raw_result.metadata or {} + metadata["__retriever"] = self.__class__.__name__ + return RetrieverResult( + items=search_items, + metadata=metadata, + ) + + @abstractmethod + async def get_search_results(self, *args: Any, **kwargs: Any) -> "RawSearchResult": + """Must be implemented in subclasses. Returns RawSearchResult.""" + pass + + def get_result_formatter(self) -> Callable[[neo4j.Record], RetrieverResultItem]: + if hasattr(self, "result_formatter"): + return self.result_formatter or self.default_record_formatter + return self.default_record_formatter + + def default_record_formatter(self, record: neo4j.Record) -> RetrieverResultItem: + return RetrieverResultItem(content=str(record), metadata=record.get("metadata")) + + class ExternalRetriever(Retriever, ABC): """ Abstract class for External Vector Stores diff --git a/src/neo4j_graphrag/types.py b/src/neo4j_graphrag/types.py index 1225b5af3..bfa02d057 100644 --- a/src/neo4j_graphrag/types.py +++ b/src/neo4j_graphrag/types.py @@ -298,6 +298,65 @@ class Text2CypherRetrieverModel(BaseModel): neo4j_database: Optional[str] = None +class AsyncNeo4jDriverModel(BaseModel): + driver: neo4j.AsyncDriver + model_config = ConfigDict(arbitrary_types_allowed=True) + + @field_validator("driver") + def check_driver(cls, value: neo4j.AsyncDriver) -> neo4j.AsyncDriver: + if not isinstance(value, neo4j.AsyncDriver): + raise ValueError("Provided driver needs to be of type neo4j.AsyncDriver") + return value + + +class AsyncVectorRetrieverModel(BaseModel): + driver_model: AsyncNeo4jDriverModel + index_name: str + embedder_model: Optional[EmbedderModel] = None + return_properties: Optional[list[str]] = None + result_formatter: Optional[Callable[[neo4j.Record], RetrieverResultItem]] = None + neo4j_database: Optional[str] = None + + +class AsyncVectorCypherRetrieverModel(BaseModel): + driver_model: AsyncNeo4jDriverModel + index_name: str + retrieval_query: str + embedder_model: Optional[EmbedderModel] = None + result_formatter: Optional[Callable[[neo4j.Record], RetrieverResultItem]] = None + neo4j_database: Optional[str] = None + + +class AsyncHybridRetrieverModel(BaseModel): + driver_model: AsyncNeo4jDriverModel + vector_index_name: str + fulltext_index_name: str + embedder_model: Optional[EmbedderModel] = None + return_properties: Optional[list[str]] = None + result_formatter: Optional[Callable[[neo4j.Record], RetrieverResultItem]] = None + neo4j_database: Optional[str] = None + + +class AsyncHybridCypherRetrieverModel(BaseModel): + driver_model: AsyncNeo4jDriverModel + vector_index_name: str + fulltext_index_name: str + retrieval_query: str + embedder_model: Optional[EmbedderModel] = None + result_formatter: Optional[Callable[[neo4j.Record], RetrieverResultItem]] = None + neo4j_database: Optional[str] = None + + +class AsyncText2CypherRetrieverModel(BaseModel): + driver_model: AsyncNeo4jDriverModel + llm_model: LLMModel + neo4j_schema_model: Optional[Neo4jSchemaModel] = None + examples: Optional[list[str]] = None + result_formatter: Optional[Callable[[neo4j.Record], RetrieverResultItem]] = None + custom_prompt: Optional[str] = None + neo4j_database: Optional[str] = None + + class Neo4jMessageHistoryModel(BaseModel): session_id: Union[str, int] driver_model: Neo4jDriverModel diff --git a/tests/unit/retrievers/test_async_vector.py b/tests/unit/retrievers/test_async_vector.py new file mode 100644 index 000000000..76bb7cf53 --- /dev/null +++ b/tests/unit/retrievers/test_async_vector.py @@ -0,0 +1,193 @@ +# 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 __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock, patch + +import neo4j +import pytest + +from neo4j_graphrag.exceptions import ( + EmbeddingRequiredError, + RetrieverInitializationError, + SearchValidationError, +) +from neo4j_graphrag.retrievers import AsyncVectorRetriever, AsyncVectorCypherRetriever +from neo4j_graphrag.types import RetrieverResult, RetrieverResultItem + + +@pytest.fixture +def async_driver() -> MagicMock: + driver = MagicMock(spec=neo4j.AsyncDriver) + driver.execute_query = AsyncMock() + return driver + + +@pytest.fixture +def embedder() -> MagicMock: + mock = MagicMock() + mock.embed_query = MagicMock(return_value=[0.1, 0.2, 0.3]) + return mock + + +def _make_retriever(async_driver, index_name="my-index", embedder=None): + """Create AsyncVectorRetriever with mocked index info already set.""" + retriever = AsyncVectorRetriever( + driver=async_driver, + index_name=index_name, + embedder=embedder, + ) + retriever._node_label = "Document" + retriever._embedding_node_property = "embedding" + retriever._embedding_dimension = 3 + retriever._filterable_properties = [] + return retriever + + +# ── Initialization ───────────────────────────────────────────────────────────── + +def test_async_vector_retriever_invalid_driver() -> None: + with pytest.raises(RetrieverInitializationError): + AsyncVectorRetriever(driver="not-a-driver", index_name="my-index") # type: ignore + + +def test_async_vector_retriever_invalid_index_name(async_driver: MagicMock) -> None: + with pytest.raises(RetrieverInitializationError): + AsyncVectorRetriever(driver=async_driver, index_name=42) # type: ignore + + +def test_async_vector_retriever_init_ok(async_driver: MagicMock) -> None: + retriever = AsyncVectorRetriever(driver=async_driver, index_name="my-index") + assert retriever.index_name == "my-index" + + +# ── async_init ───────────────────────────────────────────────────────────────── + +@pytest.mark.asyncio +async def test_async_vector_retriever_async_init(async_driver: MagicMock) -> None: + mock_record = MagicMock() + mock_record.__getitem__ = lambda self, key: { + "labels": ["Document"], + "properties": ["embedding"], + "dimensions": 3, + "filterable_properties": [], + }[key] + async_driver.execute_query.return_value = MagicMock(records=[mock_record]) + + retriever = AsyncVectorRetriever(driver=async_driver, index_name="my-index") + result = await retriever.async_init() + assert result is retriever + async_driver.execute_query.assert_called_once() + + +# ── search ───────────────────────────────────────────────────────────────────── + +@pytest.mark.asyncio +async def test_async_vector_retriever_search_with_vector(async_driver: MagicMock) -> None: + mock_record = MagicMock(spec=neo4j.Record) + mock_record.get = MagicMock(return_value=None) + async_driver.execute_query.return_value = MagicMock(records=[mock_record]) + + retriever = _make_retriever(async_driver) + with patch("neo4j_graphrag.retrievers.async_vector.supports_search_clause", return_value=False): + result = await retriever.search(query_vector=[0.1, 0.2, 0.3], top_k=2) + + assert isinstance(result, RetrieverResult) + async_driver.execute_query.assert_called_once() + + +@pytest.mark.asyncio +async def test_async_vector_retriever_search_with_text(async_driver: MagicMock, embedder: MagicMock) -> None: + mock_record = MagicMock(spec=neo4j.Record) + mock_record.get = MagicMock(return_value=None) + async_driver.execute_query.return_value = MagicMock(records=[mock_record]) + + retriever = _make_retriever(async_driver, embedder=embedder) + with patch("neo4j_graphrag.retrievers.async_vector.supports_search_clause", return_value=False): + result = await retriever.search(query_text="find something", top_k=3) + + assert isinstance(result, RetrieverResult) + embedder.embed_query.assert_called_once_with("find something") + + +@pytest.mark.asyncio +async def test_async_vector_retriever_search_no_embedder_raises(async_driver: MagicMock) -> None: + retriever = _make_retriever(async_driver) + with patch("neo4j_graphrag.retrievers.async_vector.supports_search_clause", return_value=False): + with pytest.raises(EmbeddingRequiredError): + await retriever.search(query_text="find something") + + +@pytest.mark.asyncio +async def test_async_vector_retriever_search_no_query_raises(async_driver: MagicMock) -> None: + retriever = _make_retriever(async_driver) + with pytest.raises(SearchValidationError): + await retriever.search() + + +# ── AsyncVectorCypherRetriever ───────────────────────────────────────────────── + +def test_async_vector_cypher_retriever_init_ok(async_driver: MagicMock) -> None: + retriever = AsyncVectorCypherRetriever( + driver=async_driver, + index_name="my-index", + retrieval_query="RETURN node.id AS id", + ) + assert retriever.index_name == "my-index" + assert retriever.retrieval_query == "RETURN node.id AS id" + + +def test_async_vector_cypher_retriever_invalid_driver() -> None: + with pytest.raises(RetrieverInitializationError): + AsyncVectorCypherRetriever( + driver="not-a-driver", # type: ignore + index_name="my-index", + retrieval_query="RETURN node", + ) + + +@pytest.mark.asyncio +async def test_async_vector_cypher_retriever_search(async_driver: MagicMock, embedder: MagicMock) -> None: + mock_record = MagicMock(spec=neo4j.Record) + mock_record.get = MagicMock(return_value=None) + async_driver.execute_query.return_value = MagicMock(records=[mock_record]) + + retriever = AsyncVectorCypherRetriever( + driver=async_driver, + index_name="my-index", + retrieval_query="RETURN node.id AS id", + embedder=embedder, + ) + retriever._node_label = "Document" + retriever._node_embedding_property = "embedding" + retriever._embedding_dimension = 3 + retriever._filterable_properties = [] + + with patch("neo4j_graphrag.retrievers.async_vector.supports_search_clause", return_value=False): + result = await retriever.search(query_vector=[0.1, 0.2, 0.3], top_k=2) + + assert isinstance(result, RetrieverResult) + + +# ── default_record_formatter ─────────────────────────────────────────────────── + +def test_async_vector_retriever_default_formatter(async_driver: MagicMock) -> None: + retriever = _make_retriever(async_driver) + record = MagicMock(spec=neo4j.Record) + record.get = MagicMock(side_effect=lambda k: {"score": 0.9, "node": "test-node", "nodeLabels": ["Doc"], "id": "1"}.get(k)) + item = retriever.default_record_formatter(record) + assert isinstance(item, RetrieverResultItem) + assert "test-node" in item.content + assert item.metadata["score"] == 0.9