Detailed changes
@@ -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")
+}
@@ -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)
+}
@@ -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")
+ }
+}
@@ -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 {
@@ -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))
}
@@ -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
@@ -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
@@ -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")
}
@@ -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")
}
@@ -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
@@ -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