import json
from neo4j import GraphDatabase
# --- Assume extracted_data is the JSON output from the previous step ---
# extracted_data = { "nodes": [...], "relationships": [...] }
# --- 1. Neo4j Connection Details ---
# Replace with your Neo4j AuraDB credentials or local instance details
URI = "neo4j://localhost:7687"
AUTH = ("neo4j", "your_password")
# --- 2. Define the GraphBuilder Class ---
class GraphBuilder:
def __init__(self, uri, auth):
self.driver = GraphDatabase.driver(uri, auth=auth)
def close(self):
self.driver.close()
def build_graph(self, graph_data):
nodes = graph_data.get("nodes", [])
relationships = graph_data.get("relationships", [])
with self.driver.session() as session:
# Use MERGE for nodes to avoid creating duplicates
# It will create a node if it doesn't exist, or match it if it does.
for node in nodes:
session.run("""
MERGE (n:%s {id: $id})
SET n += $properties
""" % node['type'], id=node['id'], properties=node['properties'])
print(f"Created/Merged Node: {node['id']}")
# Use MATCH for source/target nodes and MERGE for the relationship
for rel in relationships:
session.run("""
MATCH (source {id: $source_id})
MATCH (target {id: $target_id})
MERGE (source)-[r:%s]->(target)
SET r += $properties
""" % rel['type'], source_id=rel['source'], target_id=rel['target'], properties=rel.get('properties', {}))
print(f"Created/Merged Relationship: {rel['source']} -> {rel['target']}")
# --- 3. Run the Graph Building Process ---
if __name__ == "__main__":
builder = GraphBuilder(URI, AUTH)
# In a real application, you would pass the 'extracted_data' variable here
# For this example, we'll use a hardcoded version of the expected data
sample_extracted_data = {
"nodes": [{"id": "Sarah Chen", "type": "Person", "properties": {"name": "Sarah Chen", "title": "CEO"}}, {"id": "Nexus Innovations", "type": "Organization", "properties": {"name": "Nexus Innovations", "founded_year": 2021}}, {"id": "Tom Gomez", "type": "Person", "properties": {"name": "Tom Gomez", "title": "COO"}}, {"id": "QuantumLeap Inc.", "type": "Organization", "properties": {"name": "QuantumLeap Inc."}}],
"relationships": [{"source": "Sarah Chen", "target": "Nexus Innovations", "type": "WORKS_AT", "properties": {"role": "CEO"}}, {"source": "Tom Gomez", "target": "Nexus Innovations", "type": "WORKS_AT", "properties": {"role": "COO"}}]
}
builder.build_graph(sample_extracted_data)
builder.close()
print("\nKnowledge graph construction complete!")