test_copilot_api_keys.py

#!/usr/bin/env python3
"""
Copilot API Key Validation Test
Tests if all API keys work correctly
"""

import asyncio
import sys
from pathlib import Path

# Add project root to path
sys.path.insert(0, str(Path(__file__).parent))

from src.copilot.config import config


async def test_groq():
    """Test Groq API connection"""
    print("\n" + "=" * 70)
    print("TESTING GROQ API CONNECTION")
    print("=" * 70 + "\n")

    groq_keys = config.get_groq_api_keys()

    if not groq_keys:
        print("[FAIL] No Groq API keys configured")
        return False

    print(f"Testing {len(groq_keys)} Groq API key(s)...\n")

    try:
        from groq import Groq

        # Only test first 3 keys to save time
        for i, key in enumerate(groq_keys[:3], 1):
            print(f"Testing Key {i}...", end=" ")

            client = Groq(api_key=key)

            response = client.chat.completions.create(
                model="llama-3.1-8b-instant",
                messages=[
                    {
                        "role": "user",
                        "content": "Say 'Groq API is working' in exactly those words.",
                    }
                ],
                max_tokens=50,
            )

            if response.choices and response.choices[0].message.content:
                text = response.choices[0].message.content
                print(f"[OK] Working (response: {text[:40]}...)")
            else:
                print("[FAIL] No response received")
                return False

        print("\n[OK] All Groq API keys are working!")
        return True

    except Exception as e:
        print(f"[FAIL] Error: {e}")
        return False


async def test_gemini():
    """Test Gemini API connection"""
    print("\n" + "=" * 70)
    print("TESTING GEMINI API CONNECTION")
    print("=" * 70 + "\n")

    gemini_key = config.get_gemini_api_key()

    if not gemini_key:
        print("[FAIL] No Gemini API key configured")
        return False

    print("Testing Gemini API...", end=" ")

    try:
        import google.generativeai as genai

        genai.configure(api_key=gemini_key)

        # Try multiple model names (current Gemini 2.x models)
        model_names = [
            "gemini-2.5-flash",
            "gemini-2.0-flash",
            "gemini-flash-latest",
            "gemini-2.5-flash-lite",
            "gemini-2.0-flash-lite",
        ]

        response = None
        last_error = None
        for model_name in model_names:
            try:
                model = genai.GenerativeModel(model_name)
                response = model.generate_content(
                    "Say 'Gemini API is working' in exactly those words."
                )
                print(f"\n   Using model: {model_name}", end="")
                break
            except Exception as e:
                last_error = e
                continue

        if response is None:
            print(f"\n[FAIL] No working Gemini model found. Last error: {last_error}")
            return False

        if response.text:
            print(f"[OK] Working (response: {response.text[:40]}...)")
        else:
            print("[FAIL] No response received")
            return False

        print("\n[OK] Gemini API is working!")
        return True

    except Exception as e:
        print(f"[FAIL] Error: {e}")
        return False


async def test_pinecone():
    """Test Pinecone API connection (optional)"""
    print("\n" + "=" * 70)
    print("TESTING PINECONE API CONNECTION (Optional)")
    print("=" * 70 + "\n")

    pinecone_config = config.get_pinecone_config()

    if not pinecone_config["available"]:
        print("[WARN]  Pinecone not configured")
        print("   └─ This is optional, using SQLite embeddings instead")
        return True

    print("Testing Pinecone API...", end=" ")

    try:
        from pinecone import Pinecone

        pc = Pinecone(api_key=pinecone_config["api_key"])

        # Just test if we can initialize (don't need to create/list indexes)
        # as that requires the environment name
        print("[OK] Pinecone API is configured")
        print("\n[OK] Pinecone is available!")
        return True

    except Exception as e:
        print(f"[WARN]  Pinecone error (optional): {e}")
        print("   └─ Skipping Pinecone, will use SQLite embeddings")
        return True


async def test_anthropic():
    """Test Anthropic API connection (optional)"""
    print("\n" + "=" * 70)
    print("TESTING ANTHROPIC API CONNECTION (Optional)")
    print("=" * 70 + "\n")

    anthropic_key = config.get_anthropic_api_key()

    if not anthropic_key:
        print("[WARN]  Anthropic API key not configured")
        print("   └─ This is optional, using Groq as primary")
        return True

    print("Testing Anthropic API...", end=" ")

    try:
        from anthropic import Anthropic

        client = Anthropic(api_key=anthropic_key)

        response = client.messages.create(
            model="claude-3-5-haiku-20241022",
            max_tokens=50,
            messages=[
                {
                    "role": "user",
                    "content": "Say 'Anthropic API is working' in exactly those words.",
                }
            ],
        )

        if response.content:
            print(f"[OK] Working (response: {response.content[0].text[:40]}...)")
        else:
            print("[FAIL] No response received")
            return False

        print("\n[OK] Anthropic API is working!")
        return True

    except Exception as e:
        print(f"[WARN]  Anthropic error (optional): {e}")
        print("   └─ Skipping Anthropic, will use Groq as primary")
        return True


async def main():
    """Run all tests"""
    print("\n" + "=" * 70)
    print("[ROCKET] COPILOT API KEY VALIDATION TEST")
    print("=" * 70)

    config.print_status()

    # Run tests
    results = {}

    print("\n\nStarting API tests...\n")

    results["groq"] = await test_groq()
    results["gemini"] = await test_gemini()
    results["pinecone"] = await test_pinecone()
    results["anthropic"] = await test_anthropic()

    # Summary
    print("\n" + "=" * 70)
    print("TEST SUMMARY")
    print("=" * 70 + "\n")

    for provider, success in results.items():
        symbol = "[OK]" if success else "[FAIL]"
        print(f"{symbol} {provider.upper():15} {('PASS' if success else 'FAIL')}")

    print("\n" + "=" * 70)

    # Final verdict
    required_passed = results["groq"] and results["gemini"]

    if required_passed:
        print("[OK] ALL REQUIRED APIs ARE WORKING!")
        print("\n[ROCKET] Your copilot is ready to use!")
        print("\nNext steps:")
        print("   1. Start server: python server.py")
        print("   2. Open http://localhost:8000")
        print("   3. Click copilot button (bottom-right)")
        print("   4. Start asking questions!\n")
        return 0
    else:
        print("[FAIL] SOME REQUIRED APIs ARE NOT WORKING")
        print("\nPlease fix the issues above and try again.")
        print("Common fixes:")
        print("   - Check if API key is copied correctly (no extra spaces)")
        print("   - Verify key is valid and active in provider dashboard")
        print("   - Run 'python setup_copilot.py' to reconfigure\n")
        return 1

    print("=" * 70 + "\n")


if __name__ == "__main__":
    try:
        exit_code = asyncio.run(main())
        sys.exit(exit_code)
    except KeyboardInterrupt:
        print("\n\n[FAIL] Test cancelled by user")
        sys.exit(1)
    except Exception as e:
        print(f"\n[FAIL] Error during testing: {e}")
        import traceback

        traceback.print_exc()
        sys.exit(1)