A lightweight, single‑file benchmark tool for Ollama‑powered LLMs. Runs short prompts in a streaming fashion to measure key metrics: TTFT (Time to First Token), TPS (Tokens per Second), and total tokens generated. The script: * accepts a model name and a JSON‑output flag via argparse; * defaults to Gemma‑4 (latest) and a short explanatory prompt…

llmbench.py

#!/usr/bin/env python3
# Copyright (C) 2026 Cyril Bosselut
#
# llmbench.py is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# llmbench.py is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with llmbench.py  If not, see <http://www.gnu.org/licenses/>.
import argparse
import requests
import time
import json
import re
from typing import Dict, Any
from html import escape

# --- CONFIGURATION ---
OLLAMA_API_URL = "http://localhost:11434/api"
MODEL_NAME = "llama3.2:latest"  # <-- CHANGE THIS to the model you want to test
PROMPT = "Explain the concept of quantization in LLMs in simple terms, using analogies." 
MAX_TOKENS = 512  # Limit the test to a fixed number of tokens for consistent comparison
TOKEN_LIMIT = 128
NUM_RUNS = 3
SLEEP_BETWEEN_RUNS = 2.0
TIMEOUT = 30.0

def benchmark_ollama(model: str, prompt: str, max_tokens: int, timeout: float, use_json: bool = False) -> Dict[str, Any]:
    """
    Benchmarks an LLM model by streaming the response from the Ollama API 
    to accurately measure TTFT and TPS.
    """
    if not use_json:
        print(f"--- Starting Benchmark for Model: {model} ---")
        print(f"Prompt: '{prompt[:50]}...'")
    
    # 1. Prepare the API payload
    payload = {
        "model": model,
        "prompt": prompt,
        "stream": True,  # CRITICAL: Must be True for streaming metrics
        "options": {
            "temperature": 0.1, # Keep temperature low for consistent testing
            "num_predict": max_tokens
        }
    }
    
    headers = {"Content-Type": "application/json"}
    
    # 2. Initialize timing variables
    start_time = time.time()
    first_token_time = None
    total_tokens = 0
    eval_duration = 0
    
    # 3. Execute the streaming request
    try:
        with requests.post(f"{OLLAMA_API_URL}/generate", headers=headers, json=payload, stream=True, timeout=timeout) as response:
            response.raise_for_status() # Raise HTTPError for bad responses
            
            # Process the stream chunk by chunk
            for line in response.iter_lines():
                if line:
                    # Ollama streams JSON objects separated by newlines
                    try:
                        data = json.loads(line.decode('utf-8'))
                    except json.JSONDecodeError:
                        continue # Skip malformed lines
                    
                    # Check for the token data
                    if 'response' in data and data['response']:
                        token = data['response']
                        total_tokens += 1
                        
                        # Calculate TTFT on the first token
                        if first_token_time is None:
                            first_token_time = time.time()
                        
                        if not use_json:
                            # Print the token to simulate real-time output
                            print(token, end="", flush=True)
                        
                    # Check for the end of the stream
                    if data.get("done"):
                        if 'eval_count' in data:
                            total_tokens = data['eval_count']
                        if 'eval_duration' in data:
                            eval_duration = data['eval_duration']/10**9
                        break
            if not use_json:
                print("\n--- Stream Complete ---")

    except requests.exceptions.Timeout:
        if not use_json:
            print("\n\033[0;31m[ERROR]\033[0m Connection Timed Out.")
        return {"error": "Timeout", "ttft": 0.0, "tps": 0.0, "model": model, "total_tokens": 0, "total_duration": 0.0}
    except requests.exceptions.ConnectionError:
        if not use_json:
            print("\n\033[0;31m[ERROR]\033[0m Could not connect to Ollama. Ensure the service is running at http://localhost:11434.")
        return {"error": "Connection Failed", "ttft": 0.0, "tps": 0.0, "model": model, "total_tokens": 0, "total_duration": 0.0}
    except requests.exceptions.HTTPError as e:
        if not use_json:
            print(f"\n\033[0;31m[ERROR]\033[0m HTTP Error occurred: {e}. Check if the model name is correct.")
        return {"error": f"HTTP Error: {e}", "ttft": 0.0, "tps": 0.0, "model": model, "total_tokens": 0, "total_duration": 0.0}
    finally:
        if response:
            response.close()
    
    # 4. Calculate Final Metrics
    end_time = time.time()
    if eval_duration > 0:
        total_duration = eval_duration
    else:
        total_duration = end_time - start_time
    
    # Calculate TTFT (Time elapsed until the first token was received)
    if first_token_time is not None:
        ttft = first_token_time - start_time
    else:
        ttft = 0.0
        
    # Calculate TPS (Total tokens / Total duration)
    tps = total_tokens / total_duration if total_duration > 0 else 0.0
    
    return {
        "model": model,
        "total_tokens": total_tokens,
        "total_duration": total_duration,
        "ttft": ttft,
        "tps": tps
    }

def parse_int(s) -> int|None:
    try:
        return int(s.strip())
    except:
        return None

def sanitize_input(input_str: str) -> str:
    input_str = input_str.strip()
    input_str = escape(input_str.strip())
    input_str = re.sub(r'[\x00-\x1f\x7f-\x9f]', '', input_str)
    ptrs = [
        r'`',
        r'$\(',
        r'$',
        r'\|',
        r'\r',
        r'\n',
    ]
    for ptr in ptrs:
        input_str = input_str.replace(ptr, '')
    if len(input_str) > TOKEN_LIMIT:
        input_str = input_str[:TOKEN_LIMIT]
    return input_str

