#!/usr/bin/env python3
"""
Test script for Ollama OpenELM API
"""

from openai import OpenClient

def test_ollama():
    """Test Ollama with OpenAI SDK."""
    print("Testing Ollama OpenELM with OpenAI SDK...")
    print("-" * 50)
    
    client = OpenAI(
        base_url="http://127.0.0.1:11434/v1",
        api_key="ollama",
    )
    
    # Test 1: Basic generation
    print("\n[Test 1] Basic generation:")
    response = client.chat.completions.create(
        model="apple/OpenELM-3B-Instruct",
        messages=[{"role": "user", "content": "Say hello!"}],
        max_tokens=100,
        temperature=0.7
    )
    print(f"Response: {response.choices[0].message.content}")
    print(f"Tokens: {response.usage.total_tokens}")
    
    # Test 2: Multi-turn conversation
    print("\n[Test 2] Multi-turn conversation:")
    response = client.chat.completions.create(
        model="apple/OpenELM-3B-Instruct",
        messages=[
            {"role": "user", "content": "What is AI?"},
            {"role": "assistant", "content": "AI is Artificial Intelligence."},
            {"role": "user", "content": "What are examples?"}
        ],
        max_tokens=150,
        temperature=0.5
    )
    print(f"Response: {response.choices[0].message.content}")
    
    # Test 3: Creative writing
    print("\n[Test 3] Creative writing:")
    response = client.chat.completions.create(
        model="apple/OpenELM-3B-Instruct",
        messages=[{"role": "user", "content": "Write a short poem about technology."}],
        max_tokens=200,
        temperature=0.8
    )
    print(f"Response: {response.choices[0].message.content}")
    
    # Test 4: Question answering
    print("\n[Test 4] Question answering:")
    response = client.chat.completions.create(
        model="apple/OpenELM-3B-Instruct",
        messages=[{"role": "user", "content": "Explain what is machine learning in simple terms."}],
        max_tokens=250,
        temperature=0.6
    )
    print(f"Response: {response.choices[0].message.content}")
    
    print("\n" + "=" * 50)
    print("All tests completed successfully!")
    print("=" * 50)


if __name__ == "__main__":
    test_ollama()