cmd.go 22.4 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"
6
7
8
	"crypto/ed25519"
	"crypto/rand"
	"encoding/pem"
Michael Yang's avatar
Michael Yang committed
9
	"errors"
Bruce MacDonald's avatar
Bruce MacDonald committed
10
	"fmt"
Michael Yang's avatar
Michael Yang committed
11
	"io"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
12
13
14
	"log"
	"net"
	"os"
15
	"os/exec"
16
	"os/signal"
17
	"path/filepath"
18
	"runtime"
Michael Yang's avatar
Michael Yang committed
19
	"strings"
20
	"syscall"
Michael Yang's avatar
Michael Yang committed
21
	"time"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
22

Patrick Devine's avatar
Patrick Devine committed
23
24
	"github.com/dustin/go-humanize"
	"github.com/olekukonko/tablewriter"
25
	"github.com/pdevine/readline"
Michael Yang's avatar
Michael Yang committed
26
	"github.com/spf13/cobra"
27
	"golang.org/x/crypto/ssh"
28
	"golang.org/x/term"
Michael Yang's avatar
Michael Yang committed
29

Jeffrey Morgan's avatar
Jeffrey Morgan committed
30
	"github.com/jmorganca/ollama/api"
Patrick Devine's avatar
Patrick Devine committed
31
	"github.com/jmorganca/ollama/format"
32
	"github.com/jmorganca/ollama/progressbar"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
33
	"github.com/jmorganca/ollama/server"
Michael Yang's avatar
Michael Yang committed
34
	"github.com/jmorganca/ollama/version"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
35
36
)

37
type Painter struct {
38
	IsMultiLine bool
39
}
Patrick Devine's avatar
Patrick Devine committed
40

41
func (p Painter) Paint(line []rune, _ int) []rune {
Patrick Devine's avatar
Patrick Devine committed
42
	termType := os.Getenv("TERM")
43
44
45
46
47
	if termType == "xterm-256color" && len(line) == 0 {
		var prompt string
		if p.IsMultiLine {
			prompt = "Use \"\"\" to end multi-line input"
		} else {
48
			prompt = "Send a message (/? for help)"
49
		}
Patrick Devine's avatar
Patrick Devine committed
50
51
		return []rune(fmt.Sprintf("\033[38;5;245m%s\033[%dD\033[0m", prompt, len(prompt)))
	}
52
53
	// add a space and a backspace to prevent the cursor from walking up the screen
	line = append(line, []rune(" \b")...)
Patrick Devine's avatar
Patrick Devine committed
54
55
56
	return line
}

57
func CreateHandler(cmd *cobra.Command, args []string) error {
58
	filename, _ := cmd.Flags().GetString("file")
59
60
61
62
63
	filename, err := filepath.Abs(filename)
	if err != nil {
		return err
	}

64
65
66
67
	client, err := api.FromEnv()
	if err != nil {
		return err
	}
68

Michael Yang's avatar
Michael Yang committed
69
70
	var spinner *Spinner

71
72
	var currentDigest string
	var bar *progressbar.ProgressBar
Michael Yang's avatar
Michael Yang committed
73

74
75
76
77
78
79
80
	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
81
82
			switch {
			case strings.Contains(resp.Status, "embeddings"):
Michael Yang's avatar
Michael Yang committed
83
84
				bar = progressbar.Default(resp.Total, resp.Status)
				bar.Set64(resp.Completed)
85
86
87
			default:
				// pulling
				bar = progressbar.DefaultBytes(
Michael Yang's avatar
Michael Yang committed
88
					resp.Total,
89
90
					resp.Status,
				)
Michael Yang's avatar
Michael Yang committed
91
				bar.Set64(resp.Completed)
92
			}
93
		} else if resp.Digest == currentDigest && resp.Digest != "" {
Michael Yang's avatar
Michael Yang committed
94
			bar.Set64(resp.Completed)
95
96
97
98
99
100
101
102
		} else {
			currentDigest = ""
			if spinner != nil {
				spinner.Stop()
			}
			spinner = NewSpinner(resp.Status)
			go spinner.Spin(100 * time.Millisecond)
		}
103

104
105
106
107
108
109
110
		return nil
	}

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

Michael Yang's avatar
Michael Yang committed
111
112
	if spinner != nil {
		spinner.Stop()
113
114
115
		if spinner.description != "success" {
			return errors.New("unexpected end to create model")
		}
Michael Yang's avatar
Michael Yang committed
116
117
	}

118
119
120
	return nil
}

