codes/test_gemma.py

"""
Test script to verify Gemma provider works with real API calls
"""
import os
from llm_provider import get_provider, get_default_model

# Set API key
os.environ['NVIDIA_API_KEY'] = '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'

# Initialize Gemma provider
print("Initializing Gemma provider...")
provider = get_provider('gemma')
model = get_default_model('gemma')
print(f"✅ Provider initialized with model: {model}")

# Test a simple completion
print("\nTesting completion...")
messages = [
    {"role": "user", "content": "Say 'Hello, I am Gemma!' in exactly those words."}
]

try:
    completion = provider.create_completion(messages, model, max_tokens=50)
    response_text = provider.get_response_text(completion)
    usage = provider.get_usage_info(completion)
    cost = provider.calculate_cost(usage, model)
    
    print(f"\n✅ Completion successful!")
    print(f"Response: {response_text}")
    print(f"\nUsage:")
    print(f"  - Prompt tokens: {usage['prompt_tokens']}")
    print(f"  - Completion tokens: {usage['completion_tokens']}")
    print(f"  - Total tokens: {usage['total_tokens']}")
    print(f"  - Cost: ${cost:.6f}")
    
except Exception as e:
    print(f"\n❌ Completion failed: {e}")
    import traceback
    traceback.print_exc()