setup_copilot.py

#!/usr/bin/env python3
"""
COPILOT SETUP SCRIPT
Sets up API keys and validates copilot configuration

Usage:
    python setup_copilot.py
"""

import os
import sys
from pathlib import Path


def print_header(text: str):
    """Print formatted header"""
    print("\n" + "=" * 70)
    print(f"  {text}")
    print("=" * 70 + "\n")


def get_user_input(prompt: str, is_required: bool = False) -> str:
    """Get input from user with validation"""
    while True:
        value = input(prompt).strip()
        if not value and is_required:
            print("[FAIL] This field is required. Please enter a value.\n")
            continue
        return value


def setup_groq_keys():
    """Guide user through Groq API key setup"""
    print_header("GROQ API KEYS SETUP")
    print("Groq provides free access to Llama 3.1 models")
    print("Get free API keys from: https://console.groq.com/keys\n")

    groq_keys = []
    for i in range(1, 6):
        print(f"\n[NOTE] Groq API Key #{i}")
        print("   (Leave blank to skip additional keys, need at least 1)")

        key = get_user_input("Enter GROQ_API_KEY_" + str(i) + ": ")

        if not key:
            if i == 1:
                print("[FAIL] At least one Groq API key is required!")
                continue
            else:
                break

        if len(key) < 10:
            print("[FAIL] Key seems too short, are you sure? (y/n): ", end="")
            if input().lower() != "y":
                continue

        groq_keys.append((f"GROQ_API_KEY_{i}", key))
        print(f"[OK] Key {i} saved (masked: {key[:10]}...{key[-5:]})")

    return groq_keys


def setup_gemini_key():
    """Guide user through Gemini API key setup"""
    print_header("GEMINI API KEY SETUP")
    print("Gemini 1.5 Flash provides free access (1M tokens/month)")
    print("Get free API key from: https://ai.google.dev/\n")

    while True:
        key = get_user_input("Enter GEMINI_API_KEY: ", is_required=True)

        if len(key) < 10:
            print("[FAIL] Key seems too short, are you sure? (y/n): ", end="")
            if input().lower() != "y":
                continue

        print(f"[OK] Key saved (masked: {key[:10]}...{key[-5:]})")
        return [("GEMINI_API_KEY", key)]


def setup_pinecone_key():
    """Guide user through Pinecone setup (optional)"""
    print_header("PINECONE VECTOR DB SETUP (Optional)")
    print("Pinecone is optional - can use SQLite embeddings instead")
    print("Get free API key from: https://www.pinecone.io/\n")

    use_pinecone = input("Do you want to use Pinecone? (y/n) [default: n]: ").lower()

    if use_pinecone != "y":
        print("[WARN]  Skipping Pinecone - will use SQLite embeddings (local)")
        return []

    api_key = get_user_input("Enter PINECONE_API_KEY: ", is_required=True)
    index_name = get_user_input("Enter PINECONE_INDEX_NAME [default: aml-alerts]: ")
    if not index_name:
        index_name = "aml-alerts"

    environment = get_user_input("Enter PINECONE_ENVIRONMENT (optional): ")

    pinecone_config = [
        ("PINECONE_API_KEY", api_key),
        ("PINECONE_INDEX_NAME", index_name),
    ]

    if environment:
        pinecone_config.append(("PINECONE_ENVIRONMENT", environment))

    print(f"[OK] Pinecone configured (Index: {index_name})")
    return pinecone_config


def setup_anthropic_key():
    """Guide user through Anthropic setup (optional)"""
    print_header("ANTHROPIC API KEY SETUP (Optional)")
    print("Anthropic provides Claude models (optional fallback)")
    print("Get API key from: https://console.anthropic.com/\n")

    use_anthropic = input("Do you want to add Anthropic API key? (y/n) [default: n]: ").lower()

    if use_anthropic != "y":
        print("[WARN]  Skipping Anthropic - Groq will be used as primary")
        return []

    api_key = get_user_input("Enter ANTHROPIC_API_KEY: ", is_required=True)
    print(f"[OK] Anthropic key saved (masked: {api_key[:10]}...{api_key[-5:]})")
    return [("ANTHROPIC_API_KEY", api_key)]


def write_env_file(api_keys_dict: dict):
    """Write API keys to .env.copilot"""
    env_file = Path(__file__).parent / ".env.copilot"

    # Read existing template
    template_file = Path(__file__).parent / ".env.copilot"
    if template_file.exists():
        with open(template_file, 'r') as f:
            content = f.read()
    else:
        content = "# Copilot Configuration\n"

    # Update with user values
    for key, value in api_keys_dict.items():
        # Replace or add key
        lines = content.split('\n')
        found = False

        for i, line in enumerate(lines):
            if line.startswith(f"{key}="):
                lines[i] = f"{key}={value}"
                found = True
                break

        if not found:
            # Add at end (before EOF)
            lines.insert(-1, f"{key}={value}")

        content = '\n'.join(lines)

    # Write to file
    with open(env_file, 'w') as f:
        f.write(content)

    print(f"[OK] Configuration saved to: {env_file}")


def main():
    """Main setup flow"""
    os.system('clear' if os.name == 'posix' else 'cls')

    print("\n" + "=" * 70)
    print("  [ROCKET] UNION BANK AML COPILOT - API SETUP")
    print("=" * 70)

    print("\nThis wizard will help you configure API keys for the AML Copilot")
    print("All keys will be saved to .env.copilot (do not commit to git)\n")

    # Collect all keys
    all_keys = {}

    # Required: Groq
    groq_keys = setup_groq_keys()
    all_keys.update(dict(groq_keys))

    if not groq_keys:
        print("\n[FAIL] At least one Groq API key is required!")
        sys.exit(1)

    # Required: Gemini
    gemini_keys = setup_gemini_key()
    all_keys.update(dict(gemini_keys))

    # Optional: Pinecone
    pinecone_keys = setup_pinecone_key()
    all_keys.update(dict(pinecone_keys))

    # Optional: Anthropic
    anthropic_keys = setup_anthropic_key()
    all_keys.update(dict(anthropic_keys))

    # Save to file
    print_header("SAVING CONFIGURATION")
    write_env_file(all_keys)

    # Validate
    print_header("VALIDATING CONFIGURATION")
    print("Loading and validating all API keys...")
    print("(This will show the status of each provider)\n")

    # Need to reload config after writing
    from importlib import reload
    import src.copilot.config as config_module
    reload(config_module)
    from src.copilot.config import config as updated_config

    updated_config.print_status()

    if updated_config.is_ready():
        print("\n[OK] SETUP COMPLETE!")
        print("\nYour copilot is now ready to use. Next steps:")
        print("   1. Start the server: python server.py")
        print("   2. Open http://localhost:8000 in your browser")
        print("   3. Click the copilot button (bottom right)")
        print("   4. Ask any question about accounts or alerts!")
    else:
        print("\n[WARN]  Setup incomplete - some required keys are missing")
        print("Please fill in missing keys and run this script again")
        sys.exit(1)


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\n\n[FAIL] Setup cancelled by user")
        sys.exit(0)
    except Exception as e:
        print(f"\n[FAIL] Error during setup: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)