demo.py

#!/usr/bin/env python3
"""
Quick Demo Script for LLM Security Auditor
Run this to test the tool with your own prompts quickly.

Usage:
    python demo.py              # Normal mode
    python demo.py --verbose    # Show detailed DSPy operation logs
    python demo.py -v           # Short form of verbose
"""

import os
import sys
from pathlib import Path

# Add gateway to path
sys.path.append(os.path.dirname(os.path.abspath(__file__)))


def load_env_file():
    """Load environment variables from .env file with proper quote handling."""
    env_file = Path(".env")
    if env_file.exists():
        with open(env_file) as f:
            content = f.read()
            for line in content.split('\n'):
                line = line.strip()
                if line and not line.startswith("#") and "=" in line:
                    key, value = line.split("=", 1)
                    # Remove quotes if present
                    value = value.strip('"').strip("'")
                    os.environ[key] = value


def main():
    print("🛡️  LLM Security Auditor - Quick Demo")
    print("=" * 50)

    # Check for verbose mode argument
    verbose = "--verbose" in sys.argv or "-v" in sys.argv

    if verbose:
        print("🔍 Verbose mode enabled - showing detailed DSPy operation logs\n")

    # Check for API key and determine provider
    if os.getenv("GROQ_API_KEY"):
        provider = "Groq"
        model_name = "llama-3.3-70b-versatile"
        model = "groq"
    elif os.getenv("OPENAI_API_KEY"):
        provider = "OpenAI"
        model_name = "gpt-4o-mini"
        model = "openai"
    elif os.getenv("ANTHROPIC_API_KEY"):
        provider = "Anthropic"
        model_name = "claude-sonnet-4-20250514"
        model = "anthropic"
    else:
        print("❌ Error: No API key found!")
        print("Please set one of: GROQ_API_KEY, OPENAI_API_KEY, or ANTHROPIC_API_KEY")
        print("\nExample:")
        print("export GROQ_API_KEY='your-key-here'")
        print("python demo.py")
        sys.exit(1)

    print(f"✅ Using {provider} ({model_name})")

    # Get user input
    print("\n📝 Enter your system prompt to test:")
    print("(Press Ctrl+D when done, or Ctrl+C to exit)")
    
    prompt_lines = []
    try:
        while True:
            line = input()
            prompt_lines.append(line)
    except EOFError:
        pass
    except KeyboardInterrupt:
        print("\n👋 Goodbye!")
        sys.exit(0)
    
    system_prompt = "\n".join(prompt_lines).strip()
    
    if not system_prompt:
        print("❌ Error: Empty prompt. Please try again.")
        sys.exit(1)
    
    print(f"\n🎯 Testing prompt: {system_prompt[:100]}...")
    
    # Import and run auditor
    try:
        from gateway.cli.security_auditor_cli import UniversalSecurityAuditor

        # Initialize auditor with selected model
        auditor = UniversalSecurityAuditor(model=model)
        
        # Create quick config
        config = auditor.create_config_from_prompt(
            name="Demo Test",
            description="Quick demo test",
            system_prompt=system_prompt,
            business_rules=["Follow all safety guidelines"],
            custom_attacks=None
        )
        
        # Run audit
        if not verbose:
            print("\n🚀 Running security audit...")
        report = auditor.audit_security(config, verbose=verbose)
        
        # Print results
        auditor.print_report(report)
        
        # Suggest next steps
        print("\n🎉 Demo complete!")
        print("\n📖 Next steps:")
        print("1. Try verbose mode: python demo.py --verbose")
        print("2. Try the interactive mode: python audit_prompt.py interactive")
        print("3. Test the examples: python audit_prompt.py audit-config --config configs/customer_support_bot.yaml")
        print("4. Read the full guide: SETUP.md")
        
    except ImportError as e:
        print(f"❌ Import error: {e}")
        print("Please install dependencies: pip install -r requirements.txt")
    except Exception as e:
        print(f"❌ Error: {e}")
        print("Please check your API key and try again.")


if __name__ == "__main__":
    load_env_file()  # Load .env before running
    main()