scripts/setup_embeddings.py

"""
Quick script to help set up Google API key for fast embeddings
"""

from pathlib import Path


def setup_google_api_key():
    """Interactive setup for Google API key"""

    print("=" * 70)
    print("Fast Embeddings Setup - Google Gemini API")
    print("=" * 70)

    print("\nWhy Google Gemini?")
    print("   - 100x faster than local Ollama (2 mins vs 30+ mins)")
    print("   - FREE for standard usage")
    print("   - High quality embeddings")
    print("   - Automatic fallback to Ollama if unavailable")

    print("\n" + "=" * 70)
    print("Step 1: Get Your Free API Key")
    print("=" * 70)
    print("\n1. Open this URL in your browser:")
    print("   https://aistudio.google.com/app/apikey")
    print("\n2. Sign in with Google account")
    print("3. Click 'Create API Key'")
    print("4. Copy the key (starts with 'AIza...')")

    input("\nPress ENTER when you have your API key ready...")

    api_key = input("\nPaste your Google API key here: ").strip()

    if not api_key:
        print("\nNo API key provided. Using local Ollama instead.")
        return False

    if not api_key.startswith("AIza"):
        print("\nWarning: Key doesn't start with 'AIza'. Are you sure this is correct?")
        confirm = input("Continue anyway? (y/n): ").strip().lower()
        if confirm != "y":
            return False

    # Update .env file
    env_path = Path(".env")

    if env_path.exists():
        with open(env_path) as f:
            lines = f.readlines()

        # Update or add GOOGLE_API_KEY
        updated = False
        for i, line in enumerate(lines):
            if line.startswith("GOOGLE_API_KEY="):
                lines[i] = f"GOOGLE_API_KEY={api_key}\n"
                updated = True
                break

        if not updated:
            lines.insert(0, f"GOOGLE_API_KEY={api_key}\n")

        with open(env_path, "w") as f:
            f.writelines(lines)
    else:
        # Create new .env file
        with open(env_path, "w") as f:
            f.write(f"GOOGLE_API_KEY={api_key}\n")

    print("\nAPI key saved to .env file!")
    print("\n" + "=" * 70)
    print("Step 2: Build Vector Store")
    print("=" * 70)
    print("\nRun this command:")
    print("   python src/pdf_processor.py")
    print("\nChoose option 1 (Google Gemini) when prompted.")
    print("\n" + "=" * 70)

    return True


if __name__ == "__main__":
    try:
        setup_google_api_key()
    except KeyboardInterrupt:
        print("\n\nSetup cancelled.")
    except Exception as e:
        print(f"\nError: {e}")