1package ollama
2
3import (
4 "context"
5 "fmt"
6 "os/exec"
7 "strings"
8 "time"
9)
10
11// CLI-based approach for Ollama operations
12// These functions use the ollama CLI instead of HTTP requests
13
14// CLIListModels lists available models using ollama CLI
15func CLIListModels(ctx context.Context) ([]OllamaModel, error) {
16 cmd := exec.CommandContext(ctx, "ollama", "list")
17 output, err := cmd.Output()
18 if err != nil {
19 return nil, fmt.Errorf("failed to list models via CLI: %w", err)
20 }
21
22 return parseModelsList(string(output))
23}
24
25// parseModelsList parses the text output from 'ollama list'
26func parseModelsList(output string) ([]OllamaModel, error) {
27 lines := strings.Split(strings.TrimSpace(output), "\n")
28 if len(lines) < 2 {
29 return nil, fmt.Errorf("unexpected output format")
30 }
31
32 var models []OllamaModel
33 // Skip the header line
34 for i := 1; i < len(lines); i++ {
35 line := strings.TrimSpace(lines[i])
36 if line == "" {
37 continue
38 }
39
40 // Parse each line: NAME ID SIZE MODIFIED
41 fields := strings.Fields(line)
42 if len(fields) >= 4 {
43 name := fields[0]
44 models = append(models, OllamaModel{
45 Name: name,
46 Model: name,
47 Size: 0, // Size parsing from text is complex, skip for now
48 })
49 }
50 }
51
52 return models, nil
53}
54
55// CLIListRunningModels lists currently running models using ollama CLI
56func CLIListRunningModels(ctx context.Context) ([]string, error) {
57 cmd := exec.CommandContext(ctx, "ollama", "ps")
58 output, err := cmd.Output()
59 if err != nil {
60 return nil, fmt.Errorf("failed to list running models via CLI: %w", err)
61 }
62
63 return parseRunningModelsList(string(output))
64}
65
66// parseRunningModelsList parses the text output from 'ollama ps'
67func parseRunningModelsList(output string) ([]string, error) {
68 lines := strings.Split(strings.TrimSpace(output), "\n")
69 if len(lines) < 2 {
70 return []string{}, nil // No running models
71 }
72
73 var runningModels []string
74 // Skip the header line
75 for i := 1; i < len(lines); i++ {
76 line := strings.TrimSpace(lines[i])
77 if line == "" {
78 continue
79 }
80
81 // Parse each line: NAME ID SIZE PROCESSOR UNTIL
82 fields := strings.Fields(line)
83 if len(fields) >= 1 {
84 name := fields[0]
85 if name != "" {
86 runningModels = append(runningModels, name)
87 }
88 }
89 }
90
91 return runningModels, nil
92}
93
94// CLIStopModel stops a specific model using ollama CLI
95func CLIStopModel(ctx context.Context, modelName string) error {
96 cmd := exec.CommandContext(ctx, "ollama", "stop", modelName)
97 if err := cmd.Run(); err != nil {
98 return fmt.Errorf("failed to stop model %s via CLI: %w", modelName, err)
99 }
100 return nil
101}
102
103// CLIStopAllModels stops all running models using ollama CLI
104func CLIStopAllModels(ctx context.Context) error {
105 // First get list of running models
106 runningModels, err := CLIListRunningModels(ctx)
107 if err != nil {
108 return fmt.Errorf("failed to get running models: %w", err)
109 }
110
111 // Stop each model individually
112 for _, modelName := range runningModels {
113 if err := CLIStopModel(ctx, modelName); err != nil {
114 return fmt.Errorf("failed to stop model %s: %w", modelName, err)
115 }
116 }
117
118 return nil
119}
120
121// CLIIsModelRunning checks if a specific model is running using ollama CLI
122func CLIIsModelRunning(ctx context.Context, modelName string) (bool, error) {
123 runningModels, err := CLIListRunningModels(ctx)
124 if err != nil {
125 return false, err
126 }
127
128 for _, running := range runningModels {
129 if running == modelName {
130 return true, nil
131 }
132 }
133
134 return false, nil
135}
136
137// CLIStartModel starts a model using ollama CLI (similar to StartModel but using CLI)
138func CLIStartModel(ctx context.Context, modelName string) error {
139 // Use ollama run with a simple prompt that immediately exits
140 cmd := exec.CommandContext(ctx, "ollama", "run", modelName, "--verbose", "hi")
141
142 // Set a shorter timeout for the run command since we just want to load the model
143 runCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
144 defer cancel()
145
146 cmd = exec.CommandContext(runCtx, "ollama", "run", modelName, "hi")
147
148 if err := cmd.Run(); err != nil {
149 return fmt.Errorf("failed to start model %s via CLI: %w", modelName, err)
150 }
151
152 return nil
153}
154
155// CLIGetModelsCount returns the number of available models using CLI
156func CLIGetModelsCount(ctx context.Context) (int, error) {
157 models, err := CLIListModels(ctx)
158 if err != nil {
159 return 0, err
160 }
161 return len(models), nil
162}
163
164// Performance comparison helpers
165
166// BenchmarkCLIvsHTTP compares CLI vs HTTP performance
167func BenchmarkCLIvsHTTP(ctx context.Context) (map[string]time.Duration, error) {
168 results := make(map[string]time.Duration)
169
170 // Test HTTP approach
171 start := time.Now()
172 _, err := GetModels(ctx)
173 if err != nil {
174 return nil, fmt.Errorf("HTTP GetModels failed: %w", err)
175 }
176 results["HTTP_GetModels"] = time.Since(start)
177
178 // Test CLI approach
179 start = time.Now()
180 _, err = CLIListModels(ctx)
181 if err != nil {
182 return nil, fmt.Errorf("CLI ListModels failed: %w", err)
183 }
184 results["CLI_ListModels"] = time.Since(start)
185
186 // Test HTTP running models
187 start = time.Now()
188 _, err = GetRunningModels(ctx)
189 if err != nil {
190 return nil, fmt.Errorf("HTTP GetRunningModels failed: %w", err)
191 }
192 results["HTTP_GetRunningModels"] = time.Since(start)
193
194 // Test CLI running models
195 start = time.Now()
196 _, err = CLIListRunningModels(ctx)
197 if err != nil {
198 return nil, fmt.Errorf("CLI ListRunningModels failed: %w", err)
199 }
200 results["CLI_ListRunningModels"] = time.Since(start)
201
202 return results, nil
203}
204
205// CLICleanupProcesses provides CLI-based cleanup (alternative to HTTP-based cleanup)
206func CLICleanupProcesses(ctx context.Context) error {
207 return CLIStopAllModels(ctx)
208}