test_new_api.py

#!/usr/bin/env python3
"""
Quick test script for the new Google Gemini SDK implementation
"""

import os
import sys

# Test if the SDK can be imported
try:
    from google import genai
    from google.genai import types
    print("✓ google-genai SDK imported successfully")
except ImportError as e:
    print(f"✗ Failed to import google-genai SDK: {e}")
    sys.exit(1)

# Test if API key is set
api_key = os.getenv("GOOGLE_API_KEY")
if not api_key:
    print("✗ GOOGLE_API_KEY environment variable not set")
    print("  Please set it: export GOOGLE_API_KEY='your-api-key'")
    sys.exit(1)
else:
    print(f"✓ GOOGLE_API_KEY found (length: {len(api_key)})")

# Test client creation
try:
    client = genai.Client(api_key=api_key)
    print("✓ Gemini client created successfully")
except Exception as e:
    print(f"✗ Failed to create client: {e}")
    sys.exit(1)

# Test a simple API call
print("\nTesting API call with gemini-2.0-flash...")
try:
    response = client.models.generate_content(
        model='gemini-2.0-flash',
        contents='What is 2+2? Answer in one word.',
        config=types.GenerateContentConfig(
            temperature=0.1,
            max_output_tokens=10
        )
    )
    print(f"✓ API call successful!")
    print(f"  Response: {response.text.strip()}")
except Exception as e:
    print(f"✗ API call failed: {e}")
    sys.exit(1)

print("\n✓ All tests passed! The new SDK is working correctly.")