cmd.go 9.62 KB
Newer Older
Jeffrey Morgan's avatar
Jeffrey Morgan committed
1
2
3
package cmd

import (
Michael Yang's avatar
Michael Yang committed
4
	"bufio"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
5
	"context"
Michael Yang's avatar
Michael Yang committed
6
	"errors"
Bruce MacDonald's avatar
Bruce MacDonald committed
7
	"fmt"
Michael Yang's avatar
Michael Yang committed
8
	"io"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
9
10
	"log"
	"net"
Michael Yang's avatar
Michael Yang committed
11
	"net/http"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
12
	"os"
13
	"path/filepath"
Michael Yang's avatar
Michael Yang committed
14
	"strings"
Michael Yang's avatar
Michael Yang committed
15
	"time"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
16

Michael Yang's avatar
Michael Yang committed
17
	"github.com/chzyer/readline"
Patrick Devine's avatar
Patrick Devine committed
18
19
	"github.com/dustin/go-humanize"
	"github.com/olekukonko/tablewriter"
Michael Yang's avatar
Michael Yang committed
20
21
	"github.com/spf13/cobra"

Jeffrey Morgan's avatar
Jeffrey Morgan committed
22
	"github.com/jmorganca/ollama/api"
Patrick Devine's avatar
Patrick Devine committed
23
	"github.com/jmorganca/ollama/format"
24
	"github.com/jmorganca/ollama/progressbar"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
25
	"github.com/jmorganca/ollama/server"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
26
27
)

28
func CreateHandler(cmd *cobra.Command, args []string) error {
29
	filename, _ := cmd.Flags().GetString("file")
30
31
32
33
34
	filename, err := filepath.Abs(filename)
	if err != nil {
		return err
	}

35
36
	client := api.NewClient()

Michael Yang's avatar
Michael Yang committed
37
38
	var spinner *Spinner

39
40
	request := api.CreateRequest{Name: args[0], Path: filename}
	fn := func(resp api.CreateProgress) error {
Michael Yang's avatar
Michael Yang committed
41
42
43
44
45
46
47
		if spinner != nil {
			spinner.Stop()
		}

		spinner = NewSpinner(resp.Status)
		go spinner.Spin(100 * time.Millisecond)

48
49
50
51
52
53
54
		return nil
	}

	if err := client.Create(context.Background(), &request, fn); err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
55
56
57
58
	if spinner != nil {
		spinner.Stop()
	}

59
60
61
	return nil
}

62
func RunHandler(cmd *cobra.Command, args []string) error {
Patrick Devine's avatar
Patrick Devine committed
63
64
65
66
67
68
69
	mp := server.ParseModelPath(args[0])
	fp, err := mp.GetManifestPath(false)
	if err != nil {
		return err
	}

	_, err = os.Stat(fp)
Michael Yang's avatar
Michael Yang committed
70
71
	switch {
	case errors.Is(err, os.ErrNotExist):
72
		if err := pull(args[0], false); err != nil {
Michael Yang's avatar
Michael Yang committed
73
74
75
76
77
78
79
80
			var apiStatusError api.StatusError
			if !errors.As(err, &apiStatusError) {
				return err
			}

			if apiStatusError.StatusCode != http.StatusBadGateway {
				return err
			}
Michael Yang's avatar
Michael Yang committed
81
82
		}
	case err != nil:
Michael Yang's avatar
Michael Yang committed
83
84
85
86
		return err
	}

	return RunGenerate(cmd, args)
Bruce MacDonald's avatar
Bruce MacDonald committed
87
88
}

89
func PushHandler(cmd *cobra.Command, args []string) error {
90
91
	client := api.NewClient()

92
93
94
95
96
97
	insecure, err := cmd.Flags().GetBool("insecure")
	if err != nil {
		return err
	}

	request := api.PushRequest{Name: args[0], Insecure: insecure}
98
	fn := func(resp api.ProgressResponse) error {
99
100
101
102
103
104
105
106
107
108
		fmt.Println(resp.Status)
		return nil
	}

	if err := client.Push(context.Background(), &request, fn); err != nil {
		return err
	}
	return nil
}

