cmd.go 23.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
	}

Michael Yang's avatar
Michael Yang committed
64
	client, err := api.ClientFromEnvironment()
65
66
67
	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
83
84
85
86
			// pulling
			bar = progressbar.DefaultBytes(
				resp.Total,
				resp.Status,
			)
			bar.Set64(resp.Completed)
87
		} else if resp.Digest == currentDigest && resp.Digest != "" {
Michael Yang's avatar
Michael Yang committed
88
			bar.Set64(resp.Completed)
89
90
91
92
93
94
95
96
		} else {
			currentDigest = ""
			if spinner != nil {
				spinner.Stop()
			}
			spinner = NewSpinner(resp.Status)
			go spinner.Spin(100 * time.Millisecond)
		}
97

98
99
100
101
102
103
104
		return nil
	}

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

Michael Yang's avatar
Michael Yang committed
105
106
	if spinner != nil {
		spinner.Stop()
107
108
109
		if spinner.description != "success" {
			return errors.New("unexpected end to create model")
		}
Michael Yang's avatar
Michael Yang committed
110
111
	}

112
113
114
	return nil
}

115
func RunHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
116
	client, err := api.ClientFromEnvironment()
117
118
119
120
	if err != nil {
		return err
	}

121
	models, err := client.List(context.Background())
Patrick Devine's avatar
Patrick Devine committed
122
123
124
125
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
126
	canonicalModelPath := server.ParseModelPath(args[0])
127
	for _, model := range models.Models {
Michael Yang's avatar
Michael Yang committed
128
		if model.Name == canonicalModelPath.GetShortTagname() {
129
			return RunGenerate(cmd, args)
Michael Yang's avatar
Michael Yang committed
130
		}
131
132
133
	}

	if err := PullHandler(cmd, args); err != nil {
Michael Yang's avatar
Michael Yang committed
134
135
136
137
		return err
	}

	return RunGenerate(cmd, args)
Bruce MacDonald's avatar
Bruce MacDonald committed
138
139
}

140
func PushHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
141
	client, err := api.ClientFromEnvironment()
142
143
144
	if err != nil {
		return err
	}
145

146
147
148
149
150
	insecure, err := cmd.Flags().GetBool("insecure")
	if err != nil {
		return err
	}

151
152
153
	var currentDigest string
	var bar *progressbar.ProgressBar

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

Michael Yang's avatar
Michael Yang committed
163
			bar.Set64(resp.Completed)
164
		} else if resp.Digest == currentDigest && resp.Digest != "" {
Michael Yang's avatar
Michael Yang committed
165
			bar.Set64(resp.Completed)
166
167
168
169
		} else {
			currentDigest = ""
			fmt.Println(resp.Status)
		}
170
171
172
173
174
175
		return nil
	}

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

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

181
182
183
	return nil
}

184
func ListHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
185
	client, err := api.ClientFromEnvironment()
186
187
188
	if err != nil {
		return err
	}
Patrick Devine's avatar
Patrick Devine committed
189
190
191
192
193
194
195
196
197

	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
198
		if len(args) == 0 || strings.HasPrefix(m.Name, args[0]) {
Patrick Devine's avatar
Patrick Devine committed
199
			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
200
		}
Patrick Devine's avatar
Patrick Devine committed
201
202
203
	}

	table := tablewriter.NewWriter(os.Stdout)
Patrick Devine's avatar
Patrick Devine committed
204
	table.SetHeader([]string{"NAME", "ID", "SIZE", "MODIFIED"})
Patrick Devine's avatar
Patrick Devine committed
205
206
207
208
209
210
211
212
213
214
215
216
	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
}

217
func DeleteHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
218
	client, err := api.ClientFromEnvironment()
219
220
221
	if err != nil {
		return err
	}
222

223
224
225
226
227
228
	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)
229
230
231
232
	}
	return nil
}

Patrick Devine's avatar
Patrick Devine committed
233
func ShowHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
234
	client, err := api.ClientFromEnvironment()