121
func RunHandler(cmd *cobra.Command, args []string) error {
122
	client, err := api.FromEnv()
123
124
125
126
	if err != nil {
		return err
	}

127
	models, err := client.List(context.Background())
Patrick Devine's avatar
Patrick Devine committed
128
129
130
131
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
132
	canonicalModelPath := server.ParseModelPath(args[0])
133
	for _, model := range models.Models {
Michael Yang's avatar
Michael Yang committed
134
		if model.Name == canonicalModelPath.GetShortTagname() {
135
			return RunGenerate(cmd, args)
Michael Yang's avatar
Michael Yang committed
136
		}
137
138
139
	}

	if err := PullHandler(cmd, args); err != nil {
Michael Yang's avatar
Michael Yang committed
140
141
142
143
		return err
	}

	return RunGenerate(cmd, args)
Bruce MacDonald's avatar
Bruce MacDonald committed
144
145
}

146
func PushHandler(cmd *cobra.Command, args []string) error {
147
148
149
150
	client, err := api.FromEnv()
	if err != nil {
		return err
	}
151

152
153
154
155
156
	insecure, err := cmd.Flags().GetBool("insecure")
	if err != nil {
		return err
	}

157
158
159
	var currentDigest string
	var bar *progressbar.ProgressBar

160
	request := api.PushRequest{Name: args[0], Insecure: insecure}
161
	fn := func(resp api.ProgressResponse) error {
162
163
164
		if resp.Digest != currentDigest && resp.Digest != "" {
			currentDigest = resp.Digest
			bar = progressbar.DefaultBytes(
Michael Yang's avatar
Michael Yang committed
165
				resp.Total,
166
167
168
				fmt.Sprintf("pushing %s...", resp.Digest[7:19]),
			)

Michael Yang's avatar
Michael Yang committed
169
			bar.Set64(resp.Completed)
170
		} else if resp.Digest == currentDigest && resp.Digest != "" {
Michael Yang's avatar
Michael Yang committed
171
			bar.Set64(resp.Completed)
172
173
174
175
		} else {
			currentDigest = ""
			fmt.Println(resp.Status)
		}
176
177
178
179
180
181
		return nil
	}

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

	if bar != nil && !bar.IsFinished() {
		return errors.New("unexpected end to push model")
	}

187
188
189
	return nil
}

190
func ListHandler(cmd *cobra.Command, args []string) error {
191
192
193
194
	client, err := api.FromEnv()
	if err != nil {
		return err
	}
Patrick Devine's avatar
Patrick Devine committed
195
196
197
198
199
200
201
202
203

	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
204
		if len(args) == 0 || strings.HasPrefix(m.Name, args[0]) {
Patrick Devine's avatar
Patrick Devine committed
205
			data = append(data, []string{m.Name, m.Digest[:12], humanize.Bytes(uint64(m.Size)), format.HumanTime(m.ModifiedAt, "Never")})
Michael Yang's avatar
Michael Yang committed
206
		}
Patrick Devine's avatar
Patrick Devine committed
207
208
209
	}

	table := tablewriter.NewWriter(os.Stdout)
Patrick Devine's avatar
Patrick Devine committed
210
	table.SetHeader([]string{"NAME", "ID", "SIZE", "MODIFIED"})
Patrick Devine's avatar
Patrick Devine committed
211
212
213
214
215
216
217
218
219
220
221
222
	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
}

223
func DeleteHandler(cmd *cobra.Command, args []string) error {
224
225
226
227
	client, err := api.FromEnv()
	if err != nil {
		return err
	}
228

229
230
231
232
233
234
	for _, name := range args {
		req := api.DeleteRequest{Name: name}
		if err := client.Delete(context.Background(), &req); err != nil {
			return err
		}
		fmt.Printf("deleted '%s'\n", name)
235
236
237
238
	}
	return nil
}

