cmd.go 12.5 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
	var currentDigest string
	var bar *progressbar.ProgressBar
Michael Yang's avatar
Michael Yang committed
41

42
43
44
45
46
47
48
49
50
51
52
	request := api.CreateRequest{Name: args[0], Path: filename}
	fn := func(resp api.ProgressResponse) error {
		if resp.Digest != currentDigest && resp.Digest != "" {
			if spinner != nil {
				spinner.Stop()
			}
			currentDigest = resp.Digest
			bar = progressbar.DefaultBytes(
				int64(resp.Total),
				fmt.Sprintf("pulling %s...", resp.Digest[7:19]),
			)
Michael Yang's avatar
Michael Yang committed
53

54
55
56
57
58
59
60
61
62
63
64
			bar.Set(resp.Completed)
		} else if resp.Digest == currentDigest && resp.Digest != "" {
			bar.Set(resp.Completed)
		} else {
			currentDigest = ""
			if spinner != nil {
				spinner.Stop()
			}
			spinner = NewSpinner(resp.Status)
			go spinner.Spin(100 * time.Millisecond)
		}
65
66
67
68
69
70
71
		return nil
	}

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

Michael Yang's avatar
Michael Yang committed
72
73
74
75
	if spinner != nil {
		spinner.Stop()
	}

76
77
78
	return nil
}

79
80
func RunHandler(cmd *cobra.Command, args []string) error {
	mp := server.ParseModelPath(args[0])
Patrick Devine's avatar
Patrick Devine committed
81
82
83
84
85
86
	fp, err := mp.GetManifestPath(false)
	if err != nil {
		return err
	}

	_, err = os.Stat(fp)
Michael Yang's avatar
Michael Yang committed
87
88
	switch {
	case errors.Is(err, os.ErrNotExist):
89
		if err := pull(args[0], false); err != nil {
Michael Yang's avatar
Michael Yang committed
90
91
92
93
94
95
96
97
			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
98
99
		}
	case err != nil:
Michael Yang's avatar
Michael Yang committed
100
101
102
103
		return err
	}

	return RunGenerate(cmd, args)
Bruce MacDonald's avatar
Bruce MacDonald committed
104
105
}

106
func PushHandler(cmd *cobra.Command, args []string) error {
107
108
	client := api.NewClient()

109
110
111
112
113
	insecure, err := cmd.Flags().GetBool("insecure")
	if err != nil {
		return err
	}

114
115
116
	var currentDigest string
	var bar *progressbar.ProgressBar

117
	request := api.PushRequest{Name: args[0], Insecure: insecure}
118
	fn := func(resp api.ProgressResponse) error {
119
120
121
122
123
124
125
126
127
128
129
130
131
132
		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)
		}
133
134
135
136
137
138
139
140
141
		return nil
	}

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

142
func ListHandler(cmd *cobra.Command, args []string) error {
Patrick Devine's avatar
Patrick Devine committed
143
144
145
146
147
148
149
150
151
152
	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
153
154
155
		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
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
	}

	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
}

172
173
174
func DeleteHandler(cmd *cobra.Command, args []string) error {
	client := api.NewClient()

Patrick Devine's avatar
Patrick Devine committed
175
176
	req := api.DeleteRequest{Name: args[0]}
	if err := client.Delete(context.Background(), &req); err != nil {
177
178
		return err
	}
179
	fmt.Printf("deleted '%s'\n", args[0])
180
181
182
	return nil
}

Patrick Devine's avatar
Patrick Devine committed
183
184
185
186
187
188
189
190
191
192
193
func CopyHandler(cmd *cobra.Command, args []string) error {
	client := api.NewClient()

	req := api.CopyRequest{Source: args[0], Destination: args[1]}
	if err := client.Copy(context.Background(), &req); err != nil {
		return err
	}
	fmt.Printf("copied '%s' to '%s'\n", args[0], args[1])
	return nil
}

194
func PullHandler(cmd *cobra.Command, args []string) error {
195
196
197
198
199
200
	insecure, err := cmd.Flags().GetBool("insecure")
	if err != nil {
		return err
	}

	return pull(args[0], insecure)
201
202
}