Patrick Devine's avatar
Patrick Devine committed
235
236
237
238
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
	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 {
284
		return errors.New("only one of '--license', '--modelfile', '--parameters', '--system', or '--template' can be specified")
Patrick Devine's avatar
Patrick Devine committed
285
	} else if flagsSet == 0 {
286
		return errors.New("one of '--license', '--modelfile', '--parameters', '--system', or '--template' must be specified")
Patrick Devine's avatar
Patrick Devine committed
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
	}

	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
311
func CopyHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
312
	client, err := api.ClientFromEnvironment()
313
314
315
	if err != nil {
		return err
	}
Patrick Devine's avatar
Patrick Devine committed
316
317
318
319
320
321
322
323
324

	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
}

325
func PullHandler(cmd *cobra.Command, args []string) error {
326
327
328
329
330
331
	insecure, err := cmd.Flags().GetBool("insecure")
	if err != nil {
		return err
	}

	return pull(args[0], insecure)
332
333
}

334
func pull(model string, insecure bool) error {
Michael Yang's avatar
Michael Yang committed
335
	client, err := api.ClientFromEnvironment()
336
337
338
	if err != nil {
		return err
	}
339

340
	var currentDigest string
Bruce MacDonald's avatar
Bruce MacDonald committed
341
	var bar *progressbar.ProgressBar
Michael Yang's avatar
Michael Yang committed
342

343
	request := api.PullRequest{Name: model, Insecure: insecure}
344
345
346
	fn := func(resp api.ProgressResponse) error {
		if resp.Digest != currentDigest && resp.Digest != "" {
			currentDigest = resp.Digest
347
			bar = progressbar.DefaultBytes(
Michael Yang's avatar
Michael Yang committed
348
				resp.Total,
349
				fmt.Sprintf("pulling %s...", resp.Digest[7:19]),
350
			)
351

Michael Yang's avatar
Michael Yang committed
352
			bar.Set64(resp.Completed)
353
		} else if resp.Digest == currentDigest && resp.Digest != "" {
Michael Yang's avatar
Michael Yang committed
354
			bar.Set64(resp.Completed)
355
		} else {
356
			currentDigest = ""
357
358
			fmt.Println(resp.Status)
		}
359

360
361
		return nil
	}
362

363
364
365
	if err := client.Pull(context.Background(), &request, fn); err != nil {
		return err
	}
366
367
368
369
370

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

371
	return nil
Michael Yang's avatar
Michael Yang committed
372
373
}

374
func RunGenerate(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
375
	if len(args) > 1 {
Michael Yang's avatar
Michael Yang committed
376
		// join all args into a single prompt
377
378
379
380
381
382
383
384
385
386
387
388
389
390
		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
391
392
	}

Michael Yang's avatar
Michael Yang committed
393
	if readline.IsTerminal(int(os.Stdin.Fd())) {
394
		return generateInteractive(cmd, args[0])
Michael Yang's avatar
Michael Yang committed
395
396
	}

397
	return generateBatch(cmd, args[0])
Michael Yang's avatar
Michael Yang committed
398
399
}

Michael Yang's avatar
Michael Yang committed
400
type generateContextKey string
Michael Yang's avatar
Michael Yang committed
401

402
func generate(cmd *cobra.Command, model, prompt string, wordWrap bool) error {
Michael Yang's avatar
Michael Yang committed
403
	client, err := api.ClientFromEnvironment()
Patrick Devine's avatar
Patrick Devine committed
404
405
406
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
407

Patrick Devine's avatar
Patrick Devine committed
408
409
	spinner := NewSpinner("")
	go spinner.Spin(60 * time.Millisecond)
Michael Yang's avatar
Michael Yang committed
410

Patrick Devine's avatar
Patrick Devine committed
411
	var latest api.GenerateResponse
412

Patrick Devine's avatar
Patrick Devine committed
413
414
415
416
	generateContext, ok := cmd.Context().Value(generateContextKey("context")).([]int)
	if !ok {
		generateContext = []int{}
	}
Michael Yang's avatar
Michael Yang committed
417

418
419
	termWidth, _, err := term.GetSize(int(0))
	if err != nil {
420
		wordWrap = false
421
422
	}

423
424
425
426
427
428
429
430
431
432
433
434
435
	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
	}()