def list_models(timeout: float) -> Dict[str, any]:
    """
    Fetch installed models from Ollama API
    """
    headers = {"Content-Type": "application/json"}
    try:
        with requests.get(f"{OLLAMA_API_URL}/tags", headers=headers, timeout=timeout) as response:
            response.raise_for_status() # Raise HTTPError for bad responses
            try:
                data = response.json()
                if 'models' in data and data['models']:
                    models: list[str] = []
                    for model in data.get("models", []):
                        # Filter out embedding models
                        if "embed" not in model.get("name", ""):
                            # We only need model name
                            models.append(model["name"])
                    return {"models": models}
                return {"models": []}
            except json.JSONDecodeError:
                return {"error": "Can't decode response", "models": []}
    except requests.exceptions.ConnectionError:
        return {"error": "Connection Failed", "models": []}

# --- MAIN EXECUTION BLOCK ---
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Run Ollama benchmark with optional model name.")
    parser.add_argument("-m", "--model", default=MODEL_NAME, help="LLM model name to benchmark")
    parser.add_argument("-j", "--json", action="store_true", help="Output results as JSON only")
    parser.add_argument("-p", "--prompt", default=PROMPT, help="Use custom prompt")
    parser.add_argument("-l", "--ls", action="store_true", help="List available models")
    parser.add_argument("-i", "--interactive", action="store_true", help="Run in interactive mode")
    parser.add_argument("-t", "--timeout", default=TIMEOUT, help="Set request timeout in seconds (default 30s)")
    args = parser.parse_args()
    
    # Run the benchmark multiple times to average out system jitter
    results = []
    selected_model: str = args.model
    prompt = args.prompt
    response = list_models(args.timeout)
    use_json = args.json and not args.interactive
    models: list[str] = sorted(response["models"])
    
    if args.ls or args.interactive:
        if use_json:
            print(json.dumps(response, indent=2))
            exit(0)
        else:
            pos = 1
            i = 1
            for model in models:
                name = f"\033[1;3m{model} *\033[0m" if model == selected_model else model
                print(f"{pos}\033[0;32m・\033[0m{name}")
                pos += 1
            if args.interactive:
                print("--------")
                choice = input(f"Select model [1-{len(models)}]: ")
                if choice:
                    i = parse_int(choice) or -1
                    if i >= 0 and i < len(models):
                        selected_model = models[i]
                    else:
                        print("\n\033[0;31m[ERROR]\033[0m Invalid input")
                        exit(1)
                print("="*40)
                custom_prompt = sanitize_input(
                    input("Custom prompt (press return for default): "))
                if custom_prompt:
                    prompt = custom_prompt
            else:
                exit(0)
                
    if selected_model not in models:
        if use_json:
            out = {"error": "Model not found", "ttft": 0.0, "tps": 0.0, "model": selected_model, "total_tokens": 0, "total_duration": 0.0}
            print(json.dumps(out, indent=2))
        else:
            print("\n\033[0;31m[ERROR]\033[0m Model not found")
        exit(1)
        
    if not use_json:
        print("\n\033[0;32m========================================================\033[0m")
        print(f"RUNNING {NUM_RUNS} BENCHMARK ITERATIONS...")
        print("\033[0;32m========================================================\033[0m")
    
    for i in range(NUM_RUNS):
        if not use_json:
            print(f"\n\n\033[0;32m=============\033[0m RUN {i+1}/{NUM_RUNS} \033[0;32m=============\033[0m")
        # Run the benchmark and store the results
        result = benchmark_ollama(selected_model, prompt, MAX_TOKENS, args.timeout, use_json)
        results.append(result)
        time.sleep(SLEEP_BETWEEN_RUNS) # Wait a moment between runs for system stability
    
    # --- FINAL REPORTING ---
    if not use_json:
        print("\n\033[0;32m========================================================\033[0m")
        print("                ✅ BENCHMARK SUMMARY ✅                ")
        print("\033[0;32m========================================================\033[0m")
    
    # Average the results
    successful_results = []
    errors = []
    
    for result in results:
        if 'error' not in result:
            successful_results.append(result)
        else:
            errors.append(result)
            if not use_json:
                print(f"\n\033[0;31m[ERROR]\033[0m Skipped from averages: {result['error']}")
    
    avg_results = {
        "model": successful_results[0]["model"],
        "avg_ttft": sum(r['ttft'] for r in successful_results) / len(successful_results) if successful_results else 0,
        "avg_tps": sum(r['tps'] for r in successful_results) / len(successful_results) if successful_results else 0,
        "avg_total_tokens": sum(r['total_tokens'] for r in successful_results) / len(successful_results) if successful_results else 0,
        "errors": errors,
        "total_attempts": len(results)
    }
    
    if use_json:
        print(json.dumps(avg_results, indent=2))
    else:
        print(f"Model Tested: {avg_results['model']}")
        print(f"Average Tokens Generated: {avg_results['avg_total_tokens']:.0f}")
        print("-" * 40)
        print(f"🚀 Average Tokens Per Second (TPS): {avg_results['avg_tps']:.2f} tokens/sec")
        print(f"⏱️ Average Time to First Token (TTFT): {avg_results['avg_ttft']:.3f} seconds")
        print("\033[0;32m========================================================\033[0m")