Ollama REST API (requests)¶

import os
import requests
import json
import time

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

print(f"Ollama host: {OLLAMA_HOST}")
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 ===")

response = requests.get(f"{OLLAMA_HOST}/api/tags")
models_data = response.json()

if models_data.get("models"):
    for model in models_data["models"]:
        size_gb = model.get("size", 0) / (1024**3)
        print(f"  - {model['name']} ({size_gb:.2f} GB)")
else:
    print("  No models found. Run: ujust ollama pull llama3.2")

print("=== Show Model Details ===")

if not model_available:
    print()
    print("⚠ Skipping - model not available")
    print(f"  Run: ujust ollama pull {DEFAULT_MODEL.split(':')[0]}")
else:
    response = requests.post(
        f"{OLLAMA_HOST}/api/show",
        json={"model": DEFAULT_MODEL}
    )
    model_info = response.json()

    if "error" in model_info:
        print(f"✗ Error: {model_info['error']}")
    else:
        print(f"Model: {DEFAULT_MODEL}")
        print(f"\nDetails:")
        if "details" in model_info:
            details = model_info["details"]
            print(f"  Family: {details.get('family', 'N/A')}")
            print(f"  Parameter Size: {details.get('parameter_size', 'N/A')}")
            print(f"  Quantization: {details.get('quantization_level', 'N/A')}")

        print(f"\nTemplate preview:")
        template = model_info.get("template", "N/A")
        print(f"  {template[:200]}..." if len(template) > 200 else f"  {template}")

print("=== List Running Models ===")

response = requests.get(f"{OLLAMA_HOST}/api/ps")
running = response.json()

if running.get("models"):
    for model in running["models"]:
        name = model.get("name", "Unknown")
        size = model.get("size", 0) / (1024**3)
        vram = model.get("size_vram", 0) / (1024**3)
        expires = model.get("expires_at", "N/A")
        print(f"  - {name}")
        print(f"    Size: {size:.2f} GB | VRAM: {vram:.2f} GB")
        print(f"    Expires: {expires}")
else:
    print("  No models currently loaded in memory")

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()

    start_time = time.perf_counter()
    response = requests.post(
        f"{OLLAMA_HOST}/api/generate",
        json={
            "model": DEFAULT_MODEL,
            "prompt": prompt,
            "stream": False
        }
    )
    end_time = time.perf_counter()

    result = response.json()
    if "error" in result:
        print(f"✗ Error: {result['error']}")
    else:
        print(f"Response: {result['response']}")
        print()
        print(f"Latency: {end_time - start_time:.2f}s")
        print(f"Eval tokens: {result.get('eval_count', 'N/A')}")
        print(f"Eval duration: {result.get('eval_duration', 0) / 1e9:.2f}s")

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

if not model_available:
    print()
    print("⚠ Skipping - model not available")
    print(f"  Run: ujust ollama pull {DEFAULT_MODEL.split(':')[0]}")
else:
    messages = [
        {"role": "system", "content": "You are a helpful assistant. Keep responses brief."},
        {"role": "user", "content": "What is Python?"}
    ]

    response = requests.post(
        f"{OLLAMA_HOST}/api/chat",
        json={
            "model": DEFAULT_MODEL,
            "messages": messages,
            "stream": False
        }
    )

    result = response.json()
    if "error" in result:
        print(f"✗ Error: {result['error']}")
    else:
        print(f"Assistant: {result['message']['content']}")

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

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

    response = requests.post(
        f"{OLLAMA_HOST}/api/generate",
        json={
            "model": DEFAULT_MODEL,
            "prompt": "Count from 1 to 5.",
            "stream": True
        },
        stream=True
    )

    collected = []
    for line in response.iter_lines():
        if line:
            chunk = json.loads(line)
            if "error" in chunk:
                print(f"✗ Error: {chunk['error']}")
                break
            token = chunk.get("response", "")
            collected.append(token)
            if chunk.get("done"):
                break

    if collected:
        print(f"Response: {''.join(collected)}")

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

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

    response = requests.post(
        f"{OLLAMA_HOST}/api/embed",
        json={
            "model": DEFAULT_MODEL,
            "input": test_text
        }
    )
    result = response.json()

    if "error" in result:
        print(f"✗ Error: {result['error']}")
    else:
        embeddings = result.get("embeddings", [[]])[0]
        print(f"Input: '{test_text}'")
        print(f"Embedding dimensions: {len(embeddings)}")
        print(f"First 5 values: {embeddings[:5]}")
        print(f"Last 5 values: {embeddings[-5:]}")

print("=== Copy and Delete Model ===")

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

    # Step 1: Copy the model
    print(f"\n1. Copying '{DEFAULT_MODEL}' to '{COPY_NAME}'...")
    response = requests.post(
        f"{OLLAMA_HOST}/api/copy",
        json={
            "source": DEFAULT_MODEL,
            "destination": COPY_NAME
        }
    )
    if response.status_code == 200:
        print(f"   Copy successful!")
    else:
        print(f"   Copy failed: {response.text}")

    # Step 2: Verify the copy exists
    print(f"\n2. Verifying '{COPY_NAME}' exists...")
    response = requests.get(f"{OLLAMA_HOST}/api/tags")
    models = [m["name"] for m in response.json().get("models", [])]
    if COPY_NAME in models:
        print(f"   Found '{COPY_NAME}' in model list")
    else:
        print(f"   '{COPY_NAME}' not found")

    # Step 3: Delete the copy
    print(f"\n3. Deleting '{COPY_NAME}'...")
    response = requests.delete(
        f"{OLLAMA_HOST}/api/delete",
        json={"model": COPY_NAME}
    )
    if response.status_code == 200:
        print(f"   Delete successful!")
    else:
        print(f"   Delete failed: {response.text}")

    # Step 4: Verify deletion
    print(f"\n4. Verifying '{COPY_NAME}' is deleted...")
    response = requests.get(f"{OLLAMA_HOST}/api/tags")
    models = [m["name"] for m in response.json().get("models", [])]
    if COPY_NAME not in models:
        print(f"   '{COPY_NAME}' successfully removed")
    else:
        print(f"   '{COPY_NAME}' still exists")

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

# Test: Non-existent model
print("\n1. Testing non-existent model...")
response = requests.post(
    f"{OLLAMA_HOST}/api/generate",
    json={
        "model": "nonexistent-model-xyz",
        "prompt": "Hello",
        "stream": False
    }
)
if response.status_code != 200:
    print(f"   Expected error: {response.status_code} - {response.text[:100]}")
else:
    print(f"   Unexpected success")

# Test: Empty prompt
print("\n2. Testing empty prompt...")
response = requests.post(
    f"{OLLAMA_HOST}/api/generate",
    json={
        "model": DEFAULT_MODEL,
        "prompt": "",
        "stream": False
    }
)
if response.status_code == 200:
    print(f"   Empty prompts allowed (returned response)")
else:
    print(f"   Error: {response.status_code}")

print("\nError handling tests completed!")