203
func pull(model string, insecure bool) error {
Bruce MacDonald's avatar
Bruce MacDonald committed
204
	client := api.NewClient()
205

206
	var currentDigest string
Bruce MacDonald's avatar
Bruce MacDonald committed
207
	var bar *progressbar.ProgressBar
Michael Yang's avatar
Michael Yang committed
208

209
	request := api.PullRequest{Name: model, Insecure: insecure}
210
211
212
	fn := func(resp api.ProgressResponse) error {
		if resp.Digest != currentDigest && resp.Digest != "" {
			currentDigest = resp.Digest
213
214
			bar = progressbar.DefaultBytes(
				int64(resp.Total),
215
				fmt.Sprintf("pulling %s...", resp.Digest[7:19]),
216
			)
217
218
219

			bar.Set(resp.Completed)
		} else if resp.Digest == currentDigest && resp.Digest != "" {
220
221
			bar.Set(resp.Completed)
		} else {
222
			currentDigest = ""
223
224
225
226
			fmt.Println(resp.Status)
		}
		return nil
	}
227

228
229
230
231
	if err := client.Pull(context.Background(), &request, fn); err != nil {
		return err
	}
	return nil
Michael Yang's avatar
Michael Yang committed
232
233
}

234
func RunGenerate(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
235
	if len(args) > 1 {
Michael Yang's avatar
Michael Yang committed
236
		// join all args into a single prompt
237
		return generate(cmd, args[0], strings.Join(args[1:], " "))
Michael Yang's avatar
Michael Yang committed
238
239
	}

Michael Yang's avatar
Michael Yang committed
240
	if readline.IsTerminal(int(os.Stdin.Fd())) {
241
		return generateInteractive(cmd, args[0])
Michael Yang's avatar
Michael Yang committed
242
243
	}

244
	return generateBatch(cmd, args[0])
Michael Yang's avatar
Michael Yang committed
245
246
}

Michael Yang's avatar
Michael Yang committed
247
type generateContextKey string
Michael Yang's avatar
Michael Yang committed
248

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

Michael Yang's avatar
Michael Yang committed
253
254
		spinner := NewSpinner("")
		go spinner.Spin(60 * time.Millisecond)
Michael Yang's avatar
Michael Yang committed
255

256
257
		var latest api.GenerateResponse

Michael Yang's avatar
Michael Yang committed
258
		generateContext, ok := cmd.Context().Value(generateContextKey("context")).([]int)
Michael Yang's avatar
Michael Yang committed
259
260
261
262
		if !ok {
			generateContext = []int{}
		}

Michael Yang's avatar
Michael Yang committed
263
264
265
266
267
268
269
		generateSession, ok := cmd.Context().Value(generateContextKey("session")).(int64)
		if !ok {
			generateSession = 0
		}

		request := api.GenerateRequest{Model: model, Prompt: prompt, Context: generateContext, SessionID: generateSession}
		fn := func(response api.GenerateResponse) error {
Michael Yang's avatar
Michael Yang committed
270
271
272
			if !spinner.IsFinished() {
				spinner.Finish()
			}
Michael Yang's avatar
Michael Yang committed
273

Michael Yang's avatar
Michael Yang committed
274
			latest = response
275

Michael Yang's avatar
Michael Yang committed
276
			fmt.Print(response.Response)
Michael Yang's avatar
Michael Yang committed
277
			return nil
278
279
280
281
282
		}

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

Michael Yang's avatar
Michael Yang committed
284
285
		fmt.Println()
		fmt.Println()
286
287
288
289
290
291
292
293
294

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

		if verbose {
			latest.Summary()
		}
Michael Yang's avatar
Michael Yang committed
295
296
297
298
299

		ctx := cmd.Context()
		ctx = context.WithValue(ctx, generateContextKey("context"), latest.Context)
		ctx = context.WithValue(ctx, generateContextKey("session"), latest.SessionID)
		cmd.SetContext(ctx)
