A program graphrag_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_demo.py

#!/ env python3
"""
GraphRAG Practical Demo: Fleet Management Scenario
This script demonstrates the difference between Vector-based RAG and Graph-based RAG.
"""

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 the Knowledge Graph
kg = SimpleGraph()

# Add Entities
kg.add_node("Sarah", "Driver", {"name": "Sarah Connor", "experience": "10 years"})
kg.add_node("Truck42", "Vehicle", {"model": "Heavy Duty 2024", "status": "Maintenance Overdue"})
kg.add_node("ZoneAlpha", "Zone", {"name": "Zone Alpha", "security_level": "High"})
kg.add_node("WeatherWarning", "Event", {"type": "Severe Weather", "severity": "Critical"})

# Add Relationships (Edges)
kg.add_edge("Sarah", "Truck42", "DRIVES")
kg.add_edge("Sarah", "ZoneAlpha", "CURRENTLY_IN")
kg.add_edge("WeatherWarning", "ZoneAlpha", "AFFECTS")

def simulate_standard_rag(query):
    print(f"\n[Standard RAG] Query: '{query}'")
    print("Finding semantically similar chunks...")
    
    # Mocking Vector Search Results (Unconnected chunks)
    chunks = [
        "Chunk 1: Driver Sarah Connor is a veteran with 10 years experience.",
        "Chunk 2: Vehicle Truck42 has missed its safety inspection last month.",
        "Chunk 3: Severe Weather Alert issued for Zone Alpha - Extreme Caution Advised."
    ]
    
    for i, chunk in enumerate(chunks):
        print(f"  ------- Retrieved {chunk}")
        
    print("\n[RAG Result]: The LLM has the pieces, but it's 'guessing' if they are connected.")
    print("Does Sarah drive Truck42? Is Sarah in Zone Alpha? It's not explicitly stated in any single chunk.")
    
    # Synthesis Step for RAG
    facts = "\n".join(chunks)
    synthesize_answer(query, facts, mode="RAG")

def simulate_graph_rag(query):
    print(f"\n[GraphRAG] Query: '{query}'")
    print("Traversing edges for deterministic relationships...")
    
    # Starting traversal from the risk
    risk_zone = kg.get_related("WeatherWarning", "AFFECTS")[0]
    zone_name = kg.nodes[risk_zone]["properties"]["name"]
    print(f"  1. Start at Event: Severe Weather -> Affects: {zone_name}")
    
    # Find drivers in that zone
    drivers = []
    for node_id, node_data in kg.nodes.items():
        if node_data["type"] == "Driver":
            if risk_zone in kg.get_related(node_id, "CURRENTLY_IN"):
                drivers.append(node_id)
    
    for driver in drivers:
        driver_name = kg.nodes[driver]["properties"]["name"]
        print(f"  2. Find Drivers in {zone_name} -> Found: {driver_name}")
        
        # Check their vehicles
        vehicles = kg.get_related(driver, "DRIVES")
        for vehicle in vehicles:
            status = kg.nodes[vehicle]["properties"]["status"]
            print(f"  3. Check Driver's Vehicle: {vehicle} -> Status: {status}")
            
            if "Overdue" in status:
                print(f"\n[GraphRAG Success]: CRITICAL ALERT.")
                print(f"  {driver_name} is in a {zone_name} (RISK) driving {vehicle} (UNSAFE).")
                
                # Synthesis Step for GraphRAG
                facts = f"Driver: {driver_name}, Zone: {zone_name}, Risk: {kg.nodes['WeatherWarning']['properties']['type']}, Vehicle: {vehicle}, Status: {status}"
                synthesize_answer(query, facts, mode="GraphRAG")

def synthesize_answer(query, facts, mode="GraphRAG"):
    """Optional: Use Gemini to synthesize the final response if API key is available."""
    import os
    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)
        # Using the model that worked in your environment
        model = genai.GenerativeModel('gemini-2.5-flash')
        
        system_instruction = "You are a Fleet Management Assistant."
        if mode == "RAG":
            system_instruction += " You are provided with disconnected text chunks. If you cannot confirm a connection, be honest about it."
        
        prompt = f"""
        {system_instruction}
        Based on the following data, answer the user query.
        
        DATA:
        {facts}
        
        QUERY: {query}
        
        Provide a concise, professional alert for the dispatcher.
        """
        
        print(f"\n[AI Prompt ({mode})]:")
        print("-" * 20)
        print(prompt.strip())
        print("-" * 20)
        
        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 visualize_graph():
    """Optional: Visualize the graph using networkx and matplotlib."""
    try:
        import networkx as nx
        import matplotlib.pyplot as plt

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

        # Add nodes with their types as labels
        for node_id, data in kg.nodes.items():
            G.add_node(node_id, label=f"{node_id}\n({data['type']})")

        # Add edges
        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))
        
        # Draw nodes
        nx.draw(G, pos, with_labels=True, 
                node_color='skyblue', node_size=3000, 
                font_size=10, font_weight='bold', 
                labels=nx.get_node_attributes(G, 'label'),
                arrows=True)
        
        # Draw edge labels
        edge_labels = nx.get_edge_attributes(G, 'label')
        nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels, font_color='red')

        plt.title("GraphRAG Knowledge Graph: Fleet Management")
        
        # Save the visualization
        output_path = "graph_visualization.png"
        plt.savefig(output_path)
        print(f"  Success: Visualization saved to {output_path}")
        
        # In a local environment with a display, you could use plt.show()
        # plt.show()
        
    except ImportError:
        print("\n[Visualization] Skipping: 'networkx' or 'matplotlib' not installed.")
        print("  Try: pip install networkx matplotlib")
    except Exception as e:
        print(f"\n[Visualization] Error: {e}")

if __name__ == "__main__":
    query = "Find drivers in high-risk zones with unsafe vehicles."
    
    simulate_standard_rag(query)
    print("-" * 50)
    simulate_graph_rag(query)
    print("-" * 50)
    visualize_graph()