Patrick Devine's avatar
Patrick Devine committed
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
func ShowHandler(cmd *cobra.Command, args []string) error {
	client, err := api.FromEnv()
	if err != nil {
		return err
	}

	if len(args) != 1 {
		return errors.New("missing model name")
	}

	license, errLicense := cmd.Flags().GetBool("license")
	modelfile, errModelfile := cmd.Flags().GetBool("modelfile")
	parameters, errParams := cmd.Flags().GetBool("parameters")
	system, errSystem := cmd.Flags().GetBool("system")
	template, errTemplate := cmd.Flags().GetBool("template")

	for _, boolErr := range []error{errLicense, errModelfile, errParams, errSystem, errTemplate} {
		if boolErr != nil {
			return errors.New("error retrieving flags")
		}
	}

	flagsSet := 0
	showType := ""

	if license {
		flagsSet++
		showType = "license"
	}

	if modelfile {
		flagsSet++
		showType = "modelfile"
	}

	if parameters {
		flagsSet++
		showType = "parameters"
	}

	if system {
		flagsSet++
		showType = "system"
	}

	if template {
		flagsSet++
		showType = "template"
	}

	if flagsSet > 1 {
290
		return errors.New("only one of '--license', '--modelfile', '--parameters', '--system', or '--template' can be specified")
Patrick Devine's avatar
Patrick Devine committed
291
	} else if flagsSet == 0 {
292
		return errors.New("one of '--license', '--modelfile', '--parameters', '--system', or '--template' must be specified")
Patrick Devine's avatar
Patrick Devine committed
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
	}

	req := api.ShowRequest{Name: args[0]}
	resp, err := client.Show(context.Background(), &req)
	if err != nil {
		return err
	}

	switch showType {
	case "license":
		fmt.Println(resp.License)
	case "modelfile":
		fmt.Println(resp.Modelfile)
	case "parameters":
		fmt.Println(resp.Parameters)
	case "system":
		fmt.Println(resp.System)
	case "template":
		fmt.Println(resp.Template)
	}

	return nil
}

Patrick Devine's avatar
Patrick Devine committed
317
func CopyHandler(cmd *cobra.Command, args []string) error {
318
319
320
321
	client, err := api.FromEnv()
	if err != nil {
		return err
	}
Patrick Devine's avatar
Patrick Devine committed
322
323
324
325
326
327
328
329
330

	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
}

331
func PullHandler(cmd *cobra.Command, args []string) error {
332
333
334
335
336
337
	insecure, err := cmd.Flags().GetBool("insecure")
	if err != nil {
		return err
	}

	return pull(args[0], insecure)
338
339
}

340
func pull(model string, insecure bool) error {
341
342
343
344
	client, err := api.FromEnv()
	if err != nil {
		return err
	}
345

346
	var currentDigest string
Bruce MacDonald's avatar
Bruce MacDonald committed
347
	var bar *progressbar.ProgressBar
Michael Yang's avatar
Michael Yang committed
348

349
	request := api.PullRequest{Name: model, Insecure: insecure}
350
351
352
	fn := func(resp api.ProgressResponse) error {
		if resp.Digest != currentDigest && resp.Digest != "" {
			currentDigest = resp.Digest
353
			bar = progressbar.DefaultBytes(
Michael Yang's avatar
Michael Yang committed
354
				resp.Total,
355
				fmt.Sprintf("pulling %s...", resp.Digest[7:19]),
356
			)
357

Michael Yang's avatar
Michael Yang committed
358
			bar.Set64(resp.Completed)
359
		} else if resp.Digest == currentDigest && resp.Digest != "" {
Michael Yang's avatar
Michael Yang committed
360
			bar.Set64(resp.Completed)
361
		} else {
362
			currentDigest = ""
363
364
			fmt.Println(resp.Status)
		}
365

366
367
		return nil
	}
368

369
370
371
	if err := client.Pull(context.Background(), &request, fn); err != nil {
		return err
	}
372
373
374
375
376

	if bar != nil && !bar.IsFinished() {
		return errors.New("unexpected end to pull model")
	}

377
	return nil
Michael Yang's avatar
Michael Yang committed
378
379
}

380
func RunGenerate(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
381
	if len(args) > 1 {
Michael Yang's avatar
Michael Yang committed
382
		// join all args into a single prompt
383
384
385
386
387
388
389
390
391
392
393
394
395
396
		wordWrap := false
		if term.IsTerminal(int(os.Stdout.Fd())) {
			wordWrap = true
		}

		nowrap, err := cmd.Flags().GetBool("nowordwrap")
		if err != nil {
			return err
		}
		if nowrap {
			wordWrap = false
		}

		return generate(cmd, args[0], strings.Join(args[1:], " "), wordWrap)
Michael Yang's avatar
Michael Yang committed
397
398
	}

Michael Yang's avatar
Michael Yang committed
399
	if readline.IsTerminal(int(os.Stdin.Fd())) {
400
		return generateInteractive(cmd, args[0])
Michael Yang's avatar
Michael Yang committed
401
402
	}

403
	return generateBatch(cmd, args[0])
Michael Yang's avatar
Michael Yang committed
404
405
}

Michael Yang's avatar
Michael Yang committed
406
type generateContextKey string
Michael Yang's avatar
Michael Yang committed
407

