Ollama Python Library¶

import os
import time
import ollama

# === 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:
        models = ollama.list()
        print("✓ Ollama server is running!")
        model_names = [m.get("model", "") 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
    except Exception as e:
        print(f"✗ Cannot connect to Ollama server!")
        print(f"Error: {e}")
        print("To fix this, run: ujust ollama start")
        return False, False

ollama_healthy, model_available = check_ollama_health()

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

models = ollama.list()

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

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:
    try:
        model_info = ollama.show(DEFAULT_MODEL)

        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"\nModel file preview:")
        modelfile = model_info.get("modelfile", "N/A")
        print(f"  {modelfile[:300]}..." if len(modelfile) > 300 else f"  {modelfile}")
    except Exception as e:
        print(f"✗ Error: {e}")

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

running = ollama.ps()

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)
        print(f"  - {name}")
        print(f"    Size: {size:.2f} GB | VRAM: {vram:.2f} GB")
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:
    try:
        prompt = "Why is the sky blue? Answer in one sentence."
        print(f"Prompt: {prompt}")
        print()

        start_time = time.perf_counter()
        result = ollama.generate(
            model=DEFAULT_MODEL,
            prompt=prompt
        )
        end_time = time.perf_counter()

        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')}")
    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 = ollama.chat(
            model=DEFAULT_MODEL,
            messages=[
                {"role": "system", "content": "You are a helpful assistant. Keep responses brief."},
                {"role": "user", "content": "What is Python?"}
            ]
        )

        print(f"Assistant: {response['message']['content']}")
    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:
        # Turn 1
        messages = [
            {"role": "user", "content": "What is 2 + 2?"},
        ]

        response = ollama.chat(
            model=DEFAULT_MODEL,
            messages=messages
        )
        print(f"User: What is 2 + 2?")
        print(f"Assistant: {response['message']['content']}")

        # Continue conversation
        messages.append(response["message"])
        messages.append({"role": "user", "content": "And what is that multiplied by 3?"})

        response = ollama.chat(
            model=DEFAULT_MODEL,
            messages=messages
        )
        print(f"User: And what is that multiplied by 3?")
        print(f"Assistant: {response['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 = ollama.generate(
            model=DEFAULT_MODEL,
            prompt="Count from 1 to 5.",
            stream=True
        )

        collected = []
        for chunk in stream:
            collected.append(chunk["response"])

        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."

        result = ollama.embed(
            model=DEFAULT_MODEL,
            input=test_text
        )

        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:]}")
    except Exception as e:
        print(f"✗ Error: {e}")

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}'...")
    try:
        ollama.copy(source=DEFAULT_MODEL, destination=COPY_NAME)
        print(f"   Copy successful!")
    except Exception as e:
        print(f"   Copy failed: {e}")

    # Step 2: Verify the copy exists
    print(f"\n2. Verifying '{COPY_NAME}' exists...")
    models = ollama.list()
    model_names = [m["model"] for m in models.get("models", [])]
    if COPY_NAME in model_names:
        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}'...")
    try:
        ollama.delete(COPY_NAME)
        print(f"   Delete successful!")
    except Exception as e:
        print(f"   Delete failed: {e}")

    # Step 4: Verify deletion
    print(f"\n4. Verifying '{COPY_NAME}' is deleted...")
    models = ollama.list()
    model_names = [m["model"] for m in models.get("models", [])]
    if COPY_NAME not in model_names:
        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...")
try:
    result = ollama.generate(
        model="nonexistent-model-xyz",
        prompt="Hello"
    )
    print(f"   Unexpected success: {result}")
except Exception as e:
    print(f"   Expected error: {type(e).__name__}: {e}")

# Test: Empty prompt
print("\n2. Testing empty prompt...")
try:
    result = ollama.generate(
        model=DEFAULT_MODEL,
        prompt=""
    )
    print(f"   Empty prompts allowed")
except Exception as e:
    print(f"   Error: {type(e).__name__}: {e}")

print("\nError handling tests completed!")