demo.py

#!/usr/bin/env python3
"""
LightRAG + ArangoDB Demo Script
Demonstrates the complete integration for MediaTek POC
"""

import os
import sys
import asyncio
from pathlib import Path
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Add the current directory to path to import arangodb_impl
sys.path.insert(0, str(Path(__file__).parent))

# Import LightRAG and ArangoDB storage
try:
    from lightrag import LightRAG, QueryParam
    from lightrag.llm.openai import gpt_4o_mini_complete, openai_embed
    print("✅ Successfully imported LightRAG and ArangoDB storage")
except ImportError as e:
    print(f"❌ Import error: {e}")
    print("Please ensure lightrag-hku is installed: pip install lightrag-hku")
    sys.exit(1)


async def setup_lightrag():
    """Initialize LightRAG with ArangoDB storage"""
    
    # Configuration
    working_dir = os.getenv("LIGHTRAG_WORKING_DIR", "./working_dir")
    namespace = os.getenv("LIGHTRAG_NAMESPACE", "mediatek")
    
    # Create working directory
    Path(working_dir).mkdir(exist_ok=True)
    
    print(f"\n🚀 Initializing LightRAG with ArangoDB")
    print(f"   Namespace: {namespace}")
    print(f"   Working Directory: {working_dir}")
    print(f"   ArangoDB Host: {os.getenv('ARANGO_HOST', 'http://arangodb:8529')}")
    
    # Initialize LightRAG with ArangoDB storage
    # Note: Embedding dimension is auto-detected from the model in ArangoDBVectorStorage
    rag = LightRAG(
        working_dir=working_dir,
        
        # LLM and embedding functions (required)
        llm_model_func=gpt_4o_mini_complete,
        embedding_func=openai_embed,
        
        # Use ArangoDB for all storage types
        graph_storage="ArangoDBStorage",
        kv_storage="ArangoDBKVStorage", 
        vector_storage="ArangoDBVectorStorage",
        doc_status_storage="ArangoDBDocStatusStorage",
    )
    
    # Initialize storage backends
    await rag.initialize_storages()
    
    print("✅ LightRAG initialized successfully\n")
    return rag


async def ingest_sample_data(rag):
    """Ingest sample documents into the knowledge graph"""
    
    print("📄 Ingesting sample documents...")
    
    # Sample documents about AI and databases
    documents = [
        """
        ArangoDB is a native multi-model database with flexible data models for documents, 
        graphs, and key-value pairs. It has a unified query language (AQL) for all data models.
        ArangoDB is designed for high performance and scalability in production environments.
        """,
        
        """
        LightRAG is a retrieval-augmented generation framework that combines knowledge graphs
        with vector embeddings for intelligent question answering. It supports multiple storage
        backends including ArangoDB for enterprise deployments.
        """,
        
        """
        MediaTek is a leading semiconductor company specializing in chipsets for mobile devices,
        smart home products, and automotive applications. The company focuses on AI-powered
        technologies and edge computing solutions.
        """,
        
        """
        Graph databases excel at representing and querying complex relationships between entities.
        Unlike traditional relational databases, graph databases can efficiently traverse
        relationships without expensive JOIN operations.
        """
    ]
    
    # Insert documents
    for i, doc in enumerate(documents, 1):
        print(f"   Ingesting document {i}/{len(documents)}...")
        await rag.ainsert(doc.strip())
    
    print("✅ Sample data ingested successfully\n")


async def run_sample_queries(rag):
    """Run sample queries to demonstrate functionality"""
    
    queries = [
        "What is ArangoDB?",
        "How does LightRAG work?",
        "What does MediaTek specialize in?",
        "Explain the advantages of graph databases"
    ]
    
    print("🔍 Running sample queries...\n")
    
    for query in queries:
        print(f"Q: {query}")
        print("-" * 80)
        
        # Query with different modes
        try:
            # Naive mode (simple vector search)
            result = await rag.aquery(query, param=QueryParam(mode="naive"))
            print(f"A (Naive): {result}\n")
            
            # Hybrid mode (combines graph + vector)
            # result = await rag.aquery(query, param=QueryParam(mode="hybrid"))
            # print(f"A (Hybrid): {result}\n")
            
        except Exception as e:
            print(f"❌ Error querying: {e}\n")
    
    print("✅ Queries completed\n")


async def show_statistics(rag):
    """Display statistics about the knowledge graph"""
    
    print("📊 Knowledge Graph Statistics")
    print("=" * 80)
    
    try:
        # Get workspace stats from ArangoDB
        # Use the actual storage instance (chunk_entity_relation_graph), not the class
        if hasattr(rag, 'chunk_entity_relation_graph') and hasattr(rag.chunk_entity_relation_graph, 'get_workspace_stats'):
            stats = await rag.chunk_entity_relation_graph.get_workspace_stats()
            
            print(f"Graph Nodes: {stats.get('nodes', 0)}")
            print(f"Graph Edges: {stats.get('edges', 0)}")
            print(f"KV Store Items: {stats.get('kv_store', 0)}")
            print(f"Vector Embeddings: {stats.get('vectors', 0)}")
            print(f"Document Status: {stats.get('doc_status', 0)}")
        else:
            print("Statistics not available")
    
    except Exception as e:
        print(f"❌ Error getting statistics: {e}")
    
    print("=" * 80 + "\n")


async def main():
    """Main demonstration workflow"""
    
    print("\n" + "=" * 80)
    print("  LightRAG + ArangoDB Integration Demo")
    print("  MediaTek POC - January 2026")
    print("=" * 80 + "\n")
    
    # Check environment variables
    required_vars = ["ARANGO_HOST", "ARANGO_USERNAME", "ARANGO_PASSWORD"]
    missing_vars = [var for var in required_vars if not os.getenv(var)]
    
    if missing_vars:
        print(f"❌ Missing environment variables: {', '.join(missing_vars)}")
        print("Please create a .env file or set these variables\n")
        return
    
    # Clear the database for a fresh start (ensures repeatable demo)
    from arango import ArangoClient
    db_name = os.environ.get("ARANGO_DATABASE", "lightrag")
    arango_host = os.environ.get("ARANGO_HOST", "http://arangodb:8529")
    arango_password = os.environ.get("ARANGO_PASSWORD", "openSesame")
    
    print(f"🗄️  Preparing fresh '{db_name}' database...")
    client = ArangoClient(hosts=arango_host)
    sys_db = client.db("_system", username="root", password=arango_password)
    
    if sys_db.has_database(db_name):
        sys_db.delete_database(db_name)
        print(f"   🗑️  Cleared existing database")
    sys_db.create_database(db_name)
    print(f"   ✅ Created fresh '{db_name}' database\n")
    
    try:
        # Initialize LightRAG
        rag = await setup_lightrag()
        
        # Ingest sample data
        await ingest_sample_data(rag)
        
        # Run sample queries
        await run_sample_queries(rag)
        
        # Show statistics
        await show_statistics(rag)
        
        print("✅ Demo completed successfully!")
        print("\n💡 Next steps:")
        print("   1. Access ArangoDB web interface at http://localhost:8529")
        print("   2. Run test_queries.py for performance testing")
        print("   3. Customize with your own documents and queries\n")
    
    except Exception as e:
        print(f"\n❌ Demo failed: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    asyncio.run(main())