408
func generate(cmd *cobra.Command, model, prompt string, wordWrap bool) error {
Patrick Devine's avatar
Patrick Devine committed
409
410
411
412
	client, err := api.FromEnv()
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
413

Patrick Devine's avatar
Patrick Devine committed
414
415
	spinner := NewSpinner("")
	go spinner.Spin(60 * time.Millisecond)
Michael Yang's avatar
Michael Yang committed
416

Patrick Devine's avatar
Patrick Devine committed
417
	var latest api.GenerateResponse
418

Patrick Devine's avatar
Patrick Devine committed
419
420
421
422
	generateContext, ok := cmd.Context().Value(generateContextKey("context")).([]int)
	if !ok {
		generateContext = []int{}
	}
Michael Yang's avatar
Michael Yang committed
423

424
425
	termWidth, _, err := term.GetSize(int(0))
	if err != nil {
426
		wordWrap = false
427
428
	}

429
430
431
432
433
434
435
436
437
438
439
440
441
	cancelCtx, cancel := context.WithCancel(context.Background())
	defer cancel()

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT)
	var abort bool

	go func() {
		<-sigChan
		cancel()
		abort = true
	}()

442
443
444
	var currentLineLength int
	var wordBuffer string

Patrick Devine's avatar
Patrick Devine committed
445
446
447
448
449
	request := api.GenerateRequest{Model: model, Prompt: prompt, Context: generateContext}
	fn := func(response api.GenerateResponse) error {
		if !spinner.IsFinished() {
			spinner.Finish()
		}
Michael Yang's avatar
Michael Yang committed
450

Patrick Devine's avatar
Patrick Devine committed
451
		latest = response
452

453
		if wordWrap {
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
			for _, ch := range response.Response {
				if currentLineLength+1 > termWidth-5 {
					// backtrack the length of the last word and clear to the end of the line
					fmt.Printf("\x1b[%dD\x1b[K\n", len(wordBuffer))
					fmt.Printf("%s%c", wordBuffer, ch)
					currentLineLength = len(wordBuffer) + 1
				} else {
					fmt.Print(string(ch))
					currentLineLength += 1

					switch ch {
					case ' ':
						wordBuffer = ""
					case '\n':
						currentLineLength = 0
					default:
						wordBuffer += string(ch)
					}
				}
			}
		} else {
			fmt.Print(response.Response)
		}

Patrick Devine's avatar
Patrick Devine committed
478
479
		return nil
	}
480

481
	if err := client.Generate(cancelCtx, &request, fn); err != nil {
Patrick Devine's avatar
Patrick Devine committed
482
483
484
485
486
487
488
489
490
491
		if strings.Contains(err.Error(), "failed to load model") {
			// tell the user to check the server log, if it exists locally
			home, nestedErr := os.UserHomeDir()
			if nestedErr != nil {
				// return the original error
				return err
			}
			logPath := filepath.Join(home, ".ollama", "logs", "server.log")
			if _, nestedErr := os.Stat(logPath); nestedErr == nil {
				err = fmt.Errorf("%w\nFor more details, check the error logs at %s", err, logPath)
492
			}
493
494
495
		} else if strings.Contains(err.Error(), "context canceled") && abort {
			spinner.Finish()
			return nil
496
		}
Patrick Devine's avatar
Patrick Devine committed
497
498
499
		return err
	}
	if prompt != "" {
Michael Yang's avatar
Michael Yang committed
500
501
		fmt.Println()
		fmt.Println()
Patrick Devine's avatar
Patrick Devine committed
502
	}
503

Patrick Devine's avatar
Patrick Devine committed
504
	if !latest.Done {
505
506
507
		if abort {
			return nil
		}
Patrick Devine's avatar
Patrick Devine committed
508
509
		return errors.New("unexpected end of response")
	}
510

Patrick Devine's avatar
Patrick Devine committed
511
512
513
514
	verbose, err := cmd.Flags().GetBool("verbose")
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
515

Patrick Devine's avatar
Patrick Devine committed
516
517
	if verbose {
		latest.Summary()
Michael Yang's avatar
Michael Yang committed
518
	}
Michael Yang's avatar
Michael Yang committed
519

Patrick Devine's avatar
Patrick Devine committed
520
521
522
523
	ctx := cmd.Context()
	ctx = context.WithValue(ctx, generateContextKey("context"), latest.Context)
	cmd.SetContext(ctx)

Michael Yang's avatar
Michael Yang committed
524
525
526
	return nil
}