Michael Yang's avatar
Michael Yang committed
300
	}
Michael Yang's avatar
Michael Yang committed
301
302
303
304

	return nil
}

305
306
307
308
309
310
311
312
313
314
func showLayer(l *server.Layer) {
	filename, err := server.GetBlobsPath(l.Digest)
	bts, err := os.ReadFile(filename)
	if err != nil {
		fmt.Printf("Couldn't read layer")
		return
	}
	fmt.Printf(string(bts) + "\n")
}

315
func generateInteractive(cmd *cobra.Command, model string) error {
Michael Yang's avatar
Michael Yang committed
316
317
318
319
320
321
322
323
324
325
326
	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
327
328
			readline.PcItem("verbose"),
			readline.PcItem("quiet"),
Michael Yang's avatar
Michael Yang committed
329
330
331
332
333
334
			readline.PcItem("mode",
				readline.PcItem("vim"),
				readline.PcItem("emacs"),
				readline.PcItem("default"),
			),
		),
335
336
337
338
339
		readline.PcItem("/show",
			readline.PcItem("license"),
			readline.PcItem("system"),
			readline.PcItem("template"),
		),
Michael Yang's avatar
Michael Yang committed
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
		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):
367
368
369
370
			if line == "" {
				return nil
			}

Michael Yang's avatar
Michael Yang committed
371
372
			continue
		case err != nil:
Michael Yang's avatar
Michael Yang committed
373
374
375
			return err
		}

Michael Yang's avatar
Michael Yang committed
376
		line = strings.TrimSpace(line)
Michael Yang's avatar
Michael Yang committed
377

Michael Yang's avatar
Michael Yang committed
378
379
380
		switch {
		case strings.HasPrefix(line, "/list"):
			args := strings.Fields(line)
381
			if err := ListHandler(cmd, args[1:]); err != nil {
Michael Yang's avatar
Michael Yang committed
382
383
384
385
386
387
388
389
390
391
392
393
394
395
				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
396
397
398
399
400
401
				case "verbose":
					cmd.Flags().Set("verbose", "true")
					continue
				case "quiet":
					cmd.Flags().Set("verbose", "false")
					continue
Michael Yang's avatar
Michael Yang committed
402
403
404
405
406
407
408
409
410
				case "mode":
					if len(args) > 2 {
						switch args[2] {
						case "vim":
							scanner.SetVimMode(true)
							continue
						case "emacs", "default":
							scanner.SetVimMode(false)
							continue
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
						default:
							usage()
							continue
						}

					} else {
						usage()
						continue
					}
				}
			} else {
				usage()
				continue
			}
		case strings.HasPrefix(line, "/show"):
			args := strings.Fields(line)
			if len(args) > 1 {
				mp := server.ParseModelPath(model)
				manifest, err := server.GetManifest(mp)
				if err != nil {
					fmt.Printf("error: couldn't get a manifestfor this model")
					continue
				}
				switch args[1] {
				case "license":
					for _, l := range manifest.Layers {
						if l.MediaType == "application/vnd.ollama.image.license" {
							showLayer(l)
						}
					}
					continue
				case "system":
					for _, l := range manifest.Layers {
						if l.MediaType == "application/vnd.ollama.image.system" {
							showLayer(l)
						}
					}
					continue
				case "template":
					for _, l := range manifest.Layers {
						if l.MediaType == "application/vnd.ollama.image.template" {
							showLayer(l)
Michael Yang's avatar
Michael Yang committed
453
454
						}
					}
455
456
457
458
					continue
				default:
					usage()
					continue
Michael Yang's avatar
Michael Yang committed
459
				}
460
461
462
			} else {
				usage()
				continue
Michael Yang's avatar
Michael Yang committed
463
464
465
466
467
468
469
470
471
472
473
474
			}
		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
475
476
}

477
func generateBatch(cmd *cobra.Command, model string) error {
Michael Yang's avatar
Michael Yang committed
478
479
480
481
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		prompt := scanner.Text()
		fmt.Printf(">>> %s\n", prompt)
482
		if err := generate(cmd, model, prompt); err != nil {
Michael Yang's avatar
Michael Yang committed
483
484
485
486
487
488
489
490
			return err
		}
	}

	return nil
}

