demo_scenario.py

import sys
import os
import shutil
import time

# Add src to path
sys.path.append(os.path.join(os.getcwd(), "src"))

from src.cortex.storage.manager import StorageManager
from src.cortex.visualization.generator import DashboardGenerator
from src.cortex.config import config

# Initialize logging
logger = config.setup_logging()

def print_section(title):
    print("\n" + "="*60)
    print(f" {title}")
    print("="*60)

def run_demo():
    print_section("INITIALIZING CORTEX DEMO")
    
    # Clean slate - close any existing logging handlers
    import logging
    log = logging.getLogger('cortex')
    for handler in log.handlers[:]:
        handler.close()
        log.removeHandler(handler)
    
    if config.STORAGE_DIR.exists():
        print("Cleaning up previous storage...")
        import time
        import gc
        gc.collect()
        time.sleep(0.2)
        try:
            shutil.rmtree(config.STORAGE_DIR)
        except PermissionError:
            print("Warning: Could not remove old storage (files in use), continuing...")
    
    config.ensure_dirs()
    
    # Setup logging after cleanup
    logger.info("Demo started")
    
    manager = StorageManager()
    
    # --- Scenario: Multiple Projects ---
    print_section("PHASE 1: MULTI-PROJECT INCEPTION - MEMORY INGESTION")
    
    # ====== PROJECT ORION: AI Customer Support ======
    print("\n🚀 PROJECT ORION: AI-Powered Customer Support Agent")
    print("─" * 60)
    
    print("\n📋 Storing Project Orion planning memories...")
    orion_planning = [
        ("We are starting Project Orion to build a new AI-powered customer support agent.", "planning", "Orion"),
        ("Project Orion aims to reduce support ticket response time by 60%.", "objective", "Orion"),
        ("The team consists of 3 backend engineers, 2 frontend engineers, and 1 DevOps specialist.", "team", "Orion"),
        ("Timeline: MVP in 3 months, production launch in 6 months.", "timeline", "Orion"),
        ("Budget approved: $500K for initial development phase.", "budget", "Orion"),
        ("Primary stakeholder is the Customer Success department, secondary is Sales.", "stakeholders", "Orion"),
    ]
    
    for text, category, project in orion_planning:
        print(f"  ✓ [{category.upper()}] {text[:80]}...")
        manager.save_memory(text, source="planning_doc", metadata={"category": category, "project": project})
        time.sleep(0.05)
    
    print("\n🔧 Storing Orion technical architecture...")
    orion_tech = [
        ("The backend will be written in Python 3.11 using FastAPI framework for async performance.", "tech-stack", "Orion"),
        ("For the frontend, we decided to use React 18 with TypeScript for type safety.", "tech-stack", "Orion"),
        ("We chose PostgreSQL 15 for relational data storage with full-text search capabilities.", "tech-stack", "Orion"),
        ("Redis will be used for caching and session management to improve response times.", "tech-stack", "Orion"),
        ("We need to use Docker for containerization and docker-compose for local development.", "devops", "Orion"),
        ("The API will follow REST principles with OpenAPI 3.0 documentation.", "architecture", "Orion"),
        ("Authentication will use JWT tokens with RS256 signing algorithm.", "security", "Orion"),
        ("Frontend will use Material-UI component library for consistent design.", "tech-stack", "Orion"),
        ("Implementing WebSocket connections for real-time chat functionality.", "tech-stack", "Orion"),
    ]
    
    for text, category, project in orion_tech:
        print(f"  ✓ [{category.upper()}] {text[:80]}...")
        manager.save_memory(text, source="architecture_doc", metadata={"category": category, "project": project})
        time.sleep(0.05)
    
    print("\n☁️ Storing Orion DevOps information...")
    orion_devops = [
        ("Project Orion will deploy to AWS using ECS Fargate for container orchestration.", "devops", "Orion"),
        ("Infrastructure as Code will be managed with Terraform for reproducibility.", "devops", "Orion"),
        ("CI/CD pipeline will use GitHub Actions with automated testing and deployment.", "devops", "Orion"),
        ("Production environment will use multi-AZ deployment for high availability.", "infrastructure", "Orion"),
        ("Monitoring will be handled by CloudWatch with custom dashboards and alarms.", "monitoring", "Orion"),
        ("Implementing blue-green deployment strategy to minimize downtime.", "devops", "Orion"),
    ]
    
    for text, category, project in orion_devops:
        print(f"  ✓ [{category.upper()}] {text[:80]}...")
        manager.save_memory(text, source="devops_doc", metadata={"category": category, "project": project})
        time.sleep(0.05)

    print("\n🐛 Storing Orion bug reports...")
    orion_bugs = [
        ("Critical bug found: authentication API tokens expire after 5 minutes instead of 1 hour. Users getting logged out.", "bug-report", "Orion"),
        ("Performance issue: database queries taking 3+ seconds on customer search endpoint. Need indexing.", "performance", "Orion"),
        ("Security vulnerability: CORS headers allowing all origins in production. Must restrict to app domain.", "security-issue", "Orion"),
        ("Memory leak detected in WebSocket connection handler. Need to investigate.", "bug-report", "Orion"),
    ]
    
    for text, category, project in orion_bugs:
        print(f"  ⚠️  [{category.upper()}] {text[:80]}...")
        manager.save_memory(text, source="bug_tracker", metadata={"category": category, "project": project})
        time.sleep(0.05)
    
    print("\n💡 Storing Orion meeting notes...")
    orion_meetings = [
        ("Sprint Planning: Decided to prioritize authentication bug fix and database performance optimization.", "meeting", "Orion"),
        ("Architecture Review: Approved migration from monolith to microservices for better scalability.", "meeting", "Orion"),
        ("Retrospective: Team velocity increased by 20% after adopting pair programming for complex features.", "meeting", "Orion"),
        ("Daily Standup: Frontend team blocked on API spec, need backend team sync.", "meeting", "Orion"),
        ("Client demo went well, received positive feedback on the UI. Requested dark mode feature.", "meeting", "Orion"),
    ]
    
    for text, category, project in orion_meetings:
        print(f"  📝 [{category.upper()}] {text[:80]}...")
        manager.save_memory(text, source="meeting_notes", metadata={"category": category, "project": project})
        time.sleep(0.05)

    # ====== PROJECT ATLAS: Data Analytics Platform ======
    print("\n\n📊 PROJECT ATLAS: Enterprise Data Analytics Platform")
    print("─" * 60)
    
    print("\n📋 Storing Project Atlas planning memories...")
    atlas_planning = [
        ("Project Atlas is our new enterprise data analytics platform for Fortune 500 clients.", "planning", "Atlas"),
        ("Goal: Process 10TB+ of data daily with sub-second query response times.", "objective", "Atlas"),
        ("Team structure: 5 data engineers, 3 ML engineers, 2 platform engineers, 1 product manager.", "team", "Atlas"),
        ("18-month roadmap: Q1 MVP, Q2-Q3 beta testing, Q4 general availability.", "timeline", "Atlas"),
        ("Initial contracts signed with 3 pilot customers worth $2.5M ARR.", "business", "Atlas"),
    ]
    
    for text, category, project in atlas_planning:
        print(f"  ✓ [{category.upper()}] {text[:80]}...")
        manager.save_memory(text, source="planning_doc", metadata={"category": category, "project": project})
        time.sleep(0.05)
    
    print("\n🔧 Storing Atlas technical architecture...")
    atlas_tech = [
        ("Using Apache Spark 3.4 for distributed data processing and ETL pipelines.", "tech-stack", "Atlas"),
        ("Data lake built on AWS S3 with Parquet format for columnar storage.", "tech-stack", "Atlas"),
        ("Query engine: Trino (formerly Presto) for federated SQL queries across data sources.", "tech-stack", "Atlas"),
        ("Metadata catalog managed by Apache Hive Metastore with AWS Glue integration.", "tech-stack", "Atlas"),
        ("Real-time streaming pipeline using Apache Kafka and Flink.", "tech-stack", "Atlas"),
        ("ML models deployed using MLflow and served via SageMaker endpoints.", "tech-stack", "Atlas"),
        ("Frontend dashboard built with D3.js and Plotly for interactive visualizations.", "tech-stack", "Atlas"),
        ("Implementing Delta Lake for ACID transactions on data lake.", "architecture", "Atlas"),
    ]
    
    for text, category, project in atlas_tech:
        print(f"  ✓ [{category.upper()}] {text[:80]}...")
        manager.save_memory(text, source="architecture_doc", metadata={"category": category, "project": project})
        time.sleep(0.05)
    
    print("\n☁️ Storing Atlas infrastructure notes...")
    atlas_infra = [
        ("Running EMR clusters with auto-scaling based on workload. Cost optimization is critical.", "infrastructure", "Atlas"),
        ("Data encrypted at rest with KMS, in transit with TLS 1.3.", "security", "Atlas"),
        ("Multi-region replication for disaster recovery, RTO target: 4 hours.", "infrastructure", "Atlas"),
        ("Implementing data lineage tracking with Apache Atlas for compliance.", "governance", "Atlas"),
    ]
    
    for text, category, project in atlas_infra:
        print(f"  ✓ [{category.upper()}] {text[:80]}...")
        manager.save_memory(text, source="infra_doc", metadata={"category": category, "project": project})
        time.sleep(0.05)
    
    print("\n🐛 Storing Atlas issues...")
    atlas_issues = [
        ("Spark jobs failing intermittently with OutOfMemory errors on large datasets.", "bug-report", "Atlas"),
        ("Query performance degraded after schema evolution. Need to rebuild indexes.", "performance", "Atlas"),
        ("Data quality issue: 2% of records have null values in critical timestamp field.", "data-quality", "Atlas"),
    ]
    
    for text, category, project in atlas_issues:
        print(f"  ⚠️  [{category.upper()}] {text[:80]}...")
        manager.save_memory(text, source="issue_tracker", metadata={"category": category, "project": project})
        time.sleep(0.05)
    
    print("\n💡 Storing Atlas conversations...")
    atlas_convos = [
        ("Data team discussion: Should we move from batch to streaming for real-time insights? Consensus: hybrid approach.", "discussion", "Atlas"),
        ("Customer feedback session: They love the visualizations but need export to Excel feature.", "feedback", "Atlas"),
        ("Architecture debate: Snowflake vs. building custom lakehouse. Decision: custom for cost control.", "decision", "Atlas"),
    ]
    
    for text, category, project in atlas_convos:
        print(f"  📝 [{category.upper()}] {text[:80]}...")
        manager.save_memory(text, source="conversation_log", metadata={"category": category, "project": project})
        time.sleep(0.05)

    # ====== PROJECT PHOENIX: Mobile App Rebuild ======
    print("\n\n📱 PROJECT PHOENIX: Mobile App Complete Rebuild")
    print("─" * 60)
    
    print("\n📋 Storing Project Phoenix planning...")
    phoenix_planning = [
        ("Project Phoenix: Complete rewrite of our mobile app from scratch. Legacy codebase unsalvageable.", "planning", "Phoenix"),
        ("Target: 50% faster app performance, 4.5+ star rating, support iOS 15+ and Android 12+.", "objective", "Phoenix"),
        ("Team: 4 mobile engineers (2 iOS, 2 Android), 1 designer, 1 QA, 1 product owner.", "team", "Phoenix"),
        ("9-month timeline: 3 months design, 5 months dev, 1 month beta testing.", "timeline", "Phoenix"),
        ("Key requirement: Offline-first architecture for users in low-connectivity areas.", "requirement", "Phoenix"),
    ]
    
    for text, category, project in phoenix_planning:
        print(f"  ✓ [{category.upper()}] {text[:80]}...")
        manager.save_memory(text, source="planning_doc", metadata={"category": category, "project": project})
        time.sleep(0.05)
    
    print("\n🔧 Storing Phoenix technical decisions...")
    phoenix_tech = [
        ("Going native: Swift for iOS, Kotlin for Android. No React Native due to performance concerns.", "tech-stack", "Phoenix"),
        ("Shared business logic using Kotlin Multiplatform Mobile (KMM).", "tech-stack", "Phoenix"),
        ("Local database: SQLite with Room (Android) and Core Data (iOS).", "tech-stack", "Phoenix"),
        ("API client uses GraphQL with Apollo for efficient data fetching.", "tech-stack", "Phoenix"),
        ("Push notifications via Firebase Cloud Messaging for both platforms.", "tech-stack", "Phoenix"),
        ("UI design system based on Material Design 3 and iOS Human Interface Guidelines.", "design", "Phoenix"),
        ("Implementing MVVM architecture pattern for better testability.", "architecture", "Phoenix"),
    ]
    
    for text, category, project in phoenix_tech:
        print(f"  ✓ [{category.upper()}] {text[:80]}...")
        manager.save_memory(text, source="tech_doc", metadata={"category": category, "project": project})
        time.sleep(0.05)
    
    print("\n🐛 Storing Phoenix challenges...")
    phoenix_challenges = [
        ("iOS build failing on CI due to provisioning profile issues. Need DevOps help.", "bug-report", "Phoenix"),
        ("App crashes on Android 13+ when requesting location permission. Investigate new permission model.", "bug-report", "Phoenix"),
        ("Battery drain reported in beta testing. Profiling shows network requests running in background.", "performance", "Phoenix"),
        ("Design team concerned about accessibility compliance. Need WCAG audit.", "concern", "Phoenix"),
    ]
    
    for text, category, project in phoenix_challenges:
        print(f"  ⚠️  [{category.upper()}] {text[:80]}...")
        manager.save_memory(text, source="issue_log", metadata={"category": category, "project": project})
        time.sleep(0.05)
    
    print("\n💬 Storing Phoenix team conversations...")
    phoenix_convos = [
        ("Code review discussion: Should we use Combine or async/await for iOS networking? Team prefers async/await.", "discussion", "Phoenix"),
        ("Product sync: Marketing wants social sharing feature. Adding to Q2 backlog.", "meeting", "Phoenix"),
        ("Beta tester feedback: Love the new UI but miss the old widget. Plan to redesign widget.", "feedback", "Phoenix"),
        ("Team retrospective: Pair programming on complex features working well. Continue practice.", "retrospective", "Phoenix"),
        ("Technical debt discussion: Old networking layer needs refactoring. Schedule 2-week cleanup sprint.", "planning", "Phoenix"),
    ]
    
    for text, category, project in phoenix_convos:
        print(f"  📝 [{category.upper()}] {text[:80]}...")
        manager.save_memory(text, source="team_chat", metadata={"category": category, "project": project})
        time.sleep(0.05)

    # ====== CROSS-PROJECT MEMORIES ======
    print("\n\n🔗 CROSS-PROJECT: Shared Knowledge & Decisions")
    print("─" * 60)
    
    cross_project = [
        ("Company-wide decision: Standardizing on GitHub Actions for all CI/CD pipelines.", "policy", "Company"),
        ("Security policy update: All projects must implement 2FA and rotate secrets quarterly.", "security", "Company"),
        ("Engineering all-hands: Discussed hiring 10 more engineers across all teams in Q2.", "meeting", "Company"),
        ("New vendor contract signed with DataDog for observability across all products.", "vendor", "Company"),
        ("CTO directive: Every project needs comprehensive API documentation in Confluence.", "directive", "Company"),
    ]
    
    for text, category, project in cross_project:
        print(f"  ✓ [{category.upper()}] {text[:80]}...")
        manager.save_memory(text, source="company_memo", metadata={"category": category, "project": project})
        time.sleep(0.05)

    total_memories = len(orion_planning) + len(orion_tech) + len(orion_devops) + len(orion_bugs) + len(orion_meetings) + \
                     len(atlas_planning) + len(atlas_tech) + len(atlas_infra) + len(atlas_issues) + len(atlas_convos) + \
                     len(phoenix_planning) + len(phoenix_tech) + len(phoenix_challenges) + len(phoenix_convos) + \
                     len(cross_project)
    
    print(f"\n✅ Total memories stored across all projects: {total_memories}")

    # --- Scenario: Retrieval ---
    print_section("PHASE 2: INTELLIGENT RECALL - HYBRID GRAPHRAG")
    
    queries = [
        ("What is the tech stack for Project Orion?", "Retrieving Orion technology decisions..."),
        ("Tell me about the data analytics platform", "Searching for Atlas project information..."),
        ("What programming languages are we using across projects?", "Cross-project language analysis..."),
        ("Have we encountered any bugs or issues in Orion?", "Searching Orion bug tracker..."),
        ("How is Project Atlas handling data at scale?", "Retrieving Atlas architecture and infrastructure..."),
        ("What is the mobile app rebuild about?", "Searching Phoenix project details..."),
        ("Are there any security concerns across our projects?", "Security audit across all projects..."),
        ("What is the timeline for Phoenix?", "Retrieving Phoenix milestones and schedule..."),
        ("Tell me about our CI/CD practices", "Searching DevOps and deployment strategies..."),
        ("What performance issues have we identified?", "Performance analysis across projects..."),
        ("How are we handling mobile development?", "Retrieving mobile tech stack and challenges..."),
        ("What are our company-wide engineering standards?", "Searching cross-project policies..."),
        ("Tell me about Spark and data processing", "Atlas data engineering deep dive..."),
        ("What feedback did we get from users?", "User feedback compilation..."),
        ("What decisions were made in recent meetings?", "Meeting outcomes and action items..."),
    ]
    
    for q, desc in queries:
        print(f"\n{'─'*60}")
        print(f"🔍 USER QUERY: '{q}'")
        print(f"   {desc}")
        results = manager.recall(q, limit=4)
        
        if not results:
            print("   ❌ No relevant memories found.")
            continue
            
        print(f"   📊 Found {len(results)} relevant memories:\n")
        for i, res in enumerate(results, 1):
            mem = res['memory']
            topics = res.get('topics', [])
            projects = res.get('projects', [])
            score = res.get('score', 0.0)
            
            print(f"   [{i}] 📌 {mem.text}")
            print(f"       ├─ Source: {mem.source}")
            print(f"       ├─ Timestamp: {mem.timestamp.strftime('%Y-%m-%d %H:%M:%S')}")
            print(f"       ├─ Relevance Score: {score:.3f}")
            if topics:
                print(f"       ├─ Related Topics: {', '.join(topics[:5])}")
            if projects:
                print(f"       └─ Projects: {', '.join(projects)}")
            print()

    # --- System Stats ---
    print_section("PHASE 3: SYSTEM PERFORMANCE METRICS")
    
    print("\n📊 Retrieving system statistics...\n")
    
    # Get cache stats
    from src.cortex.storage.vector import VectorStore
    vector_store = manager.vector_store
    
    if hasattr(vector_store, 'embedding_cache'):
        cache = vector_store.embedding_cache
        print(f"🔹 Embedding Cache:")
        print(f"   ├─ Hits: {cache.hits}")
        print(f"   ├─ Misses: {cache.misses}")
        hit_rate = (cache.hits / (cache.hits + cache.misses) * 100) if (cache.hits + cache.misses) > 0 else 0
        print(f"   ├─ Hit Rate: {hit_rate:.1f}%")
        print(f"   └─ Cache Size: {len(cache.cache)} entries\n")
    
    if hasattr(vector_store, 'query_cache'):
        cache = vector_store.query_cache
        print(f"🔹 Query Cache:")
        print(f"   ├─ Hits: {cache.hits}")
        print(f"   ├─ Misses: {cache.misses}")
        hit_rate = (cache.hits / (cache.hits + cache.misses) * 100) if (cache.hits + cache.misses) > 0 else 0
        print(f"   ├─ Hit Rate: {hit_rate:.1f}%")
        print(f"   └─ Cache Size: {len(cache.cache)} entries\n")
    
    # Memory count
    if vector_store.table_name in vector_store.db.table_names():
        table = vector_store.db.open_table(vector_store.table_name)
        count = table.count_rows()
        print(f"🔹 Vector Store:")
        print(f"   └─ Total Memories: {count}\n")
    else:
        count = 0
        print(f"🔹 Vector Store:")
        print(f"   └─ Total Memories: 0 (table not yet created)\n")
    
    # Graph stats
    graph_store = manager.graph_store
    result = graph_store.conn.execute("MATCH (n) RETURN COUNT(n) as count")
    node_count = result.get_next()[0]
    
    result = graph_store.conn.execute("MATCH ()-[r]->() RETURN COUNT(r) as count")
    rel_count = result.get_next()[0]
    
    print(f"🔹 Graph Store:")
    print(f"   ├─ Total Nodes: {node_count}")
    print(f"   └─ Total Relationships: {rel_count}\n")

    # --- Scenario: Visualization ---
    print_section("PHASE 4: KNOWLEDGE GRAPH VISUALIZATION")
    
    print("\n🎨 Generating interactive knowledge graph dashboard...")
    viz_gen = DashboardGenerator(manager)
    path = viz_gen.generate("demo_dashboard.html")
    abs_path = os.path.abspath(path)
    
    print(f"\n✅ Dashboard generated successfully!")
    print(f"\n📁 Location: {abs_path}")
    print(f"\n🌐 Open this file in your browser to explore:")
    print(f"   • Interactive node-link diagram")
    print(f"   • Memory nodes (blue circles)")
    print(f"   • Topic nodes (green circles)")
    print(f"   • Project nodes (orange circles)")
    print(f"   • Relationships and connections")
    print(f"   • Hover over nodes for details")
    
    print_section("DEMO COMPLETE")
    print("\n🎉 Cortex multi-project demonstration finished successfully!")
    print("\n📝 Summary:")
    print(f"   • Stored {count} memories across 3 major projects + company-wide knowledge")
    print(f"   • Project Orion: AI Customer Support (planning → deployment)")
    print(f"   • Project Atlas: Data Analytics Platform (big data engineering)")
    print(f"   • Project Phoenix: Mobile App Rebuild (iOS + Android)")
    print(f"   • Created {node_count} graph nodes and {rel_count} relationships")
    print(f"   • Demonstrated hybrid semantic + graph retrieval across projects")
    print(f"   • Generated interactive visualization dashboard")
    print("\n💡 Next steps:")
    print("   1. Open demo_dashboard.html in your browser")
    print("   2. Explore different project clusters in the graph")
    print("   3. Try queries like: 'What are all our mobile technologies?'")
    print("   4. Notice how memories connect across projects via shared topics")
    print("   5. Check cortex.log for detailed operation logs")
    print("\n" + "="*60 + "\n")

if __name__ == "__main__":
    run_demo()