436
437
438
	var currentLineLength int
	var wordBuffer string

Patrick Devine's avatar
Patrick Devine committed
439
440
441
442
443
	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
444

Patrick Devine's avatar
Patrick Devine committed
445
		latest = response
446

447
		if wordWrap {
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
			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
472
473
		return nil
	}
474

475
	if err := client.Generate(cancelCtx, &request, fn); err != nil {
476
		if strings.Contains(err.Error(), "context canceled") && abort {
477
478
			spinner.Finish()
			return nil
479
		}
Patrick Devine's avatar
Patrick Devine committed
480
481
482
		return err
	}
	if prompt != "" {
Michael Yang's avatar
Michael Yang committed
483
484
		fmt.Println()
		fmt.Println()
Patrick Devine's avatar
Patrick Devine committed
485
	}
486

Patrick Devine's avatar
Patrick Devine committed
487
	if !latest.Done {
488
489
490
		if abort {
			return nil
		}
Patrick Devine's avatar
Patrick Devine committed
491
492
		return errors.New("unexpected end of response")
	}
493

Patrick Devine's avatar
Patrick Devine committed
494
495
496
497
	verbose, err := cmd.Flags().GetBool("verbose")
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
498

Patrick Devine's avatar
Patrick Devine committed
499
500
	if verbose {
		latest.Summary()
Michael Yang's avatar
Michael Yang committed
501
	}
Michael Yang's avatar
Michael Yang committed
502

Patrick Devine's avatar
Patrick Devine committed
503
504
505
506
	ctx := cmd.Context()
	ctx = context.WithValue(ctx, generateContextKey("context"), latest.Context)
	cmd.SetContext(ctx)

Michael Yang's avatar
Michael Yang committed
507
508
509
	return nil
}

510
func generateInteractive(cmd *cobra.Command, model string) error {
Michael Yang's avatar
Michael Yang committed
511
512
513
514
515
	home, err := os.UserHomeDir()
	if err != nil {
		return err
	}

Patrick Devine's avatar
Patrick Devine committed
516
	// load the model
517
	if err := generate(cmd, model, "", false); err != nil {
Patrick Devine's avatar
Patrick Devine committed
518
519
520
		return err
	}

Michael Yang's avatar
Michael Yang committed
521
522
523
524
525
526
	completer := readline.NewPrefixCompleter(
		readline.PcItem("/help"),
		readline.PcItem("/list"),
		readline.PcItem("/set",
			readline.PcItem("history"),
			readline.PcItem("nohistory"),
527
528
			readline.PcItem("wordwrap"),
			readline.PcItem("nowordwrap"),
Michael Yang's avatar
Michael Yang committed
529
530
			readline.PcItem("verbose"),
			readline.PcItem("quiet"),
Michael Yang's avatar
Michael Yang committed
531
		),
532
533
		readline.PcItem("/show",
			readline.PcItem("license"),
Patrick Devine's avatar
Patrick Devine committed
534
535
			readline.PcItem("modelfile"),
			readline.PcItem("parameters"),
536
537
538
			readline.PcItem("system"),
			readline.PcItem("template"),
		),
Michael Yang's avatar
Michael Yang committed
539
540
541
542
543
		readline.PcItem("/exit"),
		readline.PcItem("/bye"),
	)

	usage := func() {
Patrick Devine's avatar
Patrick Devine committed
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
		fmt.Fprintln(os.Stderr, "Available Commands:")
		fmt.Fprintln(os.Stderr, "  /set         Set session variables")
		fmt.Fprintln(os.Stderr, "  /show        Show model information")
		fmt.Fprintln(os.Stderr, "  /bye         Exit")
		fmt.Fprintln(os.Stderr, "  /?, /help    Help for a command")
		fmt.Fprintln(os.Stderr, "")
		fmt.Fprintln(os.Stderr, "Use \"\"\" to begin a multi-line message.")
		fmt.Fprintln(os.Stderr, "")
	}

	usageSet := func() {
		fmt.Fprintln(os.Stderr, "Available Commands:")
		fmt.Fprintln(os.Stderr, "  /set history      Enable history")
		fmt.Fprintln(os.Stderr, "  /set nohistory    Disable history")
		fmt.Fprintln(os.Stderr, "  /set wordwrap     Enable wordwrap")
		fmt.Fprintln(os.Stderr, "  /set nowordwrap   Disable wordwrap")
		fmt.Fprintln(os.Stderr, "  /set verbose      Show LLM stats")
		fmt.Fprintln(os.Stderr, "  /set quiet        Disable LLM stats")
		fmt.Fprintln(os.Stderr, "")
	}

	usageShow := func() {
		fmt.Fprintln(os.Stderr, "Available Commands:")
		fmt.Fprintln(os.Stderr, "  /show license      Show model license")
		fmt.Fprintln(os.Stderr, "  /show modelfile    Show Modelfile for this model")
		fmt.Fprintln(os.Stderr, "  /show parameters   Show parameters for this model")
		fmt.Fprintln(os.Stderr, "  /show system       Show system prompt")
		fmt.Fprintln(os.Stderr, "  /show template     Show prompt template")
		fmt.Fprintln(os.Stderr, "")
Michael Yang's avatar
Michael Yang committed
573
574
	}

575
576
	var painter Painter

Michael Yang's avatar
Michael Yang committed
577
	config := readline.Config{
578
		Painter:      &painter,
Michael Yang's avatar
Michael Yang committed
579
580
581
582
583
584
585
586
587
588
589
		Prompt:       ">>> ",
		HistoryFile:  filepath.Join(home, ".ollama", "history"),
		AutoComplete: completer,
	}

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

590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
	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
	}