109
func ListHandler(cmd *cobra.Command, args []string) error {
Patrick Devine's avatar
Patrick Devine committed
110
111
112
113
114
115
116
117
118
119
	client := api.NewClient()

	models, err := client.List(context.Background())
	if err != nil {
		return err
	}

	var data [][]string

	for _, m := range models.Models {
Michael Yang's avatar
Michael Yang committed
120
121
122
		if len(args) == 0 || strings.HasPrefix(m.Name, args[0]) {
			data = append(data, []string{m.Name, humanize.Bytes(uint64(m.Size)), format.HumanTime(m.ModifiedAt, "Never")})
		}
Patrick Devine's avatar
Patrick Devine committed
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
	}

	table := tablewriter.NewWriter(os.Stdout)
	table.SetHeader([]string{"NAME", "SIZE", "MODIFIED"})
	table.SetHeaderAlignment(tablewriter.ALIGN_LEFT)
	table.SetAlignment(tablewriter.ALIGN_LEFT)
	table.SetHeaderLine(false)
	table.SetBorder(false)
	table.SetNoWhiteSpace(true)
	table.SetTablePadding("\t")
	table.AppendBulk(data)
	table.Render()

	return nil
}

139
140
141
142
func DeleteHandler(cmd *cobra.Command, args []string) error {
	client := api.NewClient()

	request := api.DeleteRequest{Name: args[0]}
143
	if err := client.Delete(context.Background(), &request); err != nil {
144
145
		return err
	}
146
	fmt.Printf("deleted '%s'\n", args[0])
147
148
149
150
	return nil
}

func PullHandler(cmd *cobra.Command, args []string) error {
151
152
153
154
155
156
	insecure, err := cmd.Flags().GetBool("insecure")
	if err != nil {
		return err
	}

	return pull(args[0], insecure)
157
158
}

159
func pull(model string, insecure bool) error {
Bruce MacDonald's avatar
Bruce MacDonald committed
160
	client := api.NewClient()
161

162
	var currentDigest string
Bruce MacDonald's avatar
Bruce MacDonald committed
163
	var bar *progressbar.ProgressBar
Michael Yang's avatar
Michael Yang committed
164

165
	request := api.PullRequest{Name: model, Insecure: insecure}
166
167
168
	fn := func(resp api.ProgressResponse) error {
		if resp.Digest != currentDigest && resp.Digest != "" {
			currentDigest = resp.Digest
169
170
			bar = progressbar.DefaultBytes(
				int64(resp.Total),
171
				fmt.Sprintf("pulling %s...", resp.Digest[7:19]),
172
			)
173
174
175

			bar.Set(resp.Completed)
		} else if resp.Digest == currentDigest && resp.Digest != "" {
176
177
			bar.Set(resp.Completed)
		} else {
178
			currentDigest = ""
179
180
181
182
			fmt.Println(resp.Status)
		}
		return nil
	}
183

184
185
186
187
	if err := client.Pull(context.Background(), &request, fn); err != nil {
		return err
	}
	return nil
Michael Yang's avatar
Michael Yang committed
188
189
}

190
func RunGenerate(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
191
	if len(args) > 1 {
Michael Yang's avatar
Michael Yang committed
192
		// join all args into a single prompt
193
		return generate(cmd, args[0], strings.Join(args[1:], " "))
Michael Yang's avatar
Michael Yang committed
194
195
	}

Michael Yang's avatar
Michael Yang committed
196
	if readline.IsTerminal(int(os.Stdin.Fd())) {
197
		return generateInteractive(cmd, args[0])
Michael Yang's avatar
Michael Yang committed
198
199
	}

200
	return generateBatch(cmd, args[0])
Michael Yang's avatar
Michael Yang committed
201
202
}

Michael Yang's avatar
Michael Yang committed
203
204
var generateContextKey struct{}

