cmd.go 10 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
	insecure, err := cmd.Flags().GetBool("insecure")
	if err != nil {
		return err
	}

97
98
99
	var currentDigest string
	var bar *progressbar.ProgressBar

100
	request := api.PushRequest{Name: args[0], Insecure: insecure}
101
	fn := func(resp api.ProgressResponse) error {
102
103
104
105
106
107
108
109
110
111
112
113
114
115
		if resp.Digest != currentDigest && resp.Digest != "" {
			currentDigest = resp.Digest
			bar = progressbar.DefaultBytes(
				int64(resp.Total),
				fmt.Sprintf("pushing %s...", resp.Digest[7:19]),
			)

			bar.Set(resp.Completed)
		} else if resp.Digest == currentDigest && resp.Digest != "" {
			bar.Set(resp.Completed)
		} else {
			currentDigest = ""
			fmt.Println(resp.Status)
		}
116
117
118
119
120
121
122
123
124
		return nil
	}

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

125
func ListHandler(cmd *cobra.Command, args []string) error {
Patrick Devine's avatar
Patrick Devine committed
126
127
128
129
130
131
132
133
134
135
	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
136
137
138
		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
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
	}

	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
}

155
156
157
158
func DeleteHandler(cmd *cobra.Command, args []string) error {
	client := api.NewClient()

	request := api.DeleteRequest{Name: args[0]}
159
	if err := client.Delete(context.Background(), &request); err != nil {
160
161
		return err
	}
162
	fmt.Printf("deleted '%s'\n", args[0])
163
164
165
166
	return nil
}

func PullHandler(cmd *cobra.Command, args []string) error {
167
168
169
170
171
172
	insecure, err := cmd.Flags().GetBool("insecure")
	if err != nil {
		return err
	}

	return pull(args[0], insecure)
173
174
}

175
func pull(model string, insecure bool) error {
Bruce MacDonald's avatar
Bruce MacDonald committed
176
	client := api.NewClient()
177

178
	var currentDigest string
Bruce MacDonald's avatar
Bruce MacDonald committed
179
	var bar *progressbar.ProgressBar
Michael Yang's avatar
Michael Yang committed
180

181
	request := api.PullRequest{Name: model, Insecure: insecure}
182
183
184
	fn := func(resp api.ProgressResponse) error {
		if resp.Digest != currentDigest && resp.Digest != "" {
			currentDigest = resp.Digest
185
186
			bar = progressbar.DefaultBytes(
				int64(resp.Total),
187
				fmt.Sprintf("pulling %s...", resp.Digest[7:19]),
188
			)
189
190
191

			bar.Set(resp.Completed)
		} else if resp.Digest == currentDigest && resp.Digest != "" {
192
193
			bar.Set(resp.Completed)
		} else {
194
			currentDigest = ""
195
196
197
198
			fmt.Println(resp.Status)
		}
		return nil
	}
199

200
201
202
203
	if err := client.Pull(context.Background(), &request, fn); err != nil {
		return err
	}
	return nil
Michael Yang's avatar
Michael Yang committed
204
205
}

206
func RunGenerate(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
207
	if len(args) > 1 {
Michael Yang's avatar
Michael Yang committed
208
		// join all args into a single prompt
209
		return generate(cmd, args[0], strings.Join(args[1:], " "))
Michael Yang's avatar
Michael Yang committed
210
211
	}

Michael Yang's avatar
Michael Yang committed
212
	if readline.IsTerminal(int(os.Stdin.Fd())) {
213
		return generateInteractive(cmd, args[0])
Michael Yang's avatar
Michael Yang committed
214
215
	}

216
	return generateBatch(cmd, args[0])
Michael Yang's avatar
Michael Yang committed
217
218
}

Michael Yang's avatar
Michael Yang committed
219
220
var generateContextKey struct{}

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

Michael Yang's avatar
Michael Yang committed
225
226
		spinner := NewSpinner("")
		go spinner.Spin(60 * time.Millisecond)
Michael Yang's avatar
Michael Yang committed
227

228
229
		var latest api.GenerateResponse

Michael Yang's avatar
Michael Yang committed
230
231
232
233
234
235
		generateContext, ok := cmd.Context().Value(generateContextKey).([]int)
		if !ok {
			generateContext = []int{}
		}

		request := api.GenerateRequest{Model: model, Prompt: prompt, Context: generateContext}
236
		fn := func(resp api.GenerateResponse) error {
Michael Yang's avatar
Michael Yang committed
237
238
239
			if !spinner.IsFinished() {
				spinner.Finish()
			}
Michael Yang's avatar
Michael Yang committed
240

241
242
			latest = resp

Michael Yang's avatar
Michael Yang committed
243
			fmt.Print(resp.Response)
Michael Yang's avatar
Michael Yang committed
244
245

			cmd.SetContext(context.WithValue(cmd.Context(), generateContextKey, resp.Context))
Michael Yang's avatar
Michael Yang committed
246
			return nil
247
248
249
250
251
		}

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

Michael Yang's avatar
Michael Yang committed
253
254
		fmt.Println()
		fmt.Println()
255
256
257
258
259
260
261
262
263

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

		if verbose {
			latest.Summary()
		}