605
606
607
	var multiLineBuffer string
	var isMultiLine bool

Michael Yang's avatar
Michael Yang committed
608
609
610
611
612
613
	for {
		line, err := scanner.Readline()
		switch {
		case errors.Is(err, io.EOF):
			return nil
		case errors.Is(err, readline.ErrInterrupt):
614
			if line == "" {
615
				fmt.Println("Use Ctrl-D or /bye to exit.")
616
617
			}

Michael Yang's avatar
Michael Yang committed
618
619
			continue
		case err != nil:
Michael Yang's avatar
Michael Yang committed
620
621
622
			return err
		}

Michael Yang's avatar
Michael Yang committed
623
		line = strings.TrimSpace(line)
Michael Yang's avatar
Michael Yang committed
624

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

697
698
				switch args[1] {
				case "license":
699
					if resp.License == "" {
700
						fmt.Print("No license was specified for this model.\n\n")
701
702
703
					} else {
						fmt.Println(resp.License)
					}
Patrick Devine's avatar
Patrick Devine committed
704
705
706
				case "modelfile":
					fmt.Println(resp.Modelfile)
				case "parameters":
707
					if resp.Parameters == "" {
708
						fmt.Print("No parameters were specified for this model.\n\n")
709
710
711
					} else {
						fmt.Println(resp.Parameters)
					}
712
				case "system":
713
					if resp.System == "" {
714
						fmt.Print("No system prompt was specified for this model.\n\n")
715
716
717
					} else {
						fmt.Println(resp.System)
					}
718
				case "template":
719
					if resp.Template == "" {
720
						fmt.Print("No prompt template was specified for this model.\n\n")
721
722
723
					} else {
						fmt.Println(resp.Template)
					}
724
				default:
725
					fmt.Printf("Unknown command '/show %s'. Type /? for help\n", args[1])
Michael Yang's avatar
Michael Yang committed
726
				}
Patrick Devine's avatar
Patrick Devine committed
727
728
729
730
731
732
733
734
735
736
737
738
			} else {
				usageShow()
			}
		case strings.HasPrefix(line, "/help"), strings.HasPrefix(line, "/?"):
			args := strings.Fields(line)
			if len(args) > 1 {
				switch args[1] {
				case "set", "/set":
					usageSet()
				case "show", "/show":
					usageShow()
				}
739
740
			} else {
				usage()
Michael Yang's avatar
Michael Yang committed
741
742
743
			}
		case line == "/exit", line == "/bye":
			return nil
744
745
746
		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
