Update ki_gen/data_retriever.py

ki_gen/data_retriever.py

@@ -166,7 +166,11 @@ def generate_cypher(state: DocRetrieverState, config: ConfigSchema):
     """
     The node where the cypher is generated
     """
-    graph = config["configurable"].get("graph")
+    #graph = config["configurable"].get("graph")
+    NEO4J_URI = "neo4j+s://4985272f.databases.neo4j.io"
+    NEO4J_USERNAME = "neo4j"
+    NEO4J_PASSWORD = os.getenv("neo4j_password")
+    graph = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USERNAME, NEO4J_PASSWORD))
     question = state['query']
     related_concepts = get_related_concepts(graph, question)
     cyphers = []
@@ -209,7 +213,7 @@ def generate_cypher(state: DocRetrieverState, config: ConfigSchema):
         else:
             print("Warning: Cypher validation skipped, graph or schema unavailable.")
 
-
+    graph.close()
     return {"cyphers" : cyphers}
 
 # ... (generate_cypher_from_topic, get_docs remain the same)
@@ -233,7 +237,11 @@ def get_docs(state:DocRetrieverState, config:ConfigSchema):
     """
     This node retrieves docs from the graph using the generated cypher
     """
-    graph = config["configurable"].get("graph")
+    #graph = config["configurable"].get("graph")
+    NEO4J_URI = "neo4j+s://4985272f.databases.neo4j.io"
+    NEO4J_USERNAME = "neo4j"
+    NEO4J_PASSWORD = os.getenv("neo4j_password")
+    graph = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USERNAME, NEO4J_PASSWORD))
     output = []
     if graph is not None and state.get("cyphers"): # Check if cyphers exist
         for cypher in state["cyphers"]:
@@ -281,7 +289,7 @@ def get_docs(state:DocRetrieverState, config:ConfigSchema):
                   filtered_docs.append(doc)
                   seen_docs.add(doc_str)
 
-
+    graph.close()
     return {"docs": filtered_docs}