ops/scripts/seed.py

#!/usr/bin/env python3
"""
Seed script to create demo data for the agentic system.
Creates a sample workflow with instructions for testing.
"""

import requests
import json
import time
import os

# Configuration
MCP_URL = os.getenv("MCP_URL", "http://localhost:8000/mcp")
API_KEY = os.getenv("MCP_API_KEY", "dev-key-123")

def call_mcp(tool, params=None):
    """Call the MCP API"""
    response = requests.post(
        MCP_URL,
        headers={"X-API-Key": API_KEY, "Content-Type": "application/json"},
        json={"tool": tool, "params": params or {}}
    )
    return response.json()

def create_demo_workflow():
    """Create a demo workflow for testing"""
    print("🌱 Seeding demo data...")
    
    # Create demo workflow
    workflow_result = call_mcp("write_graph", {
        "action": "create_node",
        "label": "Workflow",
        "properties": {
            "id": "demo-workflow-1",
            "name": "Entity Resolution Demo",
            "description": "Demo workflow showing SQL generation and execution",
            "status": "active",
            "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ")
        }
    })
    
    if 'created' not in workflow_result:
        print(f"❌ Failed to create workflow: {workflow_result}")
        return False
    
    print("✅ Created demo workflow")
    
    # Create demo instructions
    instructions = [
        {
            "id": "demo-inst-1",
            "type": "discover_schema",
            "sequence": 1,
            "description": "Discover database schema",
            "parameters": "{}"
        },
        {
            "id": "demo-inst-2", 
            "type": "generate_sql",
            "sequence": 2,
            "description": "Generate SQL for customer orders",
            "parameters": json.dumps({
                "question": "Show me all customers who have placed orders, including their order details"
            })
        },
        {
            "id": "demo-inst-3",
            "type": "review_results", 
            "sequence": 3,
            "description": "Review and format results",
            "parameters": "{}"
        }
    ]
    
    for inst in instructions:
        inst_result = call_mcp("write_graph", {
            "action": "create_node",
            "label": "Instruction",
            "properties": {
                "id": inst["id"],
                "type": inst["type"],
                "sequence": inst["sequence"],
                "description": inst["description"],
                "status": "pending",
                "pause_duration": 300,  # 5 minutes
                "parameters": inst["parameters"],
                "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ")
            }
        })
        
        if 'created' not in inst_result:
            print(f"❌ Failed to create instruction {inst['id']}: {inst_result}")
            continue
            
        # Link instruction to workflow
        link_result = call_mcp("query_graph", {
            "query": "MATCH (w:Workflow), (i:Instruction) WHERE w.id = wid AND i.id = iid CREATE (w)-[:HAS_INSTRUCTION]->(i)",
            "parameters": {"wid": "demo-workflow-1", "iid": inst["id"]}
        })
        
        print(f"✅ Created instruction: {inst['type']}")
    
    # Create instruction chain
    for i in range(len(instructions) - 1):
        current_id = instructions[i]["id"]
        next_id = instructions[i + 1]["id"]
        
        chain_result = call_mcp("query_graph", {
            "query": "MATCH (i1:Instruction), (i2:Instruction) WHERE i1.id = id1 AND i2.id = id2 CREATE (i1)-[:NEXT_INSTRUCTION]->(i2)",
            "parameters": {"id1": current_id, "id2": next_id}
        })
    
    print("✅ Created instruction chain")
    
    # Create query templates for demo
    templates = [
        {
            "id": "template-1",
            "question": "How many customers do we have?",
            "query": "SELECT COUNT(*) as customer_count FROM customers"
        },
        {
            "id": "template-2", 
            "question": "Show me recent orders",
            "query": "SELECT o.id, o.order_date, c.name FROM orders o JOIN customers c ON o.customer_id = c.id ORDER BY o.order_date DESC LIMIT 10"
        },
        {
            "id": "template-3",
            "question": "What's the total revenue?", 
            "query": "SELECT SUM(total_amount) as total_revenue FROM orders"
        }
    ]
    
    for template in templates:
        template_result = call_mcp("write_graph", {
            "action": "create_node",
            "label": "QueryTemplate",
            "properties": {
                "id": template["id"],
                "question": template["question"],
                "query": template["query"],
                "created_at": time.strftime("%Y-%m-%dT%H:%M:%SZ")
            }
        })
    
    print("✅ Created query templates")
    
    # Verify the setup
    verification = call_mcp("query_graph", {
        "query": """
            MATCH (w:Workflow {id: 'demo-workflow-1'})-[:HAS_INSTRUCTION]->(i:Instruction)
            RETURN w.name, count(i) as instruction_count
        """,
        "parameters": {}
    })
    
    if verification.get('data'):
        workflow_name = verification['data'][0]['w.name']
        instruction_count = verification['data'][0]['instruction_count']
        print(f"✅ Verification: '{workflow_name}' with {instruction_count} instructions")
    
    print("🎉 Demo data seeded successfully!")
    print("\n📋 Next steps:")
    print("1. Open http://localhost:3000 in your browser")
    print("2. Ask a question like: 'Show me all customers who have placed orders'")
    print("3. Watch the agent process the workflow")
    print("4. During the 5-minute pause, you can edit instructions in Neo4j Browser")
    print("5. Neo4j Browser: http://localhost:7474 (username: neo4j, password: password)")
    
    return True

def check_services():
    """Check if required services are available"""
    try:
        # Check MCP health
        health_response = requests.get(f"{MCP_URL.replace('/mcp', '/health')}", timeout=5)
        if health_response.status_code != 200:
            print("❌ MCP service not available")
            return False
            
        # Check MCP API
        test_response = requests.post(
            MCP_URL,
            headers={"X-API-Key": API_KEY, "Content-Type": "application/json"},
            json={"tool": "get_schema"},
            timeout=5
        )
        if test_response.status_code != 200:
            print("❌ MCP API not responding")
            return False
            
        print("✅ Services are available")
        return True
        
    except Exception as e:
        print(f"❌ Service check failed: {e}")
        return False

if __name__ == "__main__":
    print("🚀 Starting seed process...")
    
    if not check_services():
        print("❌ Cannot proceed - services not available")
        exit(1)
    
    if create_demo_workflow():
        print("✅ Seed completed successfully")
        exit(0)
    else:
        print("❌ Seed failed")
        exit(1)