205
func generate(cmd *cobra.Command, model, prompt string) error {
Michael Yang's avatar
Michael Yang committed
206
207
208
	if len(strings.TrimSpace(prompt)) > 0 {
		client := api.NewClient()

Michael Yang's avatar
Michael Yang committed
209
210
		spinner := NewSpinner("")
		go spinner.Spin(60 * time.Millisecond)
Michael Yang's avatar
Michael Yang committed
211

212
213
		var latest api.GenerateResponse

Michael Yang's avatar
Michael Yang committed
214
215
216
217
218
219
		generateContext, ok := cmd.Context().Value(generateContextKey).([]int)
		if !ok {
			generateContext = []int{}
		}

		request := api.GenerateRequest{Model: model, Prompt: prompt, Context: generateContext}
220
		fn := func(resp api.GenerateResponse) error {
Michael Yang's avatar
Michael Yang committed
221
222
223
			if !spinner.IsFinished() {
				spinner.Finish()
			}
Michael Yang's avatar
Michael Yang committed
224

225
226
			latest = resp

Michael Yang's avatar
Michael Yang committed
227
			fmt.Print(resp.Response)
Michael Yang's avatar
Michael Yang committed
228
229

			cmd.SetContext(context.WithValue(cmd.Context(), generateContextKey, resp.Context))
Michael Yang's avatar
Michael Yang committed
230
			return nil
231
232
233
234
235
		}

		if err := client.Generate(context.Background(), &request, fn); err != nil {
			return err
		}
Michael Yang's avatar
Michael Yang committed
236

Michael Yang's avatar
Michael Yang committed
237
238
		fmt.Println()
		fmt.Println()
239
240
241
242
243
244
245
246
247

		verbose, err := cmd.Flags().GetBool("verbose")
		if err != nil {
			return err
		}

		if verbose {
			latest.Summary()
		}
Michael Yang's avatar
Michael Yang committed
248
	}
Michael Yang's avatar
Michael Yang committed
249
250
251
252

	return nil
}

253
func generateInteractive(cmd *cobra.Command, model string) error {
Michael Yang's avatar
Michael Yang committed
254
255
256
257
258
259
260
261
262
263
264
	home, err := os.UserHomeDir()
	if err != nil {
		return err
	}

	completer := readline.NewPrefixCompleter(
		readline.PcItem("/help"),
		readline.PcItem("/list"),
		readline.PcItem("/set",
			readline.PcItem("history"),
			readline.PcItem("nohistory"),
Michael Yang's avatar
Michael Yang committed
265
266
			readline.PcItem("verbose"),
			readline.PcItem("quiet"),
Michael Yang's avatar
Michael Yang committed
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
			readline.PcItem("mode",
				readline.PcItem("vim"),
				readline.PcItem("emacs"),
				readline.PcItem("default"),
			),
		),
		readline.PcItem("/exit"),
		readline.PcItem("/bye"),
	)

	usage := func() {
		fmt.Fprintln(os.Stderr, "commands:")
		fmt.Fprintln(os.Stderr, completer.Tree("  "))
	}

	config := readline.Config{
		Prompt:       ">>> ",
		HistoryFile:  filepath.Join(home, ".ollama", "history"),
		AutoComplete: completer,
	}

	scanner, err := readline.NewEx(&config)
	if err != nil {
		return err
	}
	defer scanner.Close()

	for {
		line, err := scanner.Readline()
		switch {
		case errors.Is(err, io.EOF):
			return nil
		case errors.Is(err, readline.ErrInterrupt):
300
301
302
303
			if line == "" {
				return nil
			}

Michael Yang's avatar
Michael Yang committed
304
305
			continue
		case err != nil:
Michael Yang's avatar
Michael Yang committed
306
307
308
			return err
		}

Michael Yang's avatar
Michael Yang committed
309
		line = strings.TrimSpace(line)
Michael Yang's avatar
Michael Yang committed
310

Michael Yang's avatar
Michael Yang committed
311
312
313
		switch {
		case strings.HasPrefix(line, "/list"):
			args := strings.Fields(line)
314
			if err := ListHandler(cmd, args[1:]); err != nil {
Michael Yang's avatar
Michael Yang committed
315
316
317
318
319
320
321
322
323
324
325
326
327
328
				return err
			}

			continue
		case strings.HasPrefix(line, "/set"):
			args := strings.Fields(line)
			if len(args) > 1 {
				switch args[1] {
				case "history":
					scanner.HistoryEnable()
					continue
				case "nohistory":
					scanner.HistoryDisable()
					continue
Michael Yang's avatar
Michael Yang committed
329
330
331
332
333
334
				case "verbose":
					cmd.Flags().Set("verbose", "true")
					continue
				case "quiet":
					cmd.Flags().Set("verbose", "false")
					continue
Michael Yang's avatar
Michael Yang committed
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
				case "mode":
					if len(args) > 2 {
						switch args[2] {
						case "vim":
							scanner.SetVimMode(true)
							continue
						case "emacs", "default":
							scanner.SetVimMode(false)
							continue
						}
					}
				}
			}
		case line == "/help", line == "/?":
			usage()
			continue
		case line == "/exit", line == "/bye":
			return nil
		}

		if err := generate(cmd, model, line); err != nil {
			return err
		}
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
359
360
}

361
func generateBatch(cmd *cobra.Command, model string) error {
Michael Yang's avatar
Michael Yang committed
362
363
364
365
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		prompt := scanner.Text()
		fmt.Printf(">>> %s\n", prompt)
366
		if err := generate(cmd, model, prompt); err != nil {
Michael Yang's avatar
Michael Yang committed
367
368
369
370
371
372
373
374
			return err
		}
	}

	return nil
}

