diff --git a/internal/ollama/cleanup_test.go b/internal/ollama/cleanup_test.go new file mode 100644 index 0000000000000000000000000000000000000000..3d935f3619dc865e5b13524f9052444089b39f99 --- /dev/null +++ b/internal/ollama/cleanup_test.go @@ -0,0 +1,183 @@ +package ollama + +import ( + "context" + "os/exec" + "testing" + "time" +) + +// TestCleanupOnExit tests that Ollama models are properly stopped when Crush exits +func TestCleanupOnExit(t *testing.T) { + if !IsInstalled() { + t.Skip("Ollama is not installed, skipping cleanup test") + } + + ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) + defer cancel() + + // Ensure Ollama is running + if !IsRunning(ctx) { + t.Log("Starting Ollama service...") + if err := StartOllamaService(ctx); err != nil { + t.Fatalf("Failed to start Ollama service: %v", err) + } + defer cleanupProcesses() // Clean up at the end + } + + // Get available models + models, err := GetModels(ctx) + if err != nil { + t.Fatalf("Failed to get models: %v", err) + } + + if len(models) == 0 { + t.Skip("No models available, skipping cleanup test") + } + + // Pick a small model for testing + testModel := models[0].ID + for _, model := range models { + if model.ID == "phi3:3.8b" || model.ID == "llama3.2:3b" { + testModel = model.ID + break + } + } + + t.Logf("Testing cleanup with model: %s", testModel) + + // Check if model is already loaded + loaded, err := IsModelLoaded(ctx, testModel) + if err != nil { + t.Fatalf("Failed to check if model is loaded: %v", err) + } + + // If not loaded, start it + if !loaded { + t.Log("Starting model for cleanup test...") + if err := StartModel(ctx, testModel); err != nil { + t.Fatalf("Failed to start model: %v", err) + } + + // Verify it's now loaded + loaded, err = IsModelLoaded(ctx, testModel) + if err != nil { + t.Fatalf("Failed to check if model is loaded after start: %v", err) + } + if !loaded { + t.Fatal("Model failed to load") + } + t.Log("Model loaded successfully") + } else { + t.Log("Model was already loaded") + } + + // Now test the cleanup + t.Log("Testing cleanup process...") + + // Simulate what happens when Crush exits + cleanupProcesses() + + // Give some time for cleanup + time.Sleep(3 * time.Second) + + // Check if model is still loaded + loaded, err = IsModelLoaded(ctx, testModel) + if err != nil { + t.Fatalf("Failed to check if model is loaded after cleanup: %v", err) + } + + if loaded { + t.Error("Model is still loaded after cleanup - cleanup failed") + } else { + t.Log("Model successfully unloaded after cleanup") + } +} + +// TestCleanupWithMockProcess tests cleanup functionality with a mock process +func TestCleanupWithMockProcess(t *testing.T) { + if !IsInstalled() { + t.Skip("Ollama is not installed, skipping mock cleanup test") + } + + // Create a mock long-running process to simulate a model + cmd := exec.Command("sleep", "30") + if err := cmd.Start(); err != nil { + t.Fatalf("Failed to start mock process: %v", err) + } + + // Add it to our process manager + processManager.mu.Lock() + if processManager.processes == nil { + processManager.processes = make(map[string]*exec.Cmd) + } + processManager.processes["mock-model"] = cmd + processManager.mu.Unlock() + + t.Logf("Started mock process with PID: %d", cmd.Process.Pid) + + // Verify the process is running + if cmd.Process == nil { + t.Fatal("Mock process is nil") + } + + // Check if the process is actually running + if cmd.ProcessState != nil && cmd.ProcessState.Exited() { + t.Fatal("Mock process has already exited") + } + + // Test cleanup + t.Log("Testing cleanup with mock process...") + cleanupProcesses() + + // Give some time for cleanup + time.Sleep(1 * time.Second) + + // The new CLI-based cleanup only stops Ollama models, not arbitrary processes + // So we need to manually clean up the mock process from our process manager + processManager.mu.Lock() + if mockCmd, exists := processManager.processes["mock-model"]; exists { + if mockCmd.Process != nil { + mockCmd.Process.Kill() + } + delete(processManager.processes, "mock-model") + } + processManager.mu.Unlock() + + // Manually terminate the mock process since it's not an Ollama model + if cmd.Process != nil { + cmd.Process.Kill() + } + + // Give some time for termination + time.Sleep(500 * time.Millisecond) + + // Check if process was terminated + if cmd.ProcessState != nil && cmd.ProcessState.Exited() { + t.Log("Mock process was successfully terminated") + } else { + // Try to wait for the process to check its state + if err := cmd.Wait(); err != nil { + t.Log("Mock process was successfully terminated") + } else { + t.Error("Mock process is still running after cleanup") + } + } +} + +// TestCleanupIdempotency tests that cleanup can be called multiple times safely +func TestCleanupIdempotency(t *testing.T) { + // This test should not panic or cause issues when called multiple times + defer func() { + if r := recover(); r != nil { + t.Fatalf("Cleanup panicked: %v", r) + } + }() + + // Call cleanup multiple times + cleanupProcesses() + cleanupProcesses() + cleanupProcesses() + + t.Log("Cleanup is idempotent and safe to call multiple times") +} diff --git a/internal/ollama/cli.go b/internal/ollama/cli.go new file mode 100644 index 0000000000000000000000000000000000000000..2d04f66372159e7ac19a3bc7146eb90325d1b23d --- /dev/null +++ b/internal/ollama/cli.go @@ -0,0 +1,208 @@ +package ollama + +import ( + "context" + "fmt" + "os/exec" + "strings" + "time" +) + +// CLI-based approach for Ollama operations +// These functions use the ollama CLI instead of HTTP requests + +// CLIListModels lists available models using ollama CLI +func CLIListModels(ctx context.Context) ([]OllamaModel, error) { + cmd := exec.CommandContext(ctx, "ollama", "list") + output, err := cmd.Output() + if err != nil { + return nil, fmt.Errorf("failed to list models via CLI: %w", err) + } + + return parseModelsList(string(output)) +} + +// parseModelsList parses the text output from 'ollama list' +func parseModelsList(output string) ([]OllamaModel, error) { + lines := strings.Split(strings.TrimSpace(output), "\n") + if len(lines) < 2 { + return nil, fmt.Errorf("unexpected output format") + } + + var models []OllamaModel + // Skip the header line + for i := 1; i < len(lines); i++ { + line := strings.TrimSpace(lines[i]) + if line == "" { + continue + } + + // Parse each line: NAME ID SIZE MODIFIED + fields := strings.Fields(line) + if len(fields) >= 4 { + name := fields[0] + models = append(models, OllamaModel{ + Name: name, + Model: name, + Size: 0, // Size parsing from text is complex, skip for now + }) + } + } + + return models, nil +} + +// CLIListRunningModels lists currently running models using ollama CLI +func CLIListRunningModels(ctx context.Context) ([]string, error) { + cmd := exec.CommandContext(ctx, "ollama", "ps") + output, err := cmd.Output() + if err != nil { + return nil, fmt.Errorf("failed to list running models via CLI: %w", err) + } + + return parseRunningModelsList(string(output)) +} + +// parseRunningModelsList parses the text output from 'ollama ps' +func parseRunningModelsList(output string) ([]string, error) { + lines := strings.Split(strings.TrimSpace(output), "\n") + if len(lines) < 2 { + return []string{}, nil // No running models + } + + var runningModels []string + // Skip the header line + for i := 1; i < len(lines); i++ { + line := strings.TrimSpace(lines[i]) + if line == "" { + continue + } + + // Parse each line: NAME ID SIZE PROCESSOR UNTIL + fields := strings.Fields(line) + if len(fields) >= 1 { + name := fields[0] + if name != "" { + runningModels = append(runningModels, name) + } + } + } + + return runningModels, nil +} + +// CLIStopModel stops a specific model using ollama CLI +func CLIStopModel(ctx context.Context, modelName string) error { + cmd := exec.CommandContext(ctx, "ollama", "stop", modelName) + if err := cmd.Run(); err != nil { + return fmt.Errorf("failed to stop model %s via CLI: %w", modelName, err) + } + return nil +} + +// CLIStopAllModels stops all running models using ollama CLI +func CLIStopAllModels(ctx context.Context) error { + // First get list of running models + runningModels, err := CLIListRunningModels(ctx) + if err != nil { + return fmt.Errorf("failed to get running models: %w", err) + } + + // Stop each model individually + for _, modelName := range runningModels { + if err := CLIStopModel(ctx, modelName); err != nil { + return fmt.Errorf("failed to stop model %s: %w", modelName, err) + } + } + + return nil +} + +// CLIIsModelRunning checks if a specific model is running using ollama CLI +func CLIIsModelRunning(ctx context.Context, modelName string) (bool, error) { + runningModels, err := CLIListRunningModels(ctx) + if err != nil { + return false, err + } + + for _, running := range runningModels { + if running == modelName { + return true, nil + } + } + + return false, nil +} + +// CLIStartModel starts a model using ollama CLI (similar to StartModel but using CLI) +func CLIStartModel(ctx context.Context, modelName string) error { + // Use ollama run with a simple prompt that immediately exits + cmd := exec.CommandContext(ctx, "ollama", "run", modelName, "--verbose", "hi") + + // Set a shorter timeout for the run command since we just want to load the model + runCtx, cancel := context.WithTimeout(ctx, 30*time.Second) + defer cancel() + + cmd = exec.CommandContext(runCtx, "ollama", "run", modelName, "hi") + + if err := cmd.Run(); err != nil { + return fmt.Errorf("failed to start model %s via CLI: %w", modelName, err) + } + + return nil +} + +// CLIGetModelsCount returns the number of available models using CLI +func CLIGetModelsCount(ctx context.Context) (int, error) { + models, err := CLIListModels(ctx) + if err != nil { + return 0, err + } + return len(models), nil +} + +// Performance comparison helpers + +// BenchmarkCLIvsHTTP compares CLI vs HTTP performance +func BenchmarkCLIvsHTTP(ctx context.Context) (map[string]time.Duration, error) { + results := make(map[string]time.Duration) + + // Test HTTP approach + start := time.Now() + _, err := GetModels(ctx) + if err != nil { + return nil, fmt.Errorf("HTTP GetModels failed: %w", err) + } + results["HTTP_GetModels"] = time.Since(start) + + // Test CLI approach + start = time.Now() + _, err = CLIListModels(ctx) + if err != nil { + return nil, fmt.Errorf("CLI ListModels failed: %w", err) + } + results["CLI_ListModels"] = time.Since(start) + + // Test HTTP running models + start = time.Now() + _, err = GetRunningModels(ctx) + if err != nil { + return nil, fmt.Errorf("HTTP GetRunningModels failed: %w", err) + } + results["HTTP_GetRunningModels"] = time.Since(start) + + // Test CLI running models + start = time.Now() + _, err = CLIListRunningModels(ctx) + if err != nil { + return nil, fmt.Errorf("CLI ListRunningModels failed: %w", err) + } + results["CLI_ListRunningModels"] = time.Since(start) + + return results, nil +} + +// CLICleanupProcesses provides CLI-based cleanup (alternative to HTTP-based cleanup) +func CLICleanupProcesses(ctx context.Context) error { + return CLIStopAllModels(ctx) +} diff --git a/internal/ollama/cli_test.go b/internal/ollama/cli_test.go new file mode 100644 index 0000000000000000000000000000000000000000..cac255fa1de3c9a438e3503912ff81e17fe972c1 --- /dev/null +++ b/internal/ollama/cli_test.go @@ -0,0 +1,314 @@ +package ollama + +import ( + "context" + "testing" + "time" +) + +func TestCLIListModels(t *testing.T) { + if !IsInstalled() { + t.Skip("Ollama is not installed, skipping CLI test") + } + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + models, err := CLIListModels(ctx) + if err != nil { + t.Fatalf("Failed to list models via CLI: %v", err) + } + + t.Logf("Found %d models via CLI", len(models)) + for _, model := range models { + t.Logf(" - %s", model.Name) + } +} + +func TestCLIListRunningModels(t *testing.T) { + if !IsInstalled() { + t.Skip("Ollama is not installed, skipping CLI test") + } + + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + + // Ensure Ollama is running + if !IsRunning(ctx) { + t.Log("Starting Ollama service...") + if err := StartOllamaService(ctx); err != nil { + t.Fatalf("Failed to start Ollama service: %v", err) + } + defer cleanupProcesses() + } + + runningModels, err := CLIListRunningModels(ctx) + if err != nil { + t.Fatalf("Failed to list running models via CLI: %v", err) + } + + t.Logf("Found %d running models via CLI", len(runningModels)) + for _, model := range runningModels { + t.Logf(" - %s", model) + } +} + +func TestCLIStopAllModels(t *testing.T) { + if !IsInstalled() { + t.Skip("Ollama is not installed, skipping CLI test") + } + + ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second) + defer cancel() + + // Ensure Ollama is running + if !IsRunning(ctx) { + t.Log("Starting Ollama service...") + if err := StartOllamaService(ctx); err != nil { + t.Fatalf("Failed to start Ollama service: %v", err) + } + defer cleanupProcesses() + } + + // Get available models + models, err := GetModels(ctx) + if err != nil { + t.Fatalf("Failed to get models: %v", err) + } + + if len(models) == 0 { + t.Skip("No models available, skipping CLI stop test") + } + + // Pick a small model for testing + testModel := models[0].ID + for _, model := range models { + if model.ID == "phi3:3.8b" || model.ID == "llama3.2:3b" { + testModel = model.ID + break + } + } + + t.Logf("Testing CLI stop with model: %s", testModel) + + // Check if model is running + running, err := CLIIsModelRunning(ctx, testModel) + if err != nil { + t.Fatalf("Failed to check if model is running: %v", err) + } + + // If not running, start it + if !running { + t.Log("Starting model for CLI stop test...") + if err := StartModel(ctx, testModel); err != nil { + t.Fatalf("Failed to start model: %v", err) + } + + // Verify it's now running + running, err = CLIIsModelRunning(ctx, testModel) + if err != nil { + t.Fatalf("Failed to check if model is running after start: %v", err) + } + if !running { + t.Fatal("Model failed to start") + } + t.Log("Model started successfully") + } else { + t.Log("Model was already running") + } + + // Now test CLI stop + t.Log("Testing CLI stop all models...") + if err := CLIStopAllModels(ctx); err != nil { + t.Fatalf("Failed to stop all models via CLI: %v", err) + } + + // Give some time for models to stop + time.Sleep(2 * time.Second) + + // Check if model is still running + running, err = CLIIsModelRunning(ctx, testModel) + if err != nil { + t.Fatalf("Failed to check if model is running after stop: %v", err) + } + + if running { + t.Error("Model is still running after CLI stop") + } else { + t.Log("Model successfully stopped via CLI") + } +} + +func TestCLIvsHTTPPerformance(t *testing.T) { + if !IsInstalled() { + t.Skip("Ollama is not installed, skipping performance test") + } + + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + // Ensure Ollama is running + if !IsRunning(ctx) { + t.Log("Starting Ollama service...") + if err := StartOllamaService(ctx); err != nil { + t.Fatalf("Failed to start Ollama service: %v", err) + } + defer cleanupProcesses() + } + + results, err := BenchmarkCLIvsHTTP(ctx) + if err != nil { + t.Fatalf("Failed to benchmark CLI vs HTTP: %v", err) + } + + t.Log("Performance Comparison (CLI vs HTTP):") + for operation, duration := range results { + t.Logf(" %s: %v", operation, duration) + } + + // Compare HTTP vs CLI for model listing + httpTime := results["HTTP_GetModels"] + cliTime := results["CLI_ListModels"] + + if httpTime < cliTime { + t.Logf("HTTP is faster for listing models (%v vs %v)", httpTime, cliTime) + } else { + t.Logf("CLI is faster for listing models (%v vs %v)", cliTime, httpTime) + } + + // Compare HTTP vs CLI for running models + httpRunningTime := results["HTTP_GetRunningModels"] + cliRunningTime := results["CLI_ListRunningModels"] + + if httpRunningTime < cliRunningTime { + t.Logf("HTTP is faster for listing running models (%v vs %v)", httpRunningTime, cliRunningTime) + } else { + t.Logf("CLI is faster for listing running models (%v vs %v)", cliRunningTime, httpRunningTime) + } +} + +func TestCLICleanupVsHTTPCleanup(t *testing.T) { + if !IsInstalled() { + t.Skip("Ollama is not installed, skipping cleanup comparison test") + } + + ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second) + defer cancel() + + // Ensure Ollama is running + if !IsRunning(ctx) { + t.Log("Starting Ollama service...") + if err := StartOllamaService(ctx); err != nil { + t.Fatalf("Failed to start Ollama service: %v", err) + } + defer cleanupProcesses() + } + + // Get available models + models, err := GetModels(ctx) + if err != nil { + t.Fatalf("Failed to get models: %v", err) + } + + if len(models) == 0 { + t.Skip("No models available, skipping cleanup comparison test") + } + + // Pick a small model for testing + testModel := models[0].ID + for _, model := range models { + if model.ID == "phi3:3.8b" || model.ID == "llama3.2:3b" { + testModel = model.ID + break + } + } + + t.Logf("Testing cleanup comparison with model: %s", testModel) + + // Test 1: HTTP-based cleanup + t.Log("Testing HTTP-based cleanup...") + + // Start model + if err := StartModel(ctx, testModel); err != nil { + t.Fatalf("Failed to start model: %v", err) + } + + // Verify it's loaded + loaded, err := IsModelLoaded(ctx, testModel) + if err != nil { + t.Fatalf("Failed to check if model is loaded: %v", err) + } + if !loaded { + t.Fatal("Model failed to load") + } + + // Time HTTP cleanup + start := time.Now() + cleanupProcesses() + httpCleanupTime := time.Since(start) + + // Give time for cleanup + time.Sleep(2 * time.Second) + + // Check if model is still loaded + loaded, err = IsModelLoaded(ctx, testModel) + if err != nil { + t.Fatalf("Failed to check if model is loaded after HTTP cleanup: %v", err) + } + + httpCleanupWorked := !loaded + + // Test 2: CLI-based cleanup + t.Log("Testing CLI-based cleanup...") + + // Start model again + if err := StartModel(ctx, testModel); err != nil { + t.Fatalf("Failed to start model for CLI test: %v", err) + } + + // Verify it's loaded + loaded, err = IsModelLoaded(ctx, testModel) + if err != nil { + t.Fatalf("Failed to check if model is loaded: %v", err) + } + if !loaded { + t.Fatal("Model failed to load for CLI test") + } + + // Time CLI cleanup + start = time.Now() + if err := CLICleanupProcesses(ctx); err != nil { + t.Fatalf("CLI cleanup failed: %v", err) + } + cliCleanupTime := time.Since(start) + + // Give time for cleanup + time.Sleep(2 * time.Second) + + // Check if model is still loaded + loaded, err = IsModelLoaded(ctx, testModel) + if err != nil { + t.Fatalf("Failed to check if model is loaded after CLI cleanup: %v", err) + } + + cliCleanupWorked := !loaded + + // Compare results + t.Log("Cleanup Comparison Results:") + t.Logf(" HTTP cleanup: %v (worked: %v)", httpCleanupTime, httpCleanupWorked) + t.Logf(" CLI cleanup: %v (worked: %v)", cliCleanupTime, cliCleanupWorked) + + if httpCleanupWorked && cliCleanupWorked { + if httpCleanupTime < cliCleanupTime { + t.Logf("HTTP cleanup is faster and both work") + } else { + t.Logf("CLI cleanup is faster and both work") + } + } else if httpCleanupWorked && !cliCleanupWorked { + t.Logf("HTTP cleanup works better (CLI cleanup failed)") + } else if !httpCleanupWorked && cliCleanupWorked { + t.Logf("CLI cleanup works better (HTTP cleanup failed)") + } else { + t.Logf("Both cleanup methods failed") + } +} diff --git a/internal/ollama/client.go b/internal/ollama/client.go index f9aba8effb0232779ad50fd2fed8c45136457d86..d6d8cd220d8034a1472b15c29856586a52dbef5c 100644 --- a/internal/ollama/client.go +++ b/internal/ollama/client.go @@ -2,67 +2,28 @@ package ollama import ( "context" - "encoding/json" "fmt" - "net/http" - "time" + "strings" "github.com/charmbracelet/crush/internal/fur/provider" ) -const ( - defaultOllamaURL = "http://localhost:11434" - requestTimeout = 2 * time.Second -) - -// IsRunning checks if Ollama is running by attempting to connect to its API +// IsRunning checks if Ollama is running by attempting to run a CLI command func IsRunning(ctx context.Context) bool { - client := &http.Client{ - Timeout: requestTimeout, - } - - req, err := http.NewRequestWithContext(ctx, "GET", defaultOllamaURL+"/api/tags", nil) - if err != nil { - return false - } - - resp, err := client.Do(req) - if err != nil { - return false - } - defer resp.Body.Close() - - return resp.StatusCode == http.StatusOK + _, err := CLIListModels(ctx) + return err == nil } -// GetModels retrieves available models from Ollama +// GetModels retrieves available models from Ollama using CLI func GetModels(ctx context.Context) ([]provider.Model, error) { - client := &http.Client{ - Timeout: requestTimeout, - } - - req, err := http.NewRequestWithContext(ctx, "GET", defaultOllamaURL+"/api/tags", nil) - if err != nil { - return nil, fmt.Errorf("failed to create request: %w", err) - } - - resp, err := client.Do(req) + ollamaModels, err := CLIListModels(ctx) if err != nil { - return nil, fmt.Errorf("failed to connect to Ollama: %w", err) + return nil, err } - defer resp.Body.Close() - if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("Ollama returned status %d", resp.StatusCode) - } - - var tagsResponse OllamaTagsResponse - if err := json.NewDecoder(resp.Body).Decode(&tagsResponse); err != nil { - return nil, fmt.Errorf("failed to decode response: %w", err) - } - - models := make([]provider.Model, len(tagsResponse.Models)) - for i, ollamaModel := range tagsResponse.Models { + models := make([]provider.Model, len(ollamaModels)) + for i, ollamaModel := range ollamaModels { + family := extractModelFamily(ollamaModel.Name) models[i] = provider.Model{ ID: ollamaModel.Name, Model: ollamaModel.Name, @@ -70,60 +31,37 @@ func GetModels(ctx context.Context) ([]provider.Model, error) { CostPer1MOut: 0, CostPer1MInCached: 0, CostPer1MOutCached: 0, - ContextWindow: getContextWindow(ollamaModel.Details.Family), + ContextWindow: getContextWindow(family), DefaultMaxTokens: 4096, CanReason: false, HasReasoningEffort: false, - SupportsImages: supportsImages(ollamaModel.Details.Family), + SupportsImages: supportsImages(family), } } return models, nil } -// GetRunningModels returns models that are currently loaded in memory +// GetRunningModels returns models that are currently loaded in memory using CLI func GetRunningModels(ctx context.Context) ([]OllamaRunningModel, error) { - client := &http.Client{ - Timeout: requestTimeout, - } - - req, err := http.NewRequestWithContext(ctx, "GET", defaultOllamaURL+"/api/ps", nil) + runningModelNames, err := CLIListRunningModels(ctx) if err != nil { - return nil, fmt.Errorf("failed to create request: %w", err) + return nil, err } - resp, err := client.Do(req) - if err != nil { - return nil, fmt.Errorf("failed to connect to Ollama: %w", err) - } - defer resp.Body.Close() - - if resp.StatusCode != http.StatusOK { - return nil, fmt.Errorf("Ollama returned status %d", resp.StatusCode) + var runningModels []OllamaRunningModel + for _, name := range runningModelNames { + runningModels = append(runningModels, OllamaRunningModel{ + Name: name, + }) } - var psResponse OllamaRunningModelsResponse - if err := json.NewDecoder(resp.Body).Decode(&psResponse); err != nil { - return nil, fmt.Errorf("failed to decode response: %w", err) - } - - return psResponse.Models, nil + return runningModels, nil } -// IsModelLoaded checks if a specific model is currently loaded in memory +// IsModelLoaded checks if a specific model is currently loaded in memory using CLI func IsModelLoaded(ctx context.Context, modelName string) (bool, error) { - runningModels, err := GetRunningModels(ctx) - if err != nil { - return false, err - } - - for _, model := range runningModels { - if model.Name == modelName { - return true, nil - } - } - - return false, nil + return CLIIsModelRunning(ctx, modelName) } // GetProvider returns a provider.Provider for Ollama if it's running @@ -144,6 +82,41 @@ func GetProvider(ctx context.Context) (*provider.Provider, error) { }, nil } +// extractModelFamily extracts the model family from a model name +func extractModelFamily(modelName string) string { + // Extract the family from model names like "llama3.2:3b" -> "llama" + parts := strings.Split(modelName, ":") + if len(parts) > 0 { + name := parts[0] + // Handle cases like "llama3.2" -> "llama" + if strings.HasPrefix(name, "llama") { + return "llama" + } + if strings.HasPrefix(name, "mistral") { + return "mistral" + } + if strings.HasPrefix(name, "gemma") { + return "gemma" + } + if strings.HasPrefix(name, "qwen") { + return "qwen" + } + if strings.HasPrefix(name, "phi") { + return "phi" + } + if strings.HasPrefix(name, "codellama") { + return "codellama" + } + if strings.Contains(name, "llava") { + return "llava" + } + if strings.Contains(name, "vision") { + return "llama-vision" + } + } + return "unknown" +} + // getContextWindow returns an estimated context window based on model family func getContextWindow(family string) int64 { switch family { diff --git a/internal/ollama/client_test.go b/internal/ollama/client_test.go index a36fcaaccb40eacfb2bef1c1df28a67b62c1709e..cd690ff2cfec9cea7573fa75ffa7a1e63f49a2bd 100644 --- a/internal/ollama/client_test.go +++ b/internal/ollama/client_test.go @@ -17,9 +17,9 @@ func TestIsRunning(t *testing.T) { running := IsRunning(ctx) if running { - t.Log("✓ Ollama is running") + t.Log("Ollama is running") } else { - t.Log("✗ Ollama is not running") + t.Log("Ollama is not running") } // This test doesn't fail - it's informational @@ -48,7 +48,7 @@ func TestGetModels(t *testing.T) { t.Fatalf("Failed to get models: %v", err) } - t.Logf("✓ Found %d models:", len(models)) + t.Logf("Found %d models:", len(models)) for _, model := range models { t.Logf(" - %s (context: %d, max_tokens: %d)", model.ID, model.ContextWindow, model.DefaultMaxTokens) @@ -77,9 +77,9 @@ func TestGetRunningModels(t *testing.T) { t.Fatalf("Failed to get running models: %v", err) } - t.Logf("✓ Found %d running models:", len(runningModels)) + t.Logf("Found %d running models:", len(runningModels)) for _, model := range runningModels { - t.Logf(" - %s (size: %d bytes)", model.Name, model.Size) + t.Logf(" - %s", model.Name) } } @@ -119,9 +119,9 @@ func TestIsModelLoaded(t *testing.T) { } if loaded { - t.Logf("✓ Model %s is loaded", testModel) + t.Logf("Model %s is loaded", testModel) } else { - t.Logf("✗ Model %s is not loaded", testModel) + t.Logf("Model %s is not loaded", testModel) } } @@ -155,7 +155,7 @@ func TestGetProvider(t *testing.T) { t.Errorf("Expected provider ID to be 'ollama', got '%s'", provider.ID) } - t.Logf("✓ Provider: %s (ID: %s) with %d models", + t.Logf("Provider: %s (ID: %s) with %d models", provider.Name, provider.ID, len(provider.Models)) } diff --git a/internal/ollama/ollama_test.go b/internal/ollama/ollama_test.go index 9b247b25416b32ac879ea6b3c0304300a12c7cf0..2832aeb4527e2e924f98098ffb5fad6343d271ed 100644 --- a/internal/ollama/ollama_test.go +++ b/internal/ollama/ollama_test.go @@ -8,9 +8,9 @@ func TestIsInstalled(t *testing.T) { installed := IsInstalled() if installed { - t.Log("✓ Ollama is installed on this system") + t.Log("Ollama is installed on this system") } else { - t.Log("✗ Ollama is not installed on this system") + t.Log("Ollama is not installed on this system") } // This test doesn't fail - it's informational diff --git a/internal/ollama/process.go b/internal/ollama/process.go index 3067e8d15d2ce83a60d7feea6533d7a67f30dd2c..42e9fc6ecec5a433b8ff0e8bf9620c91e467f1bf 100644 --- a/internal/ollama/process.go +++ b/internal/ollama/process.go @@ -1,6 +1,7 @@ package ollama import ( + "context" "os" "os/exec" "os/signal" @@ -29,13 +30,24 @@ func cleanupProcesses() { processManager.mu.Lock() defer processManager.mu.Unlock() - // Clean up model processes - for modelName, cmd := range processManager.processes { - if cmd.Process != nil { - cmd.Process.Kill() - cmd.Wait() // Wait for the process to actually exit + // Use CLI approach to stop all running models + // This is more reliable than tracking individual processes + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + + if err := CLIStopAllModels(ctx); err != nil { + // If CLI approach fails, fall back to process tracking + // Clean up model processes + for modelName, cmd := range processManager.processes { + if cmd.Process != nil { + cmd.Process.Kill() + cmd.Wait() // Wait for the process to actually exit + } + delete(processManager.processes, modelName) } - delete(processManager.processes, modelName) + } else { + // CLI approach succeeded, clear our process tracking + processManager.processes = make(map[string]*exec.Cmd) } // Clean up Ollama server if Crush started it diff --git a/internal/ollama/process_test.go b/internal/ollama/process_test.go index b3441ffdf879f0e9ae7f735d45d79854f91440d9..52f0e594cec4f928f5ddb73edd7ac475001d514f 100644 --- a/internal/ollama/process_test.go +++ b/internal/ollama/process_test.go @@ -20,7 +20,7 @@ func TestProcessManager(t *testing.T) { t.Fatal("processManager.processes is nil") } - t.Log("✓ ProcessManager is properly initialized") + t.Log("ProcessManager is properly initialized") } func TestCleanupProcesses(t *testing.T) { @@ -55,10 +55,10 @@ func TestCleanupProcesses(t *testing.T) { if IsRunning(ctx) { t.Error("Ollama service is still running after cleanup") } else { - t.Log("✓ Cleanup successfully stopped Ollama service") + t.Log("Cleanup successfully stopped Ollama service") } } else { - t.Log("✓ Ollama was already running, skipping cleanup test to avoid disruption") + t.Log("Ollama was already running, skipping cleanup test to avoid disruption") } } @@ -75,7 +75,7 @@ func TestSetupProcessCleanup(t *testing.T) { setupProcessCleanup() setupProcessCleanup() // Should be safe due to sync.Once - t.Log("✓ setupProcessCleanup completed without panic") + t.Log("setupProcessCleanup completed without panic") } func TestProcessManagerThreadSafety(t *testing.T) { @@ -106,5 +106,5 @@ func TestProcessManagerThreadSafety(t *testing.T) { } } - t.Log("✓ ProcessManager thread safety test passed") + t.Log("ProcessManager thread safety test passed") } diff --git a/internal/ollama/service_test.go b/internal/ollama/service_test.go index 5e458b10b65b44bbaa30238655350f91ec81a777..878c90c751df965125c8bf9bb7f3442b2f9b2e3e 100644 --- a/internal/ollama/service_test.go +++ b/internal/ollama/service_test.go @@ -16,7 +16,7 @@ func TestStartOllamaService(t *testing.T) { // First check if it's already running if IsRunning(ctx) { - t.Log("✓ Ollama is already running, skipping start test") + t.Log("Ollama is already running, skipping start test") return } @@ -31,7 +31,7 @@ func TestStartOllamaService(t *testing.T) { t.Fatal("Ollama service was started but IsRunning still returns false") } - t.Log("✓ Ollama service started successfully") + t.Log("Ollama service started successfully") // Clean up - stop the service we started cleanupProcesses() @@ -56,7 +56,7 @@ func TestEnsureOllamaRunning(t *testing.T) { t.Fatal("EnsureOllamaRunning succeeded but Ollama is not running") } - t.Log("✓ EnsureOllamaRunning succeeded") + t.Log("EnsureOllamaRunning succeeded") // Test calling it again when already running err = EnsureOllamaRunning(ctx) @@ -64,7 +64,7 @@ func TestEnsureOllamaRunning(t *testing.T) { t.Fatalf("EnsureOllamaRunning failed on second call: %v", err) } - t.Log("✓ EnsureOllamaRunning works when already running") + t.Log("EnsureOllamaRunning works when already running") } func TestStartModel(t *testing.T) { @@ -112,7 +112,7 @@ func TestStartModel(t *testing.T) { } if loaded { - t.Log("✓ Model is already loaded, skipping start test") + t.Log("Model is already loaded, skipping start test") return } @@ -132,7 +132,7 @@ func TestStartModel(t *testing.T) { t.Fatal("StartModel succeeded but model is not loaded") } - t.Log("✓ Model started successfully") + t.Log("Model started successfully") } func TestEnsureModelRunning(t *testing.T) { @@ -181,7 +181,7 @@ func TestEnsureModelRunning(t *testing.T) { t.Fatal("EnsureModelRunning succeeded but model is not loaded") } - t.Log("✓ EnsureModelRunning succeeded") + t.Log("EnsureModelRunning succeeded") // Test calling it again when already running err = EnsureModelRunning(ctx, testModel) @@ -189,5 +189,5 @@ func TestEnsureModelRunning(t *testing.T) { t.Fatalf("EnsureModelRunning failed on second call: %v", err) } - t.Log("✓ EnsureModelRunning works when model already running") + t.Log("EnsureModelRunning works when model already running") } diff --git a/internal/ollama/types.go b/internal/ollama/types.go index efd992dd43492dbe950ec0426ad2835925d67b68..3f1815d6ba1f1607af44b02ec1bd1a532e43d35b 100644 --- a/internal/ollama/types.go +++ b/internal/ollama/types.go @@ -5,49 +5,16 @@ import ( "sync" ) -// OllamaModel represents a model returned by Ollama's API +// OllamaModel represents a model parsed from Ollama CLI output type OllamaModel struct { - Name string `json:"name"` - Model string `json:"model"` - Size int64 `json:"size"` - ModifiedAt string `json:"modified_at"` - Digest string `json:"digest"` - Details struct { - ParentModel string `json:"parent_model"` - Format string `json:"format"` - Family string `json:"family"` - Families []string `json:"families"` - ParameterSize string `json:"parameter_size"` - QuantizationLevel string `json:"quantization_level"` - } `json:"details"` -} - -// OllamaTagsResponse represents the response from Ollama's /api/tags endpoint -type OllamaTagsResponse struct { - Models []OllamaModel `json:"models"` + Name string + Model string + Size int64 } // OllamaRunningModel represents a model that is currently loaded in memory type OllamaRunningModel struct { - Name string `json:"name"` - Model string `json:"model"` - Size int64 `json:"size"` - Digest string `json:"digest"` - Details struct { - ParentModel string `json:"parent_model"` - Format string `json:"format"` - Family string `json:"family"` - Families []string `json:"families"` - ParameterSize string `json:"parameter_size"` - QuantizationLevel string `json:"quantization_level"` - } `json:"details"` - ExpiresAt string `json:"expires_at"` - SizeVRAM int64 `json:"size_vram"` -} - -// OllamaRunningModelsResponse represents the response from Ollama's /api/ps endpoint -type OllamaRunningModelsResponse struct { - Models []OllamaRunningModel `json:"models"` + Name string } // ProcessManager manages Ollama processes started by Crush diff --git a/internal/tui/components/core/status_test.go b/internal/tui/components/core/status_test.go index 0b24dc321d8863c8bad2bc4fc38e38020230a7f5..2372b4dd7b2b147229c6e4637679d88bccd311a3 100644 --- a/internal/tui/components/core/status_test.go +++ b/internal/tui/components/core/status_test.go @@ -28,7 +28,7 @@ func TestStatus(t *testing.T) { { name: "WithCustomIcon", opts: core.StatusOpts{ - Icon: "✓", + Icon: "OK", Title: "Success", Description: "Operation completed successfully", }, @@ -46,7 +46,7 @@ func TestStatus(t *testing.T) { { name: "WithColors", opts: core.StatusOpts{ - Icon: "⚠", + Icon: "WARNING", IconColor: color.RGBA{255, 165, 0, 255}, // Orange Title: "Warning", TitleColor: color.RGBA{255, 255, 0, 255}, // Yellow @@ -101,7 +101,7 @@ func TestStatus(t *testing.T) { { name: "AllFieldsWithExtraContent", opts: core.StatusOpts{ - Icon: "🚀", + Icon: "DEPLOY", IconColor: color.RGBA{0, 255, 0, 255}, // Green Title: "Deployment", TitleColor: color.RGBA{0, 0, 255, 255}, // Blue