#!/usr/bin/env python3
"""Test Gemini 2.5 Pro with google-genai"""

import os
from google import genai
from google.genai import types

# Set API key
os.environ["GEMINI_API_KEY"] = "AIzaSyCqe3vjvPlo1lt_hpQ4nqAC0-_1omva1oc"

def test_generate():
    print("Testing Gemini 2.5 Pro...")
    
    client = genai.Client(
        api_key=os.environ.get("GEMINI_API_KEY"),
    )

    model = "gemini-2.5-pro"
    contents = [
        types.Content(
            role="user",
            parts=[
                types.Part.from_text(text="Say hello in exactly 5 words"),
            ],
        ),
    ]
    
    generate_content_config = types.GenerateContentConfig(
        temperature=0.7,
        max_output_tokens=100,
    )

    print(f"Calling {model}...")
    response_text = ""
    
    try:
        for chunk in client.models.generate_content_stream(
            model=model,
            contents=contents,
            config=generate_content_config,
        ):
            if chunk.text:
                response_text += chunk.text
                print(f"Chunk: {chunk.text}")
        
        print(f"\nFull response: {response_text}")
    except Exception as e:
        print(f"Error: {e}")

if __name__ == "__main__":
    test_generate()