func RunServer(_ *cobra.Command, _ []string) error {
375
376
377
378
379
380
381
382
383
384
385
	host := os.Getenv("OLLAMA_HOST")
	if host == "" {
		host = "127.0.0.1"
	}

	port := os.Getenv("OLLAMA_PORT")
	if port == "" {
		port = "11434"
	}

	ln, err := net.Listen("tcp", fmt.Sprintf("%s:%s", host, port))
Jeffrey Morgan's avatar
Jeffrey Morgan committed
386
387
388
389
390
391
392
	if err != nil {
		return err
	}

	return server.Serve(ln)
}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
393
394
395
396
func NewCLI() *cobra.Command {
	log.SetFlags(log.LstdFlags | log.Lshortfile)

	rootCmd := &cobra.Command{
Michael Yang's avatar
Michael Yang committed
397
398
399
		Use:          "ollama",
		Short:        "Large language model runner",
		SilenceUsage: true,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
400
401
402
403
404
405
406
		CompletionOptions: cobra.CompletionOptions{
			DisableDefaultCmd: true,
		},
	}

	cobra.EnableCommandSorting = false

407
408
409
410
	createCmd := &cobra.Command{
		Use:   "create MODEL",
		Short: "Create a model from a Modelfile",
		Args:  cobra.MinimumNArgs(1),
411
		RunE:  CreateHandler,
412
413
414
415
	}

	createCmd.Flags().StringP("file", "f", "Modelfile", "Name of the Modelfile (default \"Modelfile\")")

Jeffrey Morgan's avatar
Jeffrey Morgan committed
416
	runCmd := &cobra.Command{
Michael Yang's avatar
Michael Yang committed
417
		Use:   "run MODEL [PROMPT]",
Jeffrey Morgan's avatar
Jeffrey Morgan committed
418
		Short: "Run a model",
Michael Yang's avatar
Michael Yang committed
419
		Args:  cobra.MinimumNArgs(1),
420
		RunE:  RunHandler,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
421
422
	}

423
424
	runCmd.Flags().Bool("verbose", false, "Show timings for response")

Jeffrey Morgan's avatar
Jeffrey Morgan committed
425
426
427
428
	serveCmd := &cobra.Command{
		Use:     "serve",
		Aliases: []string{"start"},
		Short:   "Start ollama",
Michael Yang's avatar
Michael Yang committed
429
		RunE:    RunServer,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
430
431
	}

432
433
434
435
	pullCmd := &cobra.Command{
		Use:   "pull MODEL",
		Short: "Pull a model from a registry",
		Args:  cobra.MinimumNArgs(1),
436
		RunE:  PullHandler,
437
438
	}

439
440
	pullCmd.Flags().Bool("insecure", false, "Use an insecure registry")

441
442
443
444
	pushCmd := &cobra.Command{
		Use:   "push MODEL",
		Short: "Push a model to a registry",
		Args:  cobra.MinimumNArgs(1),
445
		RunE:  PushHandler,
446
447
	}

448
449
	pushCmd.Flags().Bool("insecure", false, "Use an insecure registry")

Patrick Devine's avatar
Patrick Devine committed
450
	listCmd := &cobra.Command{
451
		Use:     "list",
Patrick Devine's avatar
Patrick Devine committed
452
		Aliases: []string{"ls"},
453
454
		Short:   "List models",
		RunE:    ListHandler,
455
456
457
458
459
460
461
	}

	deleteCmd := &cobra.Command{
		Use:   "rm",
		Short: "Remove a model",
		Args:  cobra.MinimumNArgs(1),
		RunE:  DeleteHandler,
Patrick Devine's avatar
Patrick Devine committed
462
463
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
464
465
	rootCmd.AddCommand(
		serveCmd,
466
		createCmd,
467
		runCmd,
468
469
		pullCmd,
		pushCmd,
Patrick Devine's avatar
Patrick Devine committed
470
		listCmd,
471
		deleteCmd,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
472
473
474
475
	)

	return rootCmd
}