...
/Storing LLM Extracted Entities and Relationships in Neo4j
Storing LLM Extracted Entities and Relationships in Neo4j
Learn how to write modular code to store a large number of LLM-extracted entities and relationships for our dataset in Neo4j.
We'll cover the following...
In the previous lesson, we covered how to store two hard-coded entities and the relationship between them in Neo4j. Now, we’ll move forward by storing an entire list of entities and relationships extracted from our question-answer dataset.
A modular implementation: Storing data to Neo4j
If we look at the code below, the only new file we see is neo4j_db.py. We'll understand how this file manages the storage of entities and relationships in Neo4j. In this file, we define a class, Neo4jHandler, which is used to manage interactions with a Neo4j database. The class allows us to store entities and relationships extracted from our dataset into Neo4j in a structured and efficient way.
from neo4j import GraphDatabase
class Neo4jHandler:
def __init__(self, connection_uri, username, password):
self.driver = GraphDatabase.driver(connection_uri, auth=(username, password))
def close(self):
self.driver.close()
def store_entities_relationships(self, entities, relationships):
with self.driver.session() as session:
# Store entities
for entity in entities:
session.write_transaction(self._create_entity_node, entity)
# Store relationships
for entity1, relationship, entity2 in relationships:
session.write_transaction(self._create_relationship, entity1, relationship, entity2)
@staticmethod
def _create_entity_node(tx, entity):
query = "MERGE (e:Entity {name: $entity_name})"
tx.run(query, entity_name=entity)
@staticmethod
def _create_relationship(tx, entity1, relationship, entity2):
query = """
MATCH (e1:Entity {name: $entity1})
MATCH (e2:Entity {name: $entity2})
MERGE (e1)-[:RELATION {type: $relationship}]->(e2)
"""
tx.run(query, entity1=entity1, entity2=entity2, relationship=relationship)
Here’s a breakdown of what each part of the code does:
Lines 4–5: We write the constructor method ...