747
748
		}

Patrick Devine's avatar
Patrick Devine committed
749
		if len(line) > 0 && line[0] != '/' {
750
			if err := generate(cmd, model, line, wordWrap); err != nil {
Patrick Devine's avatar
Patrick Devine committed
751
752
				return err
			}
Michael Yang's avatar
Michael Yang committed
753
754
		}
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
755
756
}

757
func generateBatch(cmd *cobra.Command, model string) error {
Michael Yang's avatar
Michael Yang committed
758
759
760
761
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		prompt := scanner.Text()
		fmt.Printf(">>> %s\n", prompt)
762
		if err := generate(cmd, model, prompt, false); err != nil {
Michael Yang's avatar
Michael Yang committed
763
764
765
766
767
768
769
			return err
		}
	}

	return nil
}

770
func RunServer(cmd *cobra.Command, _ []string) error {
Michael Yang's avatar
Michael Yang committed
771
772
773
	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
774
		if ip := net.ParseIP(strings.Trim(os.Getenv("OLLAMA_HOST"), "[]")); ip != nil {
Michael Yang's avatar
Michael Yang committed
775
776
			host = ip.String()
		}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
777
	}
778

Michael Yang's avatar
Michael Yang committed
779
	if err := initializeKeypair(); err != nil {
780
781
782
		return err
	}

Michael Yang's avatar
Michael Yang committed
783
	ln, err := net.Listen("tcp", net.JoinHostPort(host, port))
784
785
786
	if err != nil {
		return err
	}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
787

788
789
790
791
792
	var origins []string
	if o := os.Getenv("OLLAMA_ORIGINS"); o != "" {
		origins = strings.Split(o, ",")
	}

793
794
795
796
	if noprune := os.Getenv("OLLAMA_NOPRUNE"); noprune == "" {
		if err := server.PruneLayers(); err != nil {
			return err
		}
Michael Yang's avatar
Michael Yang committed
797
798
799
800
801
802
803
804
805

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

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

Jeffrey Morgan's avatar
Jeffrey Morgan committed
808
	return server.Serve(ln, origins)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
809
810
}

811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
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
833
		err = os.MkdirAll(filepath.Dir(privKeyPath), 0o755)
834
835
836
837
		if err != nil {
			return fmt.Errorf("could not create directory %w", err)
		}

838
		err = os.WriteFile(privKeyPath, pem.EncodeToMemory(privKeyBytes), 0o600)
839
840
841
842
843
844
845
846
847
848
849
		if err != nil {
			return err
		}

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

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

850
		err = os.WriteFile(pubKeyPath, pubKeyData, 0o644)
851
852
853
854
855
856
857
858
859
		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
860
func startMacApp(client *api.Client) error {
Bruce MacDonald's avatar
Bruce MacDonald committed
861
862
863
864
865
	exe, err := os.Executable()
	if err != nil {
		return err
	}
	link, err := os.Readlink(exe)
Bruce MacDonald's avatar
Bruce MacDonald committed
866
867
868
	if err != nil {
		return err
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
869
870
871
	if !strings.Contains(link, "Ollama.app") {
		return fmt.Errorf("could not find ollama app")
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
	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
			}
		}
	}
}

891
func checkServerHeartbeat(_ *cobra.Command, _ []string) error {
Michael Yang's avatar
Michael Yang committed
892
	client, err := api.ClientFromEnvironment()
893
894
895
	if err != nil {
		return err
	}
896
	if err := client.Heartbeat(context.Background()); err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
897
898
899
900
901
		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
902
				return fmt.Errorf("could not connect to ollama app, is it running?")
903
			}