func RunServer(_ *cobra.Command, _ []string) error {
491
492
493
494
495
496
497
498
499
500
501
	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
502
503
504
505
506
507
508
	if err != nil {
		return err
	}

	return server.Serve(ln)
}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
509
510
511
512
func NewCLI() *cobra.Command {
	log.SetFlags(log.LstdFlags | log.Lshortfile)

	rootCmd := &cobra.Command{
Michael Yang's avatar
Michael Yang committed
513
514
515
		Use:          "ollama",
		Short:        "Large language model runner",
		SilenceUsage: true,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
516
517
518
519
520
521
522
		CompletionOptions: cobra.CompletionOptions{
			DisableDefaultCmd: true,
		},
	}

	cobra.EnableCommandSorting = false

523
524
525
526
	createCmd := &cobra.Command{
		Use:   "create MODEL",
		Short: "Create a model from a Modelfile",
		Args:  cobra.MinimumNArgs(1),
527
		RunE:  CreateHandler,
528
529
530
531
	}

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

Jeffrey Morgan's avatar
Jeffrey Morgan committed
532
	runCmd := &cobra.Command{
Michael Yang's avatar
Michael Yang committed
533
		Use:   "run MODEL [PROMPT]",
Jeffrey Morgan's avatar
Jeffrey Morgan committed
534
		Short: "Run a model",
Michael Yang's avatar
Michael Yang committed
535
		Args:  cobra.MinimumNArgs(1),
536
		RunE:  RunHandler,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
537
538
	}

539
540
	runCmd.Flags().Bool("verbose", false, "Show timings for response")

Jeffrey Morgan's avatar
Jeffrey Morgan committed
541
542
543
544
	serveCmd := &cobra.Command{
		Use:     "serve",
		Aliases: []string{"start"},
		Short:   "Start ollama",
Michael Yang's avatar
Michael Yang committed
545
		RunE:    RunServer,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
546
547
	}

548
549
550
551
	pullCmd := &cobra.Command{
		Use:   "pull MODEL",
		Short: "Pull a model from a registry",
		Args:  cobra.MinimumNArgs(1),
552
		RunE:  PullHandler,
553
554
	}

555
556
	pullCmd.Flags().Bool("insecure", false, "Use an insecure registry")

557
558
559
560
	pushCmd := &cobra.Command{
		Use:   "push MODEL",
		Short: "Push a model to a registry",
		Args:  cobra.MinimumNArgs(1),
561
		RunE:  PushHandler,
562
563
	}

564
565
	pushCmd.Flags().Bool("insecure", false, "Use an insecure registry")

Patrick Devine's avatar
Patrick Devine committed
566
	listCmd := &cobra.Command{
567
		Use:     "list",
Patrick Devine's avatar
Patrick Devine committed
568
		Aliases: []string{"ls"},
569
570
		Short:   "List models",
		RunE:    ListHandler,
571
572
	}

Patrick Devine's avatar
Patrick Devine committed
573
574
575
576
577
578
579
	copyCmd := &cobra.Command{
		Use:   "cp",
		Short: "Copy a model",
		Args:  cobra.MinimumNArgs(2),
		RunE:  CopyHandler,
	}

580
581
582
583
584
	deleteCmd := &cobra.Command{
		Use:   "rm",
		Short: "Remove a model",
		Args:  cobra.MinimumNArgs(1),
		RunE:  DeleteHandler,
Patrick Devine's avatar
Patrick Devine committed
585
586
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
587
588
	rootCmd.AddCommand(
		serveCmd,
589
		createCmd,
590
		runCmd,
591
592
		pullCmd,
		pushCmd,
Patrick Devine's avatar
Patrick Devine committed
593
		listCmd,
Patrick Devine's avatar
Patrick Devine committed
594
		copyCmd,
595
		deleteCmd,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
596
597
598
599
	)

	return rootCmd
}