Michael Yang's avatar
Michael Yang committed
264
	}
Michael Yang's avatar
Michael Yang committed
265
266
267
268

	return nil
}

269
func generateInteractive(cmd *cobra.Command, model string) error {
Michael Yang's avatar
Michael Yang committed
270
271
272
273
274
275
276
277
278
279
280
	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
281
282
			readline.PcItem("verbose"),
			readline.PcItem("quiet"),
Michael Yang's avatar
Michael Yang committed
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
			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):
316
317
318
319
			if line == "" {
				return nil
			}

Michael Yang's avatar
Michael Yang committed
320
321
			continue
		case err != nil:
Michael Yang's avatar
Michael Yang committed
322
323
324
			return err
		}

Michael Yang's avatar
Michael Yang committed
325
		line = strings.TrimSpace(line)
Michael Yang's avatar
Michael Yang committed
326

Michael Yang's avatar
Michael Yang committed
327
328
329
		switch {
		case strings.HasPrefix(line, "/list"):
			args := strings.Fields(line)
330
			if err := ListHandler(cmd, args[1:]); err != nil {
Michael Yang's avatar
Michael Yang committed
331
332
333
334
335
336
337
338
339
340
341
342
343
344
				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
345
346
347
348
349
350
				case "verbose":
					cmd.Flags().Set("verbose", "true")
					continue
				case "quiet":
					cmd.Flags().Set("verbose", "false")
					continue
Michael Yang's avatar
Michael Yang committed
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
				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
375
376
}

377
func generateBatch(cmd *cobra.Command, model string) error {
Michael Yang's avatar
Michael Yang committed
378
379
380
381
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		prompt := scanner.Text()
		fmt.Printf(">>> %s\n", prompt)
382
		if err := generate(cmd, model, prompt); err != nil {
Michael Yang's avatar
Michael Yang committed
383
384
385
386
387
388
389
390
			return err
		}
	}

	return nil
}

func RunServer(_ *cobra.Command, _ []string) error {
391
392
393
394
395
396
397
398
399
400
401
	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
402
403
404
405
406
407
408
	if err != nil {
		return err
	}

	return server.Serve(ln)
}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
409
410
411
412
func NewCLI() *cobra.Command {
	log.SetFlags(log.LstdFlags | log.Lshortfile)

	rootCmd := &cobra.Command{
Michael Yang's avatar
Michael Yang committed
413
414
415
		Use:          "ollama",
		Short:        "Large language model runner",
		SilenceUsage: true,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
416
417
418
419
420
421
422
		CompletionOptions: cobra.CompletionOptions{
			DisableDefaultCmd: true,
		},
	}

	cobra.EnableCommandSorting = false

423
424
425
426
	createCmd := &cobra.Command{
		Use:   "create MODEL",
		Short: "Create a model from a Modelfile",
		Args:  cobra.MinimumNArgs(1),
427
		RunE:  CreateHandler,
428
429
430
431
	}

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

Jeffrey Morgan's avatar
Jeffrey Morgan committed
432
	runCmd := &cobra.Command{
Michael Yang's avatar
Michael Yang committed
433
		Use:   "run MODEL [PROMPT]",
Jeffrey Morgan's avatar
Jeffrey Morgan committed
434
		Short: "Run a model",
Michael Yang's avatar
Michael Yang committed
435
		Args:  cobra.MinimumNArgs(1),
436
		RunE:  RunHandler,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
437
438
	}

439
440
	runCmd.Flags().Bool("verbose", false, "Show timings for response")

Jeffrey Morgan's avatar
Jeffrey Morgan committed
441
442
443
444
	serveCmd := &cobra.Command{
		Use:     "serve",
		Aliases: []string{"start"},
		Short:   "Start ollama",
Michael Yang's avatar
Michael Yang committed
445
		RunE:    RunServer,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
446
447
	}

448
449
450
451
	pullCmd := &cobra.Command{
		Use:   "pull MODEL",
		Short: "Pull a model from a registry",
		Args:  cobra.MinimumNArgs(1),
452
		RunE:  PullHandler,
453
454
	}

455
456
	pullCmd.Flags().Bool("insecure", false, "Use an insecure registry")

457
458
459
460
	pushCmd := &cobra.Command{
		Use:   "push MODEL",
		Short: "Push a model to a registry",
		Args:  cobra.MinimumNArgs(1),
461
		RunE:  PushHandler,
462
463
	}

464
465
	pushCmd.Flags().Bool("insecure", false, "Use an insecure registry")

Patrick Devine's avatar
Patrick Devine committed
466
	listCmd := &cobra.Command{
467
		Use:     "list",
Patrick Devine's avatar
Patrick Devine committed
468
		Aliases: []string{"ls"},
469
470
		Short:   "List models",
		RunE:    ListHandler,
471
472
473
474
475
476
477
	}

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

Jeffrey Morgan's avatar
Jeffrey Morgan committed
480
481
	rootCmd.AddCommand(
		serveCmd,
482
		createCmd,
483
		runCmd,
484
485
		pullCmd,
		pushCmd,
Patrick Devine's avatar
Patrick Devine committed
486
		listCmd,
487
		deleteCmd,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
488
489
490
491
	)

	return rootCmd
}