ollama_benchmark.py

#!/usr/bin/env python3
"""
Ollama Model Benchmark Tool
Modern TUI für Geschwindigkeits-Tests von Ollama-Modellen
"""

import asyncio
import time
import json
from datetime import datetime
from typing import List, Dict, Optional
from dataclasses import dataclass, asdict

import requests
from textual.app import App, ComposeResult
from textual.widgets import (
    Header, Footer, DataTable, ProgressBar, Static, 
    Button, Input, Select, Label, Log, Collapsible
)
from textual.containers import Container, Horizontal, Vertical, ScrollableContainer
from textual.reactive import reactive
from textual import work
from rich.table import Table
from rich.panel import Panel
from rich.text import Text

# Ollama API Base URL
OLLAMA_BASE_URL = "http://localhost:11434"

@dataclass
class BenchmarkResult:
    """Speichert Benchmark-Ergebnisse für ein Modell"""
    model: str
    prompt: str
    total_time: float
    tokens_generated: int
    tokens_per_second: float
    first_token_time: float
    avg_token_time: float
    timestamp: str

class OllamaBenchmark:
    """Ollama API Wrapper für Benchmarking"""
    
    def __init__(self, base_url: str = OLLAMA_BASE_URL):
        self.base_url = base_url
    
    def list_models(self) -> List[str]:
        """Listet alle verfügbaren Ollama-Modelle"""
        try:
            response = requests.get(f"{self.base_url}/api/tags", timeout=5)
            response.raise_for_status()
            models = response.json().get("models", [])
            return [model["name"] for model in models]
        except Exception as e:
            print(f"Error listing models: {e}")
            return []
    
    def benchmark_model(
        self, 
        model: str, 
        prompt: str = "Write a short story about a robot learning to paint.",
        max_tokens: int = 100
    ) -> Optional[BenchmarkResult]:
        """Benchmarkt ein Modell und gibt Metriken zurück"""
        try:
            start_time = time.time()
            first_token_time = None
            
            response = requests.post(
                f"{self.base_url}/api/generate",
                json={
                    "model": model,
                    "prompt": prompt,
                    "stream": True,
                    "options": {
                        "num_predict": max_tokens
                    }
                },
                stream=True,
                timeout=300
            )
            response.raise_for_status()
            
            tokens_generated = 0
            full_response = ""
            
            for line in response.iter_lines():
                if line:
                    try:
                        data = json.loads(line)
                        
                        # Erste Token-Zeit messen
                        if first_token_time is None and data.get("response"):
                            first_token_time = time.time() - start_time
                        
                        if data.get("response"):
                            tokens_generated += 1
                            full_response += data["response"]
                        
                        # Beende wenn fertig
                        if data.get("done", False):
                            break
                    except json.JSONDecodeError:
                        continue
            
            total_time = time.time() - start_time
            
            if tokens_generated == 0:
                return None
            
            tokens_per_second = tokens_generated / total_time if total_time > 0 else 0
            avg_token_time = total_time / tokens_generated if tokens_generated > 0 else 0
            
            return BenchmarkResult(
                model=model,
                prompt=prompt,
                total_time=total_time,
                tokens_generated=tokens_generated,
                tokens_per_second=tokens_per_second,
                first_token_time=first_token_time or 0,
                avg_token_time=avg_token_time,
                timestamp=datetime.now().isoformat()
            )
        except Exception as e:
            print(f"Error benchmarking {model}: {e}")
            return None

class BenchmarkApp(App):
    """Haupt-TUI App für Ollama Benchmarking"""
    
    CSS = """
    Screen {
        background: $surface;
    }
    
    #results-container {
        height: 1fr;
        border: solid $primary;
    }
    
    #model-select {
        width: 1fr;
    }
    
    #prompt-input {
        width: 1fr;
    }
    
    .benchmark-button {
        width: 1fr;
    }
    
    DataTable {
        height: 1fr;
    }
    
    ProgressBar {
        height: 1;
    }
    
    .status-label {
        text-style: bold;
        color: $accent;
    }
    """
    
    BINDINGS = [
        ("q", "quit", "Quit"),
        ("r", "run_benchmark", "Run Benchmark"),
        ("c", "clear_results", "Clear Results"),
        ("s", "save_results", "Save Results"),
    ]
    
    current_model = reactive("")
    current_prompt = reactive("Write a short story about a robot learning to paint.")
    benchmark_running = reactive(False)
    results: List[BenchmarkResult] = reactive([])
    
    def compose(self) -> ComposeResult:
        """Erstellt die UI-Struktur"""
        yield Header(show_clock=True)
        
        with Container(id="main-container"):
            with Vertical(id="controls"):
                yield Label("🎯 Ollama Model Benchmark Tool", classes="status-label")
                
                with Horizontal():
                    yield Label("Model:", classes="label")
                    yield Select(id="model-select", options=[], prompt="Select Model...")
                
                with Horizontal():
                    yield Label("Prompt:", classes="label")
                    yield Input(
                        id="prompt-input",
                        value=self.current_prompt,
                        placeholder="Enter test prompt..."
                    )
                
                with Horizontal():
                    yield Button("🚀 Run Benchmark", id="run-btn", variant="primary")
                    yield Button("🗑️ Clear Results", id="clear-btn", variant="default")
                    yield Button("💾 Save Results", id="save-btn", variant="default")
                
                yield ProgressBar(id="progress", show_eta=False, total=100)
                yield Label("Ready. Select a model and press 'Run Benchmark'", id="status")
            
            with ScrollableContainer(id="results-container"):
                yield DataTable(id="results-table")
                yield Log(id="log", max_lines=50)
        
        yield Footer()
    
    def on_mount(self) -> None:
        """Wird beim Start aufgerufen"""
        self.load_models()
        self.setup_table()
    
    def load_models(self) -> None:
        """Lädt verfügbare Modelle"""
        benchmark = OllamaBenchmark()
        models = benchmark.list_models()
        
        select = self.query_one("#model-select", Select)
        select.set_options([(model, model) for model in models])
        
        if models:
            self.notify(f"Loaded {len(models)} models", severity="information")
        else:
            self.notify("No models found. Make sure Ollama is running.", severity="warning")
    
    def setup_table(self) -> None:
        """Initialisiert die Results-Tabelle"""
        table = self.query_one("#results-table", DataTable)
        table.add_columns(
            "Model",
            "Tokens/sec",
            "Total Time (s)",
            "Tokens",
            "First Token (s)",
            "Avg Token (ms)",
            "Timestamp"
        )
        table.cursor_type = "row"
    
    def on_select_changed(self, event: Select.Changed) -> None:
        """Wird aufgerufen wenn ein Modell ausgewählt wird"""
        if event.control.id == "model-select":
            self.current_model = event.value or ""
    
    def on_input_changed(self, event: Input.Changed) -> None:
        """Wird aufgerufen wenn der Prompt geändert wird"""
        if event.control.id == "prompt-input":
            self.current_prompt = event.value
    
    def on_button_pressed(self, event: Button.Pressed) -> None:
        """Behandelt Button-Klicks"""
        if event.button.id == "run-btn":
            self.run_benchmark()
        elif event.button.id == "clear-btn":
            self.clear_results()
        elif event.button.id == "save-btn":
            self.save_results()
    
    def action_run_benchmark(self) -> None:
        """Startet Benchmark (via Binding)"""
        self.run_benchmark()
    
    def action_clear_results(self) -> None:
        """Löscht Results (via Binding)"""
        self.clear_results()
    
    def action_save_results(self) -> None:
        """Speichert Results (via Binding)"""
        self.save_results()
    
    @work(exclusive=True)
    async def run_benchmark(self) -> None:
        """Führt Benchmark aus"""
        if not self.current_model:
            self.notify("Please select a model first", severity="warning")
            return
        
        if self.benchmark_running:
            self.notify("Benchmark already running", severity="warning")
            return
        
        self.benchmark_running = True
        progress = self.query_one("#progress", ProgressBar)
        status = self.query_one("#status", Label)
        log = self.query_one("#log", Log)
        
        # UI Updates
        self.query_one("#run-btn", Button).disabled = True
        progress.update(progress=0)
        status.update(f"🔄 Benchmarking {self.current_model}...")
        log.write(f"[{datetime.now().strftime('%H:%M:%S')}] Starting benchmark for {self.current_model}")
        
        # Benchmark ausführen (in Thread)
        benchmark = OllamaBenchmark()
        
        # Progress Simulation
        async def update_progress():
            for i in range(0, 90, 10):
                await asyncio.sleep(0.1)
                progress.update(progress=i)
        
        progress_task = asyncio.create_task(update_progress())
        
        # Benchmark in Executor (blocking)
        loop = asyncio.get_running_loop()
        result = await loop.run_in_executor(
            None,
            benchmark.benchmark_model,
            self.current_model,
            self.current_prompt,
            100  # max_tokens
        )
        
        progress_task.cancel()
        progress.update(progress=100)
        
        if result:
            self.results.append(result)
            self.update_table(result)
            
            status.update(
                f"✅ {result.model}: {result.tokens_per_second:.2f} tokens/sec "
                f"({result.tokens_generated} tokens in {result.total_time:.2f}s)"
            )
            log.write(
                f"[{datetime.now().strftime('%H:%M:%S')}] ✅ {result.model}: "
                f"{result.tokens_per_second:.2f} tokens/sec, "
                f"{result.total_time:.2f}s total, "
                f"{result.tokens_generated} tokens"
            )
            self.notify(f"Benchmark completed: {result.tokens_per_second:.2f} tokens/sec", severity="success")
        else:
            status.update(f"❌ Benchmark failed for {self.current_model}")
            log.write(f"[{datetime.now().strftime('%H:%M:%S')}] ❌ Benchmark failed")
            self.notify("Benchmark failed. Check logs.", severity="error")
        
        self.query_one("#run-btn", Button).disabled = False
        self.benchmark_running = False
        progress.update(progress=0)
    
    def update_table(self, result: BenchmarkResult) -> None:
        """Aktualisiert die Results-Tabelle"""
        table = self.query_one("#results-table", DataTable)
        table.add_row(
            result.model,
            f"{result.tokens_per_second:.2f}",
            f"{result.total_time:.2f}",
            str(result.tokens_generated),
            f"{result.first_token_time:.3f}",
            f"{result.avg_token_time * 1000:.2f}",
            datetime.fromisoformat(result.timestamp).strftime("%H:%M:%S")
        )
    
    def clear_results(self) -> None:
        """Löscht alle Results"""
        table = self.query_one("#results-table", DataTable)
        table.clear()
        self.results.clear()
        self.query_one("#status", Label).update("Results cleared")
        self.query_one("#log", Log).clear()
        self.notify("Results cleared", severity="information")
    
    def save_results(self) -> None:
        """Speichert Results als JSON"""
        if not self.results:
            self.notify("No results to save", severity="warning")
            return
        
        filename = f"ollama_benchmark_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
        
        try:
            with open(filename, "w", encoding="utf-8") as f:
                json.dump(
                    [asdict(result) for result in self.results],
                    f,
                    indent=2,
                    ensure_ascii=False
                )
            self.notify(f"Results saved to {filename}", severity="success")
            self.query_one("#log", Log).write(f"💾 Saved results to {filename}")
        except Exception as e:
            self.notify(f"Error saving results: {e}", severity="error")

if __name__ == "__main__":
    app = BenchmarkApp()
    app.run()