Allow specifying database name for retrievers and index handling methods (#71)

* Include database name in driver.execute_query

* Update tests

* Added configurable database names to the Pinecone and Weaviate retrievers

* Fix typing for filtered vector queries (#76)

* Rename database to neo4j_database

---------

Co-authored-by: Jon Besga <jon.besga@neo4j.com>

Author: willtai

CHANGELOG.md

@@ -12,12 +12,14 @@
 -   Added LLMInterface with implementation for OpenAI LLM.
 -   Updated project configuration to support multiple Python versions (3.8 to 3.12) in CI workflows.
 -   Improved developer experience by copying the docstring from the `Retriever.get_search_results` method to the `Retriever.search` method
+-   Support for specifying database names in index handling methods and retrievers.
 
 ### Changed
 -   Refactored import paths for retrievers to neo4j_genai.retrievers.
 -   Implemented exception chaining for all re-raised exceptions to improve stack trace readability.
 -   Made error messages in `index.py` more consistent.
 -   Renamed `Retriever._get_search_results` to `Retriever.get_search_results`
+-   Updated retrievers and index handling methods to accept optional database names.
 
 ## 0.2.0
 

src/neo4j_genai/indexes.py

@@ -15,7 +15,7 @@
 from __future__ import annotations
 
 import logging
-from typing import Literal
+from typing import Literal, Optional
 
 import neo4j
 from pydantic import ValidationError
@@ -33,6 +33,7 @@ def create_vector_index(
     embedding_property: str,
     dimensions: int,
     similarity_fn: Literal["euclidean", "cosine"],
+    neo4j_database: Optional[str] = None,
 ) -> None:
     """
     This method constructs a Cypher query and executes it
@@ -77,6 +78,7 @@ def create_vector_index(
         dimensions (int): Vector embedding dimension
         similarity_fn (str): case-insensitive values for the vector similarity function:
             ``euclidean`` or ``cosine``.
+        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>`_).
 
     Raises:
         ValueError: If validation of the input arguments fail.
@@ -105,13 +107,18 @@ def create_vector_index(
         driver.execute_query(
             query,
             {"name": name, "dimensions": dimensions, "similarity_fn": similarity_fn},
+            database_=neo4j_database,
         )
     except neo4j.exceptions.ClientError as e:
         raise Neo4jIndexError(f"Neo4j vector index creation failed: {e.message}") from e
 
 
 def create_fulltext_index(
-    driver: neo4j.Driver, name: str, label: str, node_properties: list[str]
+    driver: neo4j.Driver,
+    name: str,
+    label: str,
+    node_properties: list[str],
+    neo4j_database: Optional[str] = None,
 ) -> None:
     """
     This method constructs a Cypher query and executes it
@@ -151,6 +158,7 @@ def create_fulltext_index(
         name (str): The unique name of the index.
         label (str): The node label to be indexed.
         node_properties (list[str]): The node properties to create the fulltext index on.
+        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>`_).
 
     Raises:
         ValueError: If validation of the input arguments fail.
@@ -172,14 +180,16 @@ def create_fulltext_index(
             f"[{', '.join(['n.`' + prop + '`' for prop in node_properties])}]"
         )
         logger.info(f"Creating fulltext index named '{name}'")
-        driver.execute_query(query, {"name": name})
+        driver.execute_query(query, {"name": name}, database_=neo4j_database)
     except neo4j.exceptions.ClientError as e:
         raise Neo4jIndexError(
             f"Neo4j fulltext index creation failed {e.message}"
         ) from e
 
 
-def drop_index_if_exists(driver: neo4j.Driver, name: str) -> None:
+def drop_index_if_exists(
+    driver: neo4j.Driver, name: str, neo4j_database: Optional[str] = None
+) -> None:
     """
     This method constructs a Cypher query and executes it
     to drop an index in Neo4j, if the index exists.
@@ -210,6 +220,7 @@ def drop_index_if_exists(driver: neo4j.Driver, name: str) -> None:
     Args:
         driver (neo4j.Driver): Neo4j Python driver instance.
         name (str): The name of the index to delete.
+        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>`_).
 
     Raises:
         neo4j.exceptions.ClientError: If dropping of index fails.
@@ -220,7 +231,7 @@ def drop_index_if_exists(driver: neo4j.Driver, name: str) -> None:
             "name": name,
         }
         logger.info(f"Dropping index named '{name}'")
-        driver.execute_query(query, parameters)
+        driver.execute_query(query, parameters, database_=neo4j_database)
     except neo4j.exceptions.ClientError as e:
         raise Neo4jIndexError(f"Dropping Neo4j index failed: {e.message}") from e
 
@@ -230,6 +241,7 @@ def upsert_vector(
     node_id: int,
     embedding_property: str,
     vector: list[float],
+    neo4j_database: Optional[str] = None,
 ) -> None:
     """
     This method constructs a Cypher query and executes it to upsert (insert or update) a vector property on a specific node.
@@ -260,6 +272,7 @@ def upsert_vector(
         node_id (int): The id of the node.
         embedding_property (str): The name of the property to store the vector in.
         vector (list[float]): The vector to store.
+        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>`_).
 
     Raises:
         Neo4jInsertionError: If upserting of the vector fails.
@@ -277,7 +290,7 @@ def upsert_vector(
             "embedding_property": embedding_property,
             "vector": vector,
         }
-        driver.execute_query(query, parameters)
+        driver.execute_query(query, parameters, database_=neo4j_database)
     except neo4j.exceptions.ClientError as e:
         raise Neo4jInsertionError(
             f"Upserting vector to Neo4j failed: {e.message}"

src/neo4j_genai/retrievers/base.py

@@ -81,8 +81,9 @@ class Retriever(ABC, metaclass=RetrieverMetaclass):
     index_name: str
     VERIFY_NEO4J_VERSION = True
 
-    def __init__(self, driver: neo4j.Driver):
+    def __init__(self, driver: neo4j.Driver, neo4j_database: Optional[str] = None):
         self.driver = driver
+        self.neo4j_database = neo4j_database
         if self.VERIFY_NEO4J_VERSION:
             self._verify_version()
 
@@ -95,7 +96,9 @@ def _verify_version(self) -> None:
         indexing. Raises a Neo4jMinVersionError if the connected Neo4j version is
         not supported.
         """
-        records, _, _ = self.driver.execute_query("CALL dbms.components()")
+        records, _, _ = self.driver.execute_query(
+            "CALL dbms.components()", database_=self.neo4j_database
+        )
         version = records[0]["versions"][0]
 
         if "aura" in version:
@@ -120,7 +123,9 @@ def _fetch_index_infos(self) -> None:
             "RETURN labelsOrTypes as labels, properties, "
             "options.indexConfig.`vector.dimensions` as dimensions"
         )
