A program graphrag_clinical_demo.py that demonstrates Causal Link Discovery. It shows how GraphRAG identifies that a patient's symptom is actually a side effect of their current medication—a connection standard RAG often misses

graphrag_clinical_demo.py

#!/usr/bin/env python3
"""
GraphRAG Clinical Demo: Causal Link Discovery
This script demonstrates how GraphRAG identifies drug-symptom causal links 
that standard RAG might miss due to document fragmentation.
"""
import os
import json

class SimpleGraph:
    def __init__(self):
        self.nodes = {}
        self.edges = []

    def add_node(self, id, type, properties):
        self.nodes[id] = {"type": type, "properties": properties}

    def add_edge(self, source, target, type):
        self.edges.append({"source": source, "target": target, "type": type})

    def get_related(self, source_id, edge_type=None):
        related = []
        for edge in self.edges:
            if edge["source"] == source_id:
                if edge_type is None or edge["type"] == edge_type:
                    related.append(edge["target"])
        return related

# 1. Initialize Clinical Knowledge Graph
kg = SimpleGraph()

# Add Entities
kg.add_node("P-101", "Patient", {"name": "John Doe", "age": 55})
kg.add_node("Med-Lisinopril", "Medication", {"name": "Lisinopril", "class": "ACE Inhibitor"})
kg.add_node("Cond-HTN", "Condition", {"name": "Hypertension"})
kg.add_node("Symp-Cough", "Symptom", {"name": "Persistent Dry Cough", "severity": "Moderate"})

# Add Relationships
kg.add_edge("P-101", "Cond-HTN", "HAS_CONDITION")
kg.add_edge("P-101", "Med-Lisinopril", "TAKES_MEDICATION")
kg.add_edge("Med-Lisinopril", "Cond-HTN", "TREATS")
kg.add_edge("Med-Lisinopril", "Symp-Cough", "CAUSES_SIDE_EFFECT")
kg.add_edge("P-101", "Symp-Cough", "EXHIBITS_SYMPTOM")

def synthesize_answer(query, facts, mode="GraphRAG"):
    api_key = os.environ.get("GOOGLE_API_KEY")
    if not api_key:
        print(f"\n[AI Synthesis ({mode})]: (Skipped - No GOOGLE_API_KEY found)")
        return

    try:
        import google.generativeai as genai
        genai.configure(api_key=api_key)
        model = genai.GenerativeModel('gemini-2.0-flash')
        
        system_instruction = "You are a Clinical Decision Support Assistant."
        if mode == "RAG":
            system_instruction += " You are provided with disconnected clinical notes. Be cautious about assuming causality."
        
        prompt = f"""
        {system_instruction}
        Based on the following data, answer the user query.
        
        DATA:
        {facts}
        
        QUERY: {query}
        
        Provide a concise clinical recommendation.
        """
        
        response = model.generate_content(prompt)
        print(f"\n[AI Synthesis Result ({mode})]:")
        print(f"  {response.text.strip()}")
    except Exception as e:
        print(f"\n[AI Synthesis Error ({mode})]: {e}")

def simulate_standard_rag(query):
    print(f"\n[Standard RAG] Query: '{query}'")
    print("Finding semantically similar clinical chunks...")
    
    chunks = [
        "Chunk 1: Patient John Doe (55) has a history of hypertension.",
        "Chunk 2: Patient is currently prescribed Lisinopril for blood pressure management.",
        "Chunk 3: John Doe reported a persistent dry cough during today's visit.",
        "Chunk 4: Common side effects of ACE inhibitors like Lisinopril include dry cough."
    ]
    
    for chunk in chunks:
        print(f"  ------- Retrieved {chunk}")
        
    facts = "\n".join(chunks)
    synthesize_answer(query, facts, mode="RAG")

def simulate_graph_rag(query):
    print(f"\n[GraphRAG] Query: '{query}'")
    print("Traversing clinical relationships...")
    
    # 1. Identify Patient Symptoms
    patient_id = "P-101"
    symptoms = kg.get_related(patient_id, "EXHIBITS_SYMPTOM")
    meds = kg.get_related(patient_id, "TAKES_MEDICATION")
    
    traversed_facts = []
    print(f"  1. Identify Patient: John Doe exhibits {len(symptoms)} symptoms.")
    
    for symp in symptoms:
        symp_name = kg.nodes[symp]["properties"]["name"]
        print(f"  2. Symptom: {symp_name}")
        
        for med in meds:
            med_name = kg.nodes[med]["properties"]["name"]
            # Check if this med CAUSES this symptom
            if symp in kg.get_related(med, "CAUSES_SIDE_EFFECT"):
                print(f"  3. Found Causal Link: {med_name} -> [CAUSES_SIDE_EFFECT] -> {symp_name}")
                traversed_facts.append(f"Patient takes {med_name}. {med_name} CAUSES {symp_name} as a side effect.")

    facts = "\n".join(traversed_facts)
    synthesize_answer(query, facts, mode="GraphRAG")

def visualize_graph():
    try:
        import networkx as nx
        import matplotlib.pyplot as plt

        print("\n[Visualization] Generating Clinical Graph Visualization...")
        G = nx.DiGraph()

        for node_id, data in kg.nodes.items():
            G.add_node(node_id, label=f"{data['properties'].get('name', node_id)}\n({data['type']})")

        for edge in kg.edges:
            G.add_edge(edge["source"], edge["target"], label=edge["type"])

        pos = nx.spring_layout(G, seed=42)
        plt.figure(figsize=(10, 7))
        nx.draw(G, pos, with_labels=True, 
                node_color='lightcoral', node_size=3000, 
                font_size=10, font_weight='bold', 
                labels=nx.get_node_attributes(G, 'label'),
                arrows=True)
        
        edge_labels = nx.get_edge_attributes(G, 'label')
        nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_color='blue')

        plt.title("Clinical GraphRAG: Medication & Symptom Links")
        output_path = "clinical_graph_visualization.png"
        plt.savefig(output_path)
        print(f"  Success: Visualization saved to {output_path}")
        
    except ImportError:
        print("\n[Visualization] Skipping: 'networkx' or 'matplotlib' not installed.")

if __name__ == "__main__":
    query = "Patient John Doe has a persistent dry cough. Should we prescribe a cough suppressant or evaluate his current medication?"
    
    simulate_standard_rag(query)
    print("-" * 50)
    simulate_graph_rag(query)
    print("-" * 50)
    visualize_graph()