Use the neo4j_database parameter everywhere (#216)

* Use self.neo4j_database for all queries in Neo4jWriter

* Make sure all execute_query can be run against a custom database

* Update CHANGELOG

* Update docstring + update examples not to use undocumented feature for neo4j driver

* Expose neo4j_database in SimpleKGBuilder

* Update CHANGELOG

* Simplify changelog

Author: stellasia

CHANGELOG.md

@@ -3,8 +3,15 @@
 ## Next
 
 ### Added
-- Introduced optional lexical graph configuration for SimpleKGPipeline, enhancing flexibility in customizing node labels and relationship types in the lexical graph.
-- Ability to provide description and list of properties for entities and relations in the SimpleKGPipeline constructor.
+- Introduced optional lexical graph configuration for `SimpleKGPipeline`, enhancing flexibility in customizing node labels and relationship types in the lexical graph.
+- Introduced optional `neo4j_database` parameter for `SimpleKGPipeline`, `Neo4jChunkReader`and `Text2CypherRetriever`.
+- Ability to provide description and list of properties for entities and relations in the `SimpleKGPipeline` constructor.
+
+### Fixed
+- `neo4j_database` parameter is now used for all queries in the `Neo4jWriter`.
+
+### Changed
+- Updated all examples to use `neo4j_database` parameter instead of an undocumented neo4j driver constructor.
 
 ## 1.2.0
 

examples/build_graph/simple_kg_builder_from_pdf.py

@@ -50,6 +50,7 @@ async def define_and_run_pipeline(
         entities=ENTITIES,
         relations=RELATIONS,
         potential_schema=POTENTIAL_SCHEMA,
+        neo4j_database=DATABASE,
     )
     return await kg_builder.run_async(file_path=str(file_path))
 
@@ -62,7 +63,7 @@ async def main() -> PipelineResult:
             "response_format": {"type": "json_object"},
         },
     )
-    with neo4j.GraphDatabase.driver(URI, auth=AUTH, database=DATABASE) as driver:
+    with neo4j.GraphDatabase.driver(URI, auth=AUTH) as driver:
         res = await define_and_run_pipeline(driver, llm)
     await llm.async_client.close()
     return res

examples/build_graph/simple_kg_builder_from_text.py

@@ -21,7 +21,7 @@
 # Neo4j db infos
 URI = "neo4j://localhost:7687"
 AUTH = ("neo4j", "password")
