#!/usr/bin/env python3
"""
Setup script for initializing a basic cardiology ontology in Neo4j.
Run this after setting up your Neo4j database to get started with some base concepts.
"""
import os
from neo4j import GraphDatabase
from typing import List, Dict
# Connection details
URI = os.getenv("NEO4J_URI", "bolt://localhost:7687")
USER = os.getenv("NEO4J_USERNAME", "neo4j")
PASS = os.getenv("NEO4J_PASSWORD", "password")
def create_cardiology_ontology():
"""Create basic cardiology entities and relationships."""
# Basic cardiology entities
entities = [
# Anatomical structures
{"name": "heart", "label": "Anatomy", "properties": {"type": "organ"}},
{"name": "left ventricle", "label": "Anatomy", "properties": {"type": "chamber"}},
{"name": "right ventricle", "label": "Anatomy", "properties": {"type": "chamber"}},
{"name": "left atrium", "label": "Anatomy", "properties": {"type": "chamber"}},
{"name": "right atrium", "label": "Anatomy", "properties": {"type": "chamber"}},
{"name": "aortic valve", "label": "Anatomy", "properties": {"type": "valve"}},
{"name": "mitral valve", "label": "Anatomy", "properties": {"type": "valve"}},
{"name": "tricuspid valve", "label": "Anatomy", "properties": {"type": "valve"}},
{"name": "pulmonary valve", "label": "Anatomy", "properties": {"type": "valve"}},
{"name": "coronary arteries", "label": "Anatomy", "properties": {"type": "blood vessel"}},
# Physiological processes
{"name": "cardiac cycle", "label": "Process", "properties": {"type": "physiological"}},
{"name": "systole", "label": "Process", "properties": {"type": "cardiac phase"}},
{"name": "diastole", "label": "Process", "properties": {"type": "cardiac phase"}},
{"name": "atrial contraction", "label": "Process", "properties": {"type": "cardiac event"}},
{"name": "ventricular contraction", "label": "Process", "properties": {"type": "cardiac event"}},
{"name": "electrical conduction", "label": "Process", "properties": {"type": "cardiac event"}},
# Common conditions
{"name": "atrial fibrillation", "label": "Condition", "properties": {"type": "arrhythmia"}},
{"name": "heart failure", "label": "Condition", "properties": {"type": "syndrome"}},
{"name": "myocardial infarction", "label": "Condition", "properties": {"type": "acute event"}},
{"name": "hypertension", "label": "Condition", "properties": {"type": "chronic"}},
{"name": "coronary artery disease", "label": "Condition", "properties": {"type": "vascular"}},
{"name": "left ventricular hypertrophy", "label": "Condition", "properties": {"type": "structural"}},
# Common medications
{"name": "metoprolol", "label": "Medication", "properties": {"class": "beta-blocker"}},
{"name": "lisinopril", "label": "Medication", "properties": {"class": "ACE inhibitor"}},
{"name": "atorvastatin", "label": "Medication", "properties": {"class": "statin"}},
{"name": "warfarin", "label": "Medication", "properties": {"class": "anticoagulant"}},
{"name": "aspirin", "label": "Medication", "properties": {"class": "antiplatelet"}},
# Diagnostic procedures
{"name": "echocardiogram", "label": "Procedure", "properties": {"type": "imaging"}},
{"name": "electrocardiogram", "label": "Procedure", "properties": {"type": "diagnostic"}},
{"name": "cardiac catheterization", "label": "Procedure", "properties": {"type": "invasive"}},
{"name": "stress test", "label": "Procedure", "properties": {"type": "functional"}},
{"name": "chest X-ray", "label": "Procedure", "properties": {"type": "imaging"}},
]
# Basic relationships
relationships = [
# Anatomical relationships
{"source": "left ventricle", "target": "heart", "type": "PART_OF"},
{"source": "right ventricle", "target": "heart", "type": "PART_OF"},
{"source": "left atrium", "target": "heart", "type": "PART_OF"},
{"source": "right atrium", "target": "heart", "type": "PART_OF"},
{"source": "aortic valve", "target": "heart", "type": "PART_OF"},
{"source": "mitral valve", "target": "heart", "type": "PART_OF"},
{"source": "tricuspid valve", "target": "heart", "type": "PART_OF"},
{"source": "pulmonary valve", "target": "heart", "type": "PART_OF"},
# Process relationships
{"source": "systole", "target": "cardiac cycle", "type": "PART_OF"},
{"source": "diastole", "target": "cardiac cycle", "type": "PART_OF"},
{"source": "atrial contraction", "target": "systole", "type": "PART_OF"},
{"source": "ventricular contraction", "target": "systole", "type": "PART_OF"},
{"source": "diastole", "target": "systole", "type": "FOLLOWS"},
# Condition relationships
{"source": "hypertension", "target": "left ventricular hypertrophy", "type": "CAUSES"},
{"source": "coronary artery disease", "target": "myocardial infarction", "type": "CAUSES"},
{"source": "myocardial infarction", "target": "heart failure", "type": "CAUSES"},
{"source": "atrial fibrillation", "target": "heart failure", "type": "CAUSES"},
# Treatment relationships
{"source": "metoprolol", "target": "atrial fibrillation", "type": "TREATS"},
{"source": "metoprolol", "target": "hypertension", "type": "TREATS"},
{"source": "lisinopril", "target": "hypertension", "type": "TREATS"},
{"source": "lisinopril", "target": "heart failure", "type": "TREATS"},
{"source": "atorvastatin", "target": "coronary artery disease", "type": "TREATS"},
{"source": "warfarin", "target": "atrial fibrillation", "type": "TREATS"},
{"source": "aspirin", "target": "coronary artery disease", "type": "TREATS"},
# Diagnostic relationships
{"source": "echocardiogram", "target": "heart failure", "type": "DIAGNOSES"},
{"source": "electrocardiogram", "target": "atrial fibrillation", "type": "DIAGNOSES"},
{"source": "cardiac catheterization", "target": "coronary artery disease", "type": "DIAGNOSES"},
{"source": "echocardiogram", "target": "left ventricular hypertrophy", "type": "DIAGNOSES"},
# Effects relationships
{"source": "hypertension", "target": "left ventricle", "type": "AFFECTS"},
{"source": "myocardial infarction", "target": "left ventricle", "type": "AFFECTS"},
{"source": "atrial fibrillation", "target": "atrial contraction", "type": "AFFECTS"},
]
with GraphDatabase.driver(URI, auth=(USER, PASS)) as driver:
print("🏗️ Setting up cardiology ontology...")
# Clear existing data (optional)
print("Clearing existing data...")
driver.execute_query("MATCH (n) DETACH DELETE n")
# Create entities
print(f"Creating {len(entities)} entities...")
for entity in entities:
cypher = f"""
CREATE (n:{entity['label']} {{name: $name}})
SET n += $properties
"""
driver.execute_query(
cypher,
name=entity['name'],
properties=entity['properties']
)
# Create relationships
print(f"Creating {len(relationships)} relationships...")
for rel in relationships:
cypher = f"""
MATCH (source {{name: $source_name}})
MATCH (target {{name: $target_name}})
CREATE (source)-[r:{rel['type']}]->(target)
"""
driver.execute_query(
cypher,
source_name=rel['source'],
target_name=rel['target']
)
# Create some useful indexes
print("Creating indexes...")
try:
driver.execute_query("CREATE INDEX entity_name_index FOR (n) ON (n.name)")
except:
pass # Index might already exist
# Get final counts
result = driver.execute_query("MATCH (n) RETURN count(n) as node_count")
node_count = result.records[0]["node_count"]
result = driver.execute_query("MATCH ()-[r]->() RETURN count(r) as rel_count")
rel_count = result.records[0]["rel_count"]
print(f"✅ Cardiology ontology setup complete!")
print(f" 📊 Created {node_count} nodes and {rel_count} relationships")
print(f" 🎯 Ready for document ingestion and querying!")
if __name__ == "__main__":
print("🫀 Cardiology Knowledge Graph Setup")
print("=" * 40)
try:
create_cardiology_ontology()
except Exception as e:
print(f"❌ Error setting up ontology: {str(e)}")
print("\nTroubleshooting:")
print("1. Make sure Neo4j is running")
print("2. Check your connection details in .env")
print("3. Verify database credentials")
print("\nNext steps:")
print("1. Start your MCP server: python mcp_server.py")
print("2. Configure Claude Desktop with your MCP server")
print("3. Try ingesting a cardiology document!")
