Skip to content

Commit

Permalink
Add specific driver closing to avoid issues with neo4
Browse files Browse the repository at this point in the history
  • Loading branch information
Vasilije1990 committed Jan 5, 2024
1 parent f7d2dd4 commit 0120421
Showing 1 changed file with 31 additions and 0 deletions.
31 changes: 31 additions & 0 deletions main.py
Original file line number Diff line number Diff line change
Expand Up @@ -184,6 +184,7 @@ async def user_query_to_graph_db(session: AsyncSession, user_id: str, query_inpu
neo4j_graph_db = Neo4jGraphDB(url=config.graph_database_url, username=config.graph_database_username, password=config.graph_database_password)
cypher_query = await neo4j_graph_db.generate_cypher_query_for_user_prompt_decomposition(user_id,query_input)
result = neo4j_graph_db.query(cypher_query)
neo4j_graph_db.close()

await update_entity(session, Operation, job_id, "SUCCESS")

Expand Down Expand Up @@ -281,27 +282,41 @@ async def add_documents_to_graph_db(session: AsyncSession, user_id: str= None, d
if document_memory_types == ['PUBLIC']:
await create_public_memory(user_id=user_id, labels=['sr'], topic="PublicMemory")
ids = neo4j_graph_db.retrieve_node_id_for_memory_type(topic="PublicMemory")
neo4j_graph_db.close()
print(ids)
else:
ids = neo4j_graph_db.retrieve_node_id_for_memory_type(topic="SemanticMemory")
neo4j_graph_db.close()
print(ids)

for id in ids:
print(id.get('memoryId'))
neo4j_graph_db = Neo4jGraphDB(url=config.graph_database_url, username=config.graph_database_username,
password=config.graph_database_password)
if document_memory_types == ['PUBLIC']:

rs = neo4j_graph_db.create_document_node_cypher(classification, user_id, public_memory_id=id.get('memoryId'))
neo4j_graph_db.close()
else:
rs = neo4j_graph_db.create_document_node_cypher(classification, user_id, memory_type='SemanticMemory')
neo4j_graph_db.close()
logging.info("Cypher query is", rs)
neo4j_graph_db = Neo4jGraphDB(url=config.graph_database_url, username=config.graph_database_username,
password=config.graph_database_password)
neo4j_graph_db.query(rs)
neo4j_graph_db.close()
logging.info("WE GOT HERE")
neo4j_graph_db = Neo4jGraphDB(url=config.graph_database_url, username=config.graph_database_username,
password=config.graph_database_password)
if memory_details[0][1] == "PUBLIC":

neo4j_graph_db.update_document_node_with_db_ids( vectordb_namespace=memory_details[0][0],
document_id=doc_id)
neo4j_graph_db.close()
else:
neo4j_graph_db.update_document_node_with_db_ids( vectordb_namespace=memory_details[0][0],
document_id=doc_id, user_id=user_id)
neo4j_graph_db.close()
# await update_entity_graph_summary(session, DocsModel, doc_id, True)
except Exception as e:
return e
Expand Down Expand Up @@ -366,7 +381,11 @@ async def user_context_enrichment(session, user_id:str, query:str, generative_re
await user_query_to_graph_db(session, user_id, query)

semantic_mem = neo4j_graph_db.retrieve_semantic_memory(user_id=user_id)
neo4j_graph_db.close()
neo4j_graph_db = Neo4jGraphDB(url=config.graph_database_url, username=config.graph_database_username,
password=config.graph_database_password)
episodic_mem = neo4j_graph_db.retrieve_episodic_memory(user_id=user_id)
neo4j_graph_db.close()
# public_mem = neo4j_graph_db.retrieve_public_memory(user_id=user_id)


Expand Down Expand Up @@ -485,6 +504,7 @@ async def create_public_memory(user_id: str=None, labels:list=None, topic:str=No
# Assuming the topic for public memory is predefined, e.g., "PublicMemory"
# Create the memory node
memory_id = await neo4j_graph_db.create_memory_node(labels=labels, topic=topic)
neo4j_graph_db.close()
return memory_id
except Neo4jError as e:
logging.error(f"Error creating public memory node: {e}")
Expand Down Expand Up @@ -522,9 +542,14 @@ async def attach_user_to_memory(user_id: str=None, labels:list=None, topic:str=N

# Assuming the topic for public memory is predefined, e.g., "PublicMemory"
ids = neo4j_graph_db.retrieve_node_id_for_memory_type(topic=topic)
neo4j_graph_db.close()

for id in ids:
neo4j_graph_db = Neo4jGraphDB(url=config.graph_database_url,
username=config.graph_database_username,
password=config.graph_database_password)
linked_memory = neo4j_graph_db.link_public_memory_to_user(memory_id=id.get('memoryId'), user_id=user_id)
neo4j_graph_db.close()
return 1
except Neo4jError as e:
logging.error(f"Error creating public memory node: {e}")
Expand Down Expand Up @@ -561,8 +586,14 @@ async def unlink_user_from_memory(user_id: str=None, labels:list=None, topic:str

# Assuming the topic for public memory is predefined, e.g., "PublicMemory"
ids = neo4j_graph_db.retrieve_node_id_for_memory_type(topic=topic)
neo4j_graph_db.close()

for id in ids:
neo4j_graph_db = Neo4jGraphDB(url=config.graph_database_url,
username=config.graph_database_username,
password=config.graph_database_password)
linked_memory = neo4j_graph_db.unlink_memory_from_user(memory_id=id.get('memoryId'), user_id=user_id)
neo4j_graph_db.close()
return 1
except Neo4jError as e:
logging.error(f"Error creating public memory node: {e}")
Expand Down

0 comments on commit 0120421

Please sign in to comment.