Bruce MacDonald's avatar
Bruce MacDonald committed
904
		} else {
905
906
907
908
909
910
			return fmt.Errorf("could not connect to ollama server, run 'ollama serve' to start it")
		}
	}
	return nil
}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
911
912
913
914
func NewCLI() *cobra.Command {
	log.SetFlags(log.LstdFlags | log.Lshortfile)

	rootCmd := &cobra.Command{
915
916
917
918
		Use:           "ollama",
		Short:         "Large language model runner",
		SilenceUsage:  true,
		SilenceErrors: true,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
919
920
921
		CompletionOptions: cobra.CompletionOptions{
			DisableDefaultCmd: true,
		},
Michael Yang's avatar
Michael Yang committed
922
		Version: version.Version,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
923
924
925
926
	}

	cobra.EnableCommandSorting = false

927
	createCmd := &cobra.Command{
928
929
930
931
932
		Use:     "create MODEL",
		Short:   "Create a model from a Modelfile",
		Args:    cobra.MinimumNArgs(1),
		PreRunE: checkServerHeartbeat,
		RunE:    CreateHandler,
933
934
935
936
	}

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

Patrick Devine's avatar
Patrick Devine committed
937
938
939
940
941
942
943
944
945
946
947
948
949
950
	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
951
	runCmd := &cobra.Command{
952
953
954
955
956
		Use:     "run MODEL [PROMPT]",
		Short:   "Run a model",
		Args:    cobra.MinimumNArgs(1),
		PreRunE: checkServerHeartbeat,
		RunE:    RunHandler,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
957
958
	}

959
	runCmd.Flags().Bool("verbose", false, "Show timings for response")
960
	runCmd.Flags().Bool("insecure", false, "Use an insecure registry")
961
	runCmd.Flags().Bool("nowordwrap", false, "Don't wrap words to the next line automatically")
962

Jeffrey Morgan's avatar
Jeffrey Morgan committed
963
964
965
966
	serveCmd := &cobra.Command{
		Use:     "serve",
		Aliases: []string{"start"},
		Short:   "Start ollama",
Michael Yang's avatar
Michael Yang committed
967
		RunE:    RunServer,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
968
969
	}

970
	pullCmd := &cobra.Command{
971
972
973
974
975
		Use:     "pull MODEL",
		Short:   "Pull a model from a registry",
		Args:    cobra.MinimumNArgs(1),
		PreRunE: checkServerHeartbeat,
		RunE:    PullHandler,
976
977
	}

978
979
	pullCmd.Flags().Bool("insecure", false, "Use an insecure registry")

980
	pushCmd := &cobra.Command{
981
982
983
984
985
		Use:     "push MODEL",
		Short:   "Push a model to a registry",
		Args:    cobra.MinimumNArgs(1),
		PreRunE: checkServerHeartbeat,
		RunE:    PushHandler,
986
987
	}

988
989
	pushCmd.Flags().Bool("insecure", false, "Use an insecure registry")

Patrick Devine's avatar
Patrick Devine committed
990
	listCmd := &cobra.Command{
991
		Use:     "list",
Patrick Devine's avatar
Patrick Devine committed
992
		Aliases: []string{"ls"},
993
		Short:   "List models",
994
		PreRunE: checkServerHeartbeat,
995
		RunE:    ListHandler,
996
997
	}

Patrick Devine's avatar
Patrick Devine committed
998
	copyCmd := &cobra.Command{
999
1000
1001
1002
1003
		Use:     "cp",
		Short:   "Copy a model",
		Args:    cobra.MinimumNArgs(2),
		PreRunE: checkServerHeartbeat,
		RunE:    CopyHandler,
Patrick Devine's avatar
Patrick Devine committed
1004
1005
	}

1006
	deleteCmd := &cobra.Command{
1007
1008
1009
1010
1011
		Use:     "rm",
		Short:   "Remove a model",
		Args:    cobra.MinimumNArgs(1),
		PreRunE: checkServerHeartbeat,
		RunE:    DeleteHandler,
Patrick Devine's avatar
Patrick Devine committed
1012
1013
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
1014
1015
	rootCmd.AddCommand(
		serveCmd,
1016
		createCmd,
Patrick Devine's avatar
Patrick Devine committed
1017
		showCmd,
1018
		runCmd,
1019
1020
		pullCmd,
		pushCmd,
Patrick Devine's avatar
Patrick Devine committed
1021
		listCmd,
Patrick Devine's avatar
Patrick Devine committed
1022
		copyCmd,
1023
		deleteCmd,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
1024
1025
1026
1027
	)

	return rootCmd
}