-        query_result = self.driver.execute_query(query, {"index_name": self.index_name})
+        query_result = self.driver.execute_query(
+            query, {"index_name": self.index_name}, database_=self.neo4j_database
+        )
         try:
             result = query_result.records[0]
             self._node_label = result["labels"][0]
@@ -185,11 +190,16 @@ class ExternalRetriever(Retriever, ABC):
     VERIFY_NEO4J_VERSION = False
 
     def __init__(
-        self, driver: neo4j.Driver, id_property_external: str, id_property_neo4j: str
+        self,
+        driver: neo4j.Driver,
+        id_property_external: str,
+        id_property_neo4j: str,
+        neo4j_database: Optional[str] = None,
     ):
         super().__init__(driver)
         self.id_property_external = id_property_external
         self.id_property_neo4j = id_property_neo4j
+        self.neo4j_database = neo4j_database
 
     @abstractmethod
     def get_search_results(

src/neo4j_genai/retrievers/external/pinecone/pinecone.py

@@ -79,6 +79,7 @@ class PineconeNeo4jRetriever(ExternalRetriever):
         embedder (Optional[Embedder]): Embedder object to embed query text.
         return_properties (Optional[list[str]]): List of node properties to return.
         result_formatter (Optional[Callable[[Any], Any]]): Function to transform a neo4j.Record to a RetrieverResultItem.
+        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>`_).
 
     Raises:
         RetrieverInitializationError: If validation of the input arguments fail.
@@ -94,6 +95,7 @@ def __init__(
         return_properties: Optional[list[str]] = None,
         retrieval_query: Optional[str] = None,
         result_formatter: Optional[Callable[[Any], Any]] = None,
+        neo4j_database: Optional[str] = None,
     ):
         try:
             driver_model = Neo4jDriverModel(driver=driver)
@@ -108,6 +110,7 @@ def __init__(
                 return_properties=return_properties,
                 retrieval_query=retrieval_query,
                 result_formatter=result_formatter,
+                neo4j_database=neo4j_database,
             )
         except ValidationError as e:
             raise RetrieverInitializationError(e.errors()) from e
@@ -116,6 +119,7 @@ def __init__(
             driver=driver,
             id_property_external="id",
             id_property_neo4j=validated_data.id_property_neo4j,
+            neo4j_database=neo4j_database,
         )
         self.driver = validated_data.driver_model.driver
         self.client = validated_data.client_model.client
@@ -224,6 +228,8 @@ def get_search_results(
         logger.debug("Pinecone Store Cypher parameters: %s", parameters)
         logger.debug("Pinecone Store Cypher query: %s", search_query)
 
-        records, _, _ = self.driver.execute_query(search_query, parameters)
+        records, _, _ = self.driver.execute_query(
+            search_query, parameters, database_=self.neo4j_database
+        )
 
         return RawSearchResult(records=records)

src/neo4j_genai/retrievers/external/pinecone/types.py

@@ -53,3 +53,4 @@ class PineconeNeo4jRetrieverModel(BaseModel):
     return_properties: Optional[list[str]] = None
     retrieval_query: Optional[str] = None
     result_formatter: Optional[Callable[[neo4j.Record], str]] = None
+    neo4j_database: Optional[str] = None

src/neo4j_genai/retrievers/external/weaviate/types.py

@@ -51,6 +51,7 @@ class WeaviateNeo4jRetrieverModel(BaseModel):
     return_properties: Optional[list[str]] = None
     retrieval_query: Optional[str] = None
     result_formatter: Optional[Callable[[neo4j.Record], str]] = None
+    neo4j_database: Optional[str] = None
 
 
 class WeaviateNeo4jSearchModel(VectorSearchModel):

src/neo4j_genai/retrievers/external/weaviate/weaviate.py

@@ -70,6 +70,7 @@ class WeaviateNeo4jRetriever(ExternalRetriever):
         embedder (Optional[Embedder]): Embedder object to embed query text.
         return_properties (Optional[list[str]]): List of node properties to return.
         result_formatter (Optional[Callable[[Any], Any]]): Function to transform a neo4j.Record to a RetrieverResultItem.
+        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>`_).
 
     Raises:
         RetrieverInitializationError: If validation of the input arguments fail.