527
func generateInteractive(cmd *cobra.Command, model string) error {
Michael Yang's avatar
Michael Yang committed
528
529
530
531
532
	home, err := os.UserHomeDir()
	if err != nil {
		return err
	}

Patrick Devine's avatar
Patrick Devine committed
533
	// load the model
534
	if err := generate(cmd, model, "", false); err != nil {
Patrick Devine's avatar
Patrick Devine committed
535
536
537
		return err
	}

Michael Yang's avatar
Michael Yang committed
538
539
540
541
542
543
	completer := readline.NewPrefixCompleter(
		readline.PcItem("/help"),
		readline.PcItem("/list"),
		readline.PcItem("/set",
			readline.PcItem("history"),
			readline.PcItem("nohistory"),
544
545
			readline.PcItem("wordwrap"),
			readline.PcItem("nowordwrap"),
Michael Yang's avatar
Michael Yang committed
546
547
			readline.PcItem("verbose"),
			readline.PcItem("quiet"),
Michael Yang's avatar
Michael Yang committed
548
		),
549
550
		readline.PcItem("/show",
			readline.PcItem("license"),
Patrick Devine's avatar
Patrick Devine committed
551
552
			readline.PcItem("modelfile"),
			readline.PcItem("parameters"),
553
554
555
			readline.PcItem("system"),
			readline.PcItem("template"),
		),
Michael Yang's avatar
Michael Yang committed
556
557
558
559
560
561
562
563
564
		readline.PcItem("/exit"),
		readline.PcItem("/bye"),
	)

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

565
566
	var painter Painter

Michael Yang's avatar
Michael Yang committed
567
	config := readline.Config{
568
		Painter:      &painter,
Michael Yang's avatar
Michael Yang committed
569
570
571
572
573
574
575
576
577
578
579
		Prompt:       ">>> ",
		HistoryFile:  filepath.Join(home, ".ollama", "history"),
		AutoComplete: completer,
	}

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

580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
	var wordWrap bool
	termType := os.Getenv("TERM")
	if termType == "xterm-256color" {
		wordWrap = true
	}

	// override wrapping if the user turned it off
	nowrap, err := cmd.Flags().GetBool("nowordwrap")
	if err != nil {
		return err
	}
	if nowrap {
		wordWrap = false
	}

595
596
597
	var multiLineBuffer string
	var isMultiLine bool

Michael Yang's avatar
Michael Yang committed
598
599
600
601
602
603
	for {
		line, err := scanner.Readline()
		switch {
		case errors.Is(err, io.EOF):
			return nil
		case errors.Is(err, readline.ErrInterrupt):
604
			if line == "" {
605
				fmt.Println("Use Ctrl-D or /bye to exit.")
606
607
			}

Michael Yang's avatar
Michael Yang committed
608
609
			continue
		case err != nil:
Michael Yang's avatar
Michael Yang committed
610
611
612
			return err
		}

Michael Yang's avatar
Michael Yang committed
613
		line = strings.TrimSpace(line)
Michael Yang's avatar
Michael Yang committed
614

Michael Yang's avatar
Michael Yang committed
615
		switch {
616
617
618
		case isMultiLine:
			if strings.HasSuffix(line, `"""`) {
				isMultiLine = false
619
				painter.IsMultiLine = isMultiLine
620
				multiLineBuffer += strings.TrimSuffix(line, `"""`)
621
				line = multiLineBuffer
622
623
624
625
626
627
628
629
				multiLineBuffer = ""
				scanner.SetPrompt(">>> ")
			} else {
				multiLineBuffer += line + " "
				continue
			}
		case strings.HasPrefix(line, `"""`):
			isMultiLine = true
630
			painter.IsMultiLine = isMultiLine
631
632
633
			multiLineBuffer = strings.TrimPrefix(line, `"""`) + " "
			scanner.SetPrompt("... ")
			continue
Michael Yang's avatar
Michael Yang committed
634
635
		case strings.HasPrefix(line, "/list"):
			args := strings.Fields(line)
636
			if err := ListHandler(cmd, args[1:]); err != nil {
Michael Yang's avatar
Michael Yang committed
637
638
639
640
641
642
643
644
645
646
				return err
			}
		case strings.HasPrefix(line, "/set"):
			args := strings.Fields(line)
			if len(args) > 1 {
				switch args[1] {
				case "history":
					scanner.HistoryEnable()
				case "nohistory":
					scanner.HistoryDisable()
647
				case "wordwrap":
648
					wordWrap = true
649
650
					fmt.Println("Set 'wordwrap' mode.")
				case "nowordwrap":
651
					wordWrap = false
652
					fmt.Println("Set 'nowordwrap' mode.")
Michael Yang's avatar
Michael Yang committed
653
654
				case "verbose":
					cmd.Flags().Set("verbose", "true")
655
					fmt.Println("Set 'verbose' mode.")
Michael Yang's avatar
Michael Yang committed
656
657
				case "quiet":
					cmd.Flags().Set("verbose", "false")
658
					fmt.Println("Set 'quiet' mode.")
Michael Yang's avatar
Michael Yang committed
659
660
661
662
663
664
665
				case "mode":
					if len(args) > 2 {
						switch args[2] {
						case "vim":
							scanner.SetVimMode(true)
						case "emacs", "default":
							scanner.SetVimMode(false)
666
667
668
669
670
671
						default:
							usage()
						}
					} else {
						usage()
					}
672
673
				default:
					fmt.Printf("Unknown command '/set %s'. Type /? for help\n", args[1])
674
675
676
677
678
679
680
				}
			} else {
				usage()
			}
		case strings.HasPrefix(line, "/show"):
			args := strings.Fields(line)
			if len(args) > 1 {
Patrick Devine's avatar
Patrick Devine committed
681
				resp, err := server.GetModelInfo(model)
682
				if err != nil {
Patrick Devine's avatar
Patrick Devine committed
683
					fmt.Println("error: couldn't get model")
684
					return err
685
				}
Patrick Devine's avatar
Patrick Devine committed
686

687
688
				switch args[1] {
				case "license":
689
690
691
692
693
					if resp.License == "" {
						fmt.Println("No license was specified for this model.\n")
					} else {
						fmt.Println(resp.License)
					}
Patrick Devine's avatar
Patrick Devine committed
694
695
696
				case "modelfile":
					fmt.Println(resp.Modelfile)
				case "parameters":
697
698
699
700
701
					if resp.Parameters == "" {
						fmt.Println("No parameters were specified for this model.\n")
					} else {
						fmt.Println(resp.Parameters)
					}
702
				case "system":
703
704
705
706
707
					if resp.System == "" {
						fmt.Println("No system prompt was specified for this model.\n")
					} else {
						fmt.Println(resp.System)
					}
708
				case "template":
709
710
711
712
713
					if resp.Template == "" {
						fmt.Println("No prompt template was specified for this model.\n")
					} else {
						fmt.Println(resp.Template)
					}
714
				default:
715
					fmt.Printf("Unknown command '/show %s'. Type /? for help\n", args[1])
Michael Yang's avatar
Michael Yang committed
716
				}
717
718
			} else {
				usage()
Michael Yang's avatar
Michael Yang committed
719
720
721
722
723
			}
		case line == "/help", line == "/?":
			usage()
		case line == "/exit", line == "/bye":
			return nil
724
725
726
		case strings.HasPrefix(line, "/"):
			args := strings.Fields(line)
			fmt.Printf("Unknown command '%s'. Type /? for help\n", args[0])
Michael Yang's avatar
Michael Yang committed
727
728
		}

Patrick Devine's avatar
Patrick Devine committed
729
		if len(line) > 0 && line[0] != '/' {
730
			if err := generate(cmd, model, line, wordWrap); err != nil {
Patrick Devine's avatar
Patrick Devine committed
731
732
				return err
			}
Michael Yang's avatar
Michael Yang committed
733
734
		}
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
735
736
}

737
func generateBatch(cmd *cobra.Command, model string) error {
Michael Yang's avatar
Michael Yang committed
738
739
740
741
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		prompt := scanner.Text()
		fmt.Printf(">>> %s\n", prompt)
742
		if err := generate(cmd, model, prompt, false); err != nil {
Michael Yang's avatar
Michael Yang committed
743
744
745
746
747
748
749
			return err
		}
	}

	return nil
}