-DATABASE = "neo4j"
+DATABASE = "newdb"
 
 # Text to process
 TEXT = """The son of Duke Leto Atreides and the Lady Jessica, Paul is the heir of House Atreides,
@@ -67,6 +67,7 @@ async def define_and_run_pipeline(
         relations=RELATIONS,
         potential_schema=POTENTIAL_SCHEMA,
         from_pdf=False,
+        neo4j_database=DATABASE,
     )
     return await kg_builder.run_async(text=TEXT)
 
@@ -79,7 +80,7 @@ async def main() -> PipelineResult:
             "response_format": {"type": "json_object"},
         },
     )
-    with neo4j.GraphDatabase.driver(URI, auth=AUTH, database=DATABASE) as driver:
+    with neo4j.GraphDatabase.driver(URI, auth=AUTH) as driver:
         res = await define_and_run_pipeline(driver, llm)
     await llm.async_client.close()
     return res

examples/customize/answer/custom_prompt.py

@@ -23,7 +23,6 @@
 driver = neo4j.GraphDatabase.driver(
     URI,
     auth=AUTH,
-    database=DATABASE,
 )
 
 embedder = OpenAIEmbeddings()
@@ -33,6 +32,7 @@
     index_name=INDEX,
     retrieval_query="WITH node, score RETURN node.title as title, node.plot as plot",
     embedder=embedder,
+    neo4j_database=DATABASE,
 )
 
 llm = OpenAILLM(model_name="gpt-4o", model_params={"temperature": 0})

examples/customize/answer/langchain_compatiblity.py

@@ -21,7 +21,6 @@
 driver = neo4j.GraphDatabase.driver(
     URI,
     auth=AUTH,
-    database=DATABASE,
 )
 
 embedder = OpenAIEmbeddings(model="text-embedding-ada-002")
@@ -31,6 +30,7 @@
     index_name=INDEX,
     retrieval_query="WITH node, score RETURN node.title as title, node.plot as plot",
     embedder=embedder,  # type: ignore[arg-type, unused-ignore]
+    neo4j_database=DATABASE,
 )
 
 llm = ChatOpenAI(model="gpt-4o", temperature=0)

examples/customize/retrievers/result_formatter_vector_cypher_retriever.py

@@ -38,7 +38,7 @@ def my_result_formatter(record: neo4j.Record) -> RetrieverResultItem:
     )
 
 
-with neo4j.GraphDatabase.driver(URI, auth=AUTH, database=DATABASE) as driver:
+with neo4j.GraphDatabase.driver(URI, auth=AUTH) as driver:
     # Initialize the retriever
     retriever = VectorCypherRetriever(
         driver=driver,
@@ -48,7 +48,7 @@ def my_result_formatter(record: neo4j.Record) -> RetrieverResultItem:
         retrieval_query=RETRIEVAL_QUERY,
         result_formatter=my_result_formatter,
         # optionally, set neo4j database
-        # neo4j_database="neo4j",
+        neo4j_database=DATABASE,
     )
 
     # Perform the similarity search for a text query

examples/customize/retrievers/result_formatter_vector_retriever.py

@@ -35,7 +35,7 @@
 
 
 # Connect to Neo4j database
-driver = neo4j.GraphDatabase.driver(URI, auth=AUTH, database=DATABASE)
+driver = neo4j.GraphDatabase.driver(URI, auth=AUTH)
 
 
 query_text = "Find a movie about astronauts"
@@ -52,6 +52,7 @@
     index_name=INDEX_NAME,
     embedder=OpenAIEmbeddings(),
     return_properties=["title", "plot"],
+    neo4j_database=DATABASE,
 )
 print(retriever.search(query_text=query_text, top_k=top_k_results))
 print()

examples/question_answering/graphrag.py

@@ -35,7 +35,6 @@ def formatter(record: neo4j.Record) -> RetrieverResultItem:
 driver = neo4j.GraphDatabase.driver(
     URI,
     auth=AUTH,
-    database=DATABASE,
 )
 
 embedder = OpenAIEmbeddings()
@@ -46,6 +45,7 @@ def formatter(record: neo4j.Record) -> RetrieverResultItem:
     retrieval_query="with node, score return node.title as title, node.plot as plot",
     result_formatter=formatter,
     embedder=embedder,
+    neo4j_database=DATABASE,
 )
 
 llm = OpenAILLM(model_name="gpt-4o", model_params={"temperature": 0})

examples/retrieve/hybrid_cypher_retriever.py

@@ -24,7 +24,7 @@
 # the name of all actors starring in that movie
 RETRIEVAL_QUERY = " MATCH (node)<-[:ACTED_IN]-(p:Person) RETURN node.title as movieTitle, node.plot as moviePlot, collect(p.name) as actors, score as similarityScore"
 
-with neo4j.GraphDatabase.driver(URI, auth=AUTH, database=DATABASE) as driver:
+with neo4j.GraphDatabase.driver(URI, auth=AUTH) as driver:
     # Initialize the retriever
     retriever = HybridCypherRetriever(
         driver=driver,
@@ -37,7 +37,7 @@
         # (see corresponding example in 'customize' directory)
         # result_formatter=None,
         # optionally, set neo4j database
-        # neo4j_database="neo4j",
+        neo4j_database=DATABASE,
     )
 
     # Perform the similarity search for a text query

examples/retrieve/hybrid_retriever.py

@@ -18,7 +18,7 @@
 FULLTEXT_INDEX_NAME = "movieFulltext"
 
 
-with neo4j.GraphDatabase.driver(URI, auth=AUTH, database=DATABASE) as driver:
+with neo4j.GraphDatabase.driver(URI, auth=AUTH) as driver:
     # Initialize the retriever
     retriever = HybridRetriever(
         driver=driver,
@@ -31,7 +31,7 @@
         # (see corresponding example in 'customize' directory)
         # result_formatter=None,
         # optionally, set neo4j database
-        # neo4j_database="neo4j",
+        neo4j_database=DATABASE,
     )
 
     # Perform the similarity search for a text query