@@ -86,6 +87,7 @@ def __init__(
         return_properties: Optional[list[str]] = None,
         retrieval_query: Optional[str] = None,
         result_formatter: Optional[Callable[[Any], Any]] = None,
+        neo4j_database: Optional[str] = None,
     ):
         try:
             driver_model = Neo4jDriverModel(driver=driver)
@@ -101,11 +103,14 @@ def __init__(
                 return_properties=return_properties,
                 retrieval_query=retrieval_query,
                 result_formatter=result_formatter,
+                neo4j_database=neo4j_database,
             )
         except ValidationError as e:
             raise RetrieverInitializationError(e.errors()) from e
 
-        super().__init__(driver, id_property_external, id_property_neo4j)
+        super().__init__(
+            driver, id_property_external, id_property_neo4j, neo4j_database
+        )
         self.client = validated_data.client_model.client
         collection = validated_data.collection
         self.search_collection = self.client.collections.get(collection)
@@ -227,6 +232,8 @@ def get_search_results(
         logger.debug("Weaviate Store Cypher parameters: %s", parameters)
         logger.debug("Weaviate Store Cypher query: %s", search_query)
 
-        records, _, _ = self.driver.execute_query(search_query, parameters)
+        records, _, _ = self.driver.execute_query(
+            search_query, parameters, database_=self.neo4j_database
+        )
 
         return RawSearchResult(records=records)

src/neo4j_genai/retrievers/hybrid.py

@@ -69,6 +69,8 @@ class HybridRetriever(Retriever):
         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.
+        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>`_).
+
     """
 
     def __init__(
@@ -78,6 +80,7 @@ def __init__(
         fulltext_index_name: str,
         embedder: Optional[Embedder] = None,
         return_properties: Optional[list[str]] = None,
+        neo4j_database: Optional[str] = None,
     ) -> None:
         try:
             driver_model = Neo4jDriverModel(driver=driver)
@@ -88,11 +91,14 @@ def __init__(
                 fulltext_index_name=fulltext_index_name,
                 embedder_model=embedder_model,
                 return_properties=return_properties,
+                neo4j_database=neo4j_database,
             )
         except ValidationError as e:
             raise RetrieverInitializationError(e.errors()) from e
 
-        super().__init__(validated_data.driver_model.driver)
+        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
@@ -173,7 +179,9 @@ def get_search_results(
         logger.debug("HybridRetriever Cypher parameters: %s", parameters)
         logger.debug("HybridRetriever Cypher query: %s", search_query)
 
-        records, _, _ = self.driver.execute_query(search_query, parameters)
+        records, _, _ = self.driver.execute_query(
+            search_query, parameters, database_=self.neo4j_database
+        )
         return RawSearchResult(
             records=records,
         )
@@ -212,6 +220,7 @@ class HybridCypherRetriever(Retriever):
         retrieval_query (str): Cypher query that gets appended.
         embedder (Optional[Embedder]): Embedder object to embed query text.
         result_formatter (Optional[Callable[[Any], Any]]): Provided custom function to transform a neo4j.Record to a RetrieverResultItem.
+        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>`_).
 
     Raises:
         RetrieverInitializationError: If validation of the input arguments fail.
@@ -225,6 +234,7 @@ def __init__(
         retrieval_query: str,
         embedder: Optional[Embedder] = None,
         result_formatter: Optional[Callable[[Any], Any]] = None,
+        neo4j_database: Optional[str] = None,
     ) -> None:
         try:
             driver_model = Neo4jDriverModel(driver=driver)
@@ -235,11 +245,14 @@ def __init__(
                 fulltext_index_name=fulltext_index_name,
                 retrieval_query=retrieval_query,
                 embedder_model=embedder_model,
+                neo4j_database=neo4j_database,
             )
         except ValidationError as e:
             raise RetrieverInitializationError(e.errors()) from e
 
-        super().__init__(validated_data.driver_model.driver)
+        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
@@ -316,7 +329,9 @@ def get_search_results(
         logger.debug("HybridCypherRetriever Cypher parameters: %s", parameters)
         logger.debug("HybridCypherRetriever Cypher query: %s", search_query)
 
-        records, _, _ = self.driver.execute_query(search_query, parameters)
+        records, _, _ = self.driver.execute_query(
+            search_query, parameters, database_=self.neo4j_database
+        )
         return RawSearchResult(
             records=records,
         )