750
func RunServer(cmd *cobra.Command, _ []string) error {
Michael Yang's avatar
Michael Yang committed
751
752
753
	host, port, err := net.SplitHostPort(os.Getenv("OLLAMA_HOST"))
	if err != nil {
		host, port = "127.0.0.1", "11434"
Michael Yang's avatar
Michael Yang committed
754
		if ip := net.ParseIP(strings.Trim(os.Getenv("OLLAMA_HOST"), "[]")); ip != nil {
Michael Yang's avatar
Michael Yang committed
755
756
			host = ip.String()
		}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
757
	}
758

Michael Yang's avatar
Michael Yang committed
759
	if err := initializeKeypair(); err != nil {
760
761
762
		return err
	}

Michael Yang's avatar
Michael Yang committed
763
	ln, err := net.Listen("tcp", net.JoinHostPort(host, port))
764
765
766
	if err != nil {
		return err
	}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
767

768
769
770
771
772
	var origins []string
	if o := os.Getenv("OLLAMA_ORIGINS"); o != "" {
		origins = strings.Split(o, ",")
	}

773
774
775
776
	if noprune := os.Getenv("OLLAMA_NOPRUNE"); noprune == "" {
		if err := server.PruneLayers(); err != nil {
			return err
		}
Michael Yang's avatar
Michael Yang committed
777
778
779
780
781
782
783
784
785

		manifestsPath, err := server.GetManifestPath()
		if err != nil {
			return err
		}

		if err := server.PruneDirectory(manifestsPath); err != nil {
			return err
		}
786
787
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
788
	return server.Serve(ln, origins)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
789
790
}

791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
func initializeKeypair() error {
	home, err := os.UserHomeDir()
	if err != nil {
		return err
	}

	privKeyPath := filepath.Join(home, ".ollama", "id_ed25519")
	pubKeyPath := filepath.Join(home, ".ollama", "id_ed25519.pub")

	_, err = os.Stat(privKeyPath)
	if os.IsNotExist(err) {
		fmt.Printf("Couldn't find '%s'. Generating new private key.\n", privKeyPath)
		_, privKey, err := ed25519.GenerateKey(rand.Reader)
		if err != nil {
			return err
		}

		privKeyBytes, err := format.OpenSSHPrivateKey(privKey, "")
		if err != nil {
			return err
		}

Michael Yang's avatar
Michael Yang committed
813
		err = os.MkdirAll(filepath.Dir(privKeyPath), 0o755)
814
815
816
817
		if err != nil {
			return fmt.Errorf("could not create directory %w", err)
		}

818
		err = os.WriteFile(privKeyPath, pem.EncodeToMemory(privKeyBytes), 0o600)
819
820
821
822
823
824
825
826
827
828
829
		if err != nil {
			return err
		}

		sshPrivateKey, err := ssh.NewSignerFromKey(privKey)
		if err != nil {
			return err
		}

		pubKeyData := ssh.MarshalAuthorizedKey(sshPrivateKey.PublicKey())

830
		err = os.WriteFile(pubKeyPath, pubKeyData, 0o644)
831
832
833
834
835
836
837
838
839
		if err != nil {
			return err
		}

		fmt.Printf("Your new public key is: \n\n%s\n", string(pubKeyData))
	}
	return nil
}

Bruce MacDonald's avatar
Bruce MacDonald committed
840
func startMacApp(client *api.Client) error {
Bruce MacDonald's avatar
Bruce MacDonald committed
841
842
843
844
845
	exe, err := os.Executable()
	if err != nil {
		return err
	}
	link, err := os.Readlink(exe)
Bruce MacDonald's avatar
Bruce MacDonald committed
846
847
848
	if err != nil {
		return err
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
849
850
851
	if !strings.Contains(link, "Ollama.app") {
		return fmt.Errorf("could not find ollama app")
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
	path := strings.Split(link, "Ollama.app")
	if err := exec.Command("/usr/bin/open", "-a", path[0]+"Ollama.app").Run(); err != nil {
		return err
	}
	// wait for the server to start
	timeout := time.After(5 * time.Second)
	tick := time.Tick(500 * time.Millisecond)
	for {
		select {
		case <-timeout:
			return errors.New("timed out waiting for server to start")
		case <-tick:
			if err := client.Heartbeat(context.Background()); err == nil {
				return nil // server has started
			}
		}
	}
}

871
func checkServerHeartbeat(_ *cobra.Command, _ []string) error {
872
873
874
875
	client, err := api.FromEnv()
	if err != nil {
		return err
	}
876
	if err := client.Heartbeat(context.Background()); err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
877
878
879
880
881
		if !strings.Contains(err.Error(), "connection refused") {
			return err
		}
		if runtime.GOOS == "darwin" {
			if err := startMacApp(client); err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
882
				return fmt.Errorf("could not connect to ollama app, is it running?")
883
			}
Bruce MacDonald's avatar
Bruce MacDonald committed
884
		} else {
885
886
887
888
889
890
			return fmt.Errorf("could not connect to ollama server, run 'ollama serve' to start it")
		}
	}
	return nil
}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
891
892
893
894
func NewCLI() *cobra.Command {
	log.SetFlags(log.LstdFlags | log.Lshortfile)

	rootCmd := &cobra.Command{
895
896
897
898
		Use:           "ollama",
		Short:         "Large language model runner",
		SilenceUsage:  true,
		SilenceErrors: true,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
899
900
901
		CompletionOptions: cobra.CompletionOptions{
			DisableDefaultCmd: true,
		},
Michael Yang's avatar
Michael Yang committed
902
		Version: version.Version,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
903
904
905
906
	}

	cobra.EnableCommandSorting = false

907
	createCmd := &cobra.Command{
908
909
910
911
912
		Use:     "create MODEL",
		Short:   "Create a model from a Modelfile",
		Args:    cobra.MinimumNArgs(1),
		PreRunE: checkServerHeartbeat,
		RunE:    CreateHandler,
913
914
915
916
	}

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

Patrick Devine's avatar
Patrick Devine committed
917
918
919
920
921
922
923
924
925
926
927
928
929
930
	showCmd := &cobra.Command{
		Use:     "show MODEL",
		Short:   "Show information for a model",
		Args:    cobra.MinimumNArgs(1),
		PreRunE: checkServerHeartbeat,
		RunE:    ShowHandler,
	}

	showCmd.Flags().Bool("license", false, "Show license of a model")
	showCmd.Flags().Bool("modelfile", false, "Show Modelfile of a model")
	showCmd.Flags().Bool("parameters", false, "Show parameters of a model")
	showCmd.Flags().Bool("template", false, "Show template of a model")
	showCmd.Flags().Bool("system", false, "Show system prompt of a model")

Jeffrey Morgan's avatar
Jeffrey Morgan committed
931
	runCmd := &cobra.Command{
932
933
934
935
936
		Use:     "run MODEL [PROMPT]",
		Short:   "Run a model",
		Args:    cobra.MinimumNArgs(1),
		PreRunE: checkServerHeartbeat,
		RunE:    RunHandler,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
937
938
	}

939
	runCmd.Flags().Bool("verbose", false, "Show timings for response")
940
	runCmd.Flags().Bool("insecure", false, "Use an insecure registry")
941
	runCmd.Flags().Bool("nowordwrap", false, "Don't wrap words to the next line automatically")
942

Jeffrey Morgan's avatar
Jeffrey Morgan committed
943
944
945
946
	serveCmd := &cobra.Command{
		Use:     "serve",
		Aliases: []string{"start"},
		Short:   "Start ollama",
Michael Yang's avatar
Michael Yang committed
947
		RunE:    RunServer,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
948
949
	}

950
	pullCmd := &cobra.Command{
951
952
953
954
955
		Use:     "pull MODEL",
		Short:   "Pull a model from a registry",
		Args:    cobra.MinimumNArgs(1),
		PreRunE: checkServerHeartbeat,
		RunE:    PullHandler,
956
957
	}

958
959
	pullCmd.Flags().Bool("insecure", false, "Use an insecure registry")

960
	pushCmd := &cobra.Command{
961
962
963
964
965
		Use:     "push MODEL",
		Short:   "Push a model to a registry",
		Args:    cobra.MinimumNArgs(1),
		PreRunE: checkServerHeartbeat,
		RunE:    PushHandler,
966
967
	}

968
969
	pushCmd.Flags().Bool("insecure", false, "Use an insecure registry")

Patrick Devine's avatar
Patrick Devine committed
970
	listCmd := &cobra.Command{
971
		Use:     "list",
Patrick Devine's avatar
Patrick Devine committed
972
		Aliases: []string{"ls"},
973
		Short:   "List models",
974
		PreRunE: checkServerHeartbeat,
975
		RunE:    ListHandler,
976
977
	}

Patrick Devine's avatar
Patrick Devine committed
978
	copyCmd := &cobra.Command{
979
980
981
982
983
		Use:     "cp",
		Short:   "Copy a model",
		Args:    cobra.MinimumNArgs(2),
		PreRunE: checkServerHeartbeat,
		RunE:    CopyHandler,
Patrick Devine's avatar
Patrick Devine committed
984
985
	}

986
	deleteCmd := &cobra.Command{
987
988
989
990
991
		Use:     "rm",
		Short:   "Remove a model",
		Args:    cobra.MinimumNArgs(1),
		PreRunE: checkServerHeartbeat,
		RunE:    DeleteHandler,
Patrick Devine's avatar
Patrick Devine committed
992
993
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
994
995
	rootCmd.AddCommand(
		serveCmd,
996
		createCmd,
Patrick Devine's avatar
Patrick Devine committed
997
		showCmd,
998
		runCmd,
999
1000
		pullCmd,
		pushCmd,
Patrick Devine's avatar
Patrick Devine committed
1001
		listCmd,
Patrick Devine's avatar
Patrick Devine committed
1002
		copyCmd,
1003
		deleteCmd,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
1004
1005
1006
1007
	)

	return rootCmd
}