Ollama Local AI Testing

from locust import task, HttpUser
import json
import random

class OllamaUser(HttpUser):
    def on_start(self):
        # Ollama typically runs on localhost:11434
        self.base_url = "http://localhost:11434"
        
        # Available models (install with: ollama pull model-name)
        self.models = [
            "llama2:7b",
            "mistral:7b", 
            "codellama:7b",
            "phi:2.7b"
        ]
        
        # Test prompts
        self.prompts = [
            "Write a short product description for a smartphone.",
            "Explain machine learning in simple terms.",
            "Create a brief email to schedule a meeting.",
            "Write a Python function to calculate fibonacci numbers.",
            "Describe the benefits of renewable energy."
        ]

    @task(3)
    def generate_text(self):
        """Generate text using Ollama"""
        model = random.choice(self.models)
        prompt = random.choice(self.prompts)
        
        payload = {
            "model": model,
            "prompt": prompt,
            "stream": False,
            "options": {
                "temperature": 0.7,
                "num_predict": 100
            }
        }
        
        with self.client.post(
            f"{self.base_url}/api/generate",
            json=payload,
            name=f"Generate - {model}"
        ) as response:
            if response.status_code == 200:
                try:
                    data = response.json()
                    response_text = data.get("response", "")
                    done = data.get("done", False)
                    
                    if done and response_text:
                        print(f"{model}: Generated {len(response_text)} characters")
                    else:
                        response.failure("Incomplete response from Ollama")
                        
                except json.JSONDecodeError:
                    response.failure("Invalid JSON response from Ollama")
            else:
                response.failure(f"Ollama error: {response.status_code}")

    @task(2)
    def chat_completion(self):
        """Test chat completion with Ollama"""
        model = random.choice(self.models)
        prompt = random.choice(self.prompts)
        
        payload = {
            "model": model,
            "messages": [
                {"role": "user", "content": prompt}
            ],
            "stream": False,
            "options": {
                "temperature": 0.7,
                "num_predict": 100
            }
        }
        
        with self.client.post(
            f"{self.base_url}/api/chat",
            json=payload,
            name=f"Chat - {model}"
        ) as response:
            if response.status_code == 200:
                try:
                    data = response.json()
                    message = data.get("message", {})
                    content = message.get("content", "")
                    done = data.get("done", False)
                    
                    if done and content:
                        print(f"{model} Chat: {len(content)} characters")
                    else:
                        response.failure("Incomplete chat response")
                        
                except json.JSONDecodeError:
                    response.failure("Invalid JSON response from Ollama")
            else:
                response.failure(f"Ollama chat error: {response.status_code}")

    @task(1)
    def list_models(self):
        """List available models"""
        with self.client.get(
            f"{self.base_url}/api/tags",
            name="List Models"
        ) as response:
            if response.status_code == 200:
                try:
                    data = response.json()
                    models = data.get("models", [])
                    print(f"Available models: {len(models)}")
                    
                    for model in models[:3]:  # Show first 3
                        name = model.get("name", "Unknown")
                        size = model.get("size", 0) / (1024**3)  # Convert to GB
                        print(f"  - {name}: {size:.1f}GB")
                        
                except json.JSONDecodeError:
                    response.failure("Invalid JSON response from Ollama")
            else:
                response.failure(f"Failed to list models: {response.status_code}")

    @task(1)
    def model_info(self):
        """Get information about a specific model"""
        model = random.choice(self.models)
        
        payload = {"name": model}
        
        with self.client.post(
            f"{self.base_url}/api/show",
            json=payload,
            name=f"Model Info - {model}"
        ) as response:
            if response.status_code == 200:
                try:
                    data = response.json()
                    modelfile = data.get("modelfile", "")
                    parameters = data.get("parameters", "")
                    print(f"{model} info retrieved")
                    
                except json.JSONDecodeError:
                    response.failure("Invalid JSON response from Ollama")
            else:
                response.failure(f"Failed to get model info: {response.status_code}")