Ollama OpenAI Compatibility API¶

import os
import time
import requests
from openai import OpenAI

# === Configuration ===
OLLAMA_HOST = os.getenv("OLLAMA_HOST", "http://ollama:11434")
DEFAULT_MODEL = "llama3.2:latest"

# Initialize OpenAI client pointing to Ollama
client = OpenAI(
    base_url=f"{OLLAMA_HOST}/v1",
    api_key="ollama"  # Required by library but ignored by Ollama
)

print(f"Ollama host: {OLLAMA_HOST}")
print(f"OpenAI base URL: {OLLAMA_HOST}/v1")
print(f"Default model: {DEFAULT_MODEL}")

def check_ollama_health() -> tuple[bool, bool]:
    """Check if Ollama server is running and model is available.
    
    Returns:
        tuple: (server_healthy, model_available)
    """
    try:
        response = requests.get(f"{OLLAMA_HOST}/api/tags", timeout=5)
        if response.status_code == 200:
            print("✓ Ollama server is running!")
            models = response.json()
            model_names = [m.get("name", "") for m in models.get("models", [])]
            
            if DEFAULT_MODEL in model_names:
                print(f"✓ Model '{DEFAULT_MODEL}' is available")
                return True, True
            else:
                print(f"✗ Model '{DEFAULT_MODEL}' not found!")
                print()
                if model_names:
                    print("Available models:")
                    for name in model_names:
                        print(f"  - {name}")
                else:
                    print("No models installed.")
                print()
                print("To fix this, run:")
                print(f"  ujust ollama pull {DEFAULT_MODEL.split(':')[0]}")
                return True, False
        else:
            print(f"Ollama returned unexpected status: {response.status_code}")
            return False, False
    except requests.exceptions.ConnectionError:
        print("✗ Cannot connect to Ollama server!")
        print("To fix this, run: ujust ollama start")
        return False, False
    except requests.exceptions.Timeout:
        print("✗ Connection to Ollama timed out!")
        return False, False

ollama_healthy, model_available = check_ollama_health()

print("=== List Available Models ===")

models = client.models.list()
for model in models.data:
    print(f"  - {model.id}")

print("=== Generate Response ===")

if not model_available:
    print()
    print("⚠ Skipping - model not available")
    print(f"  Run: ujust ollama pull {DEFAULT_MODEL.split(':')[0]}")
else:
    prompt = "Why is the sky blue? Answer in one sentence."
    print(f"Prompt: {prompt}")
    print()

    try:
        start_time = time.perf_counter()
        response = client.completions.create(
            model=DEFAULT_MODEL,
            prompt=prompt,
            max_tokens=100
        )
        end_time = time.perf_counter()

        print(f"Response: {response.choices[0].text}")
        print()
        print(f"Latency: {end_time - start_time:.2f}s")
        print(f"Completion tokens: {response.usage.completion_tokens}")
    except Exception as e:
        print(f"✗ Error: {e}")

print("=== Chat Completion ===")

if not model_available:
    print()
    print("⚠ Skipping - model not available")
    print(f"  Run: ujust ollama pull {DEFAULT_MODEL.split(':')[0]}")
else:
    try:
        response = client.chat.completions.create(
            model=DEFAULT_MODEL,
            messages=[
                {"role": "system", "content": "You are a helpful assistant. Keep responses brief."},
                {"role": "user", "content": "Explain machine learning in one sentence."}
            ],
            temperature=0.7,
            max_tokens=100
        )

        print(f"Assistant: {response.choices[0].message.content}")
        print(f"\nTokens used: {response.usage.total_tokens}")
    except Exception as e:
        print(f"✗ Error: {e}")

print("=== Multi-turn Conversation ===")

if not model_available:
    print()
    print("⚠ Skipping - model not available")
    print(f"  Run: ujust ollama pull {DEFAULT_MODEL.split(':')[0]}")
else:
    try:
        messages = [
            {"role": "system", "content": "You are a helpful math tutor."}
        ]

        # Turn 1
        messages.append({"role": "user", "content": "What is 2 + 2?"})
        response = client.chat.completions.create(
            model=DEFAULT_MODEL,
            messages=messages,
            max_tokens=50
        )
        assistant_msg = response.choices[0].message.content
        messages.append({"role": "assistant", "content": assistant_msg})
        print(f"User: What is 2 + 2?")
        print(f"Assistant: {assistant_msg}")

        # Turn 2
        messages.append({"role": "user", "content": "And what is that multiplied by 3?"})
        response = client.chat.completions.create(
            model=DEFAULT_MODEL,
            messages=messages,
            max_tokens=50
        )
        print(f"User: And what is that multiplied by 3?")
        print(f"Assistant: {response.choices[0].message.content}")
    except Exception as e:
        print(f"✗ Error: {e}")

print("=== Streaming Response ===")

if not model_available:
    print()
    print("⚠ Skipping - model not available")
    print(f"  Run: ujust ollama pull {DEFAULT_MODEL.split(':')[0]}")
else:
    try:
        print()

        stream = client.chat.completions.create(
            model=DEFAULT_MODEL,
            messages=[{"role": "user", "content": "Count from 1 to 5."}],
            stream=True
        )

        collected = []
        for chunk in stream:
            if chunk.choices[0].delta.content:
                collected.append(chunk.choices[0].delta.content)

        print(f"Response: {''.join(collected)}")
    except Exception as e:
        print(f"✗ Error: {e}")

print("=== Generate Embeddings ===")

if not model_available:
    print()
    print("⚠ Skipping - model not available")
    print(f"  Run: ujust ollama pull {DEFAULT_MODEL.split(':')[0]}")
else:
    try:
        test_text = "Ollama makes running LLMs locally easy and efficient."

        response = client.embeddings.create(
            model=DEFAULT_MODEL,
            input=test_text
        )

        embedding = response.data[0].embedding
        print(f"Input: '{test_text}'")
        print(f"Embedding dimensions: {len(embedding)}")
        print(f"First 5 values: {embedding[:5]}")
        print(f"Last 5 values: {embedding[-5:]}")
    except Exception as e:
        print(f"✗ Error: {e}")

print("=== Error Handling ===")

# Test: Non-existent model
print("\n1. Testing non-existent model...")
try:
    response = client.chat.completions.create(
        model="invalid-model",
        messages=[{"role": "user", "content": "Hello"}]
    )
    print(f"   Unexpected success")
except Exception as e:
    print(f"   Expected error: {type(e).__name__}")

# Test: Empty messages
print("\n2. Testing empty messages...")
try:
    response = client.chat.completions.create(
        model=DEFAULT_MODEL,
        messages=[]
    )
    print(f"   Empty messages allowed")
except Exception as e:
    print(f"   Error: {type(e).__name__}")

print("\nError handling tests completed!")