cmd.go 23.6 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 {
691
692
693
694
695
696
				client, err := api.ClientFromEnvironment()
				if err != nil {
					fmt.Println("error: couldn't connect to ollama server")
					return err
				}
				resp, err := client.Show(cmd.Context(), &api.ShowRequest{Name: model})
697
				if err != nil {
Patrick Devine's avatar
Patrick Devine committed
698
					fmt.Println("error: couldn't get model")
699
					return err
700
				}
Patrick Devine's avatar
Patrick Devine committed
701

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

Patrick Devine's avatar
Patrick Devine committed
754
		if len(line) > 0 && line[0] != '/' {
755
			if err := generate(cmd, model, line, wordWrap); err != nil {
Patrick Devine's avatar
Patrick Devine committed
756
757
				return err
			}
Michael Yang's avatar
Michael Yang committed
758
759
		}
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
760
761
}

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

	return nil
}

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

Michael Yang's avatar
Michael Yang committed
784
	if err := initializeKeypair(); err != nil {
785
786
787
		return err
	}

Michael Yang's avatar
Michael Yang committed
788
	ln, err := net.Listen("tcp", net.JoinHostPort(host, port))
789
790
791
	if err != nil {
		return err
	}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
792

793
794
795
796
797
	var origins []string
	if o := os.Getenv("OLLAMA_ORIGINS"); o != "" {
		origins = strings.Split(o, ",")
	}

798
799
800
801
	if noprune := os.Getenv("OLLAMA_NOPRUNE"); noprune == "" {
		if err := server.PruneLayers(); err != nil {
			return err
		}
Michael Yang's avatar
Michael Yang committed
802
803
804
805
806
807
808
809
810

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

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

Jeffrey Morgan's avatar
Jeffrey Morgan committed
813
	return server.Serve(ln, origins)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
814
815
}

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

843
		err = os.WriteFile(privKeyPath, pem.EncodeToMemory(privKeyBytes), 0o600)
844
845
846
847
848
849
850
851
852
853
854
		if err != nil {
			return err
		}

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

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

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

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

Jeffrey Morgan's avatar
Jeffrey Morgan committed
916
917
918
919
func NewCLI() *cobra.Command {
	log.SetFlags(log.LstdFlags | log.Lshortfile)

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

	cobra.EnableCommandSorting = false

932
	createCmd := &cobra.Command{
933
934
935
936
937
		Use:     "create MODEL",
		Short:   "Create a model from a Modelfile",
		Args:    cobra.MinimumNArgs(1),
		PreRunE: checkServerHeartbeat,
		RunE:    CreateHandler,
938
939
940
941
	}

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

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

964
	runCmd.Flags().Bool("verbose", false, "Show timings for response")
965
	runCmd.Flags().Bool("insecure", false, "Use an insecure registry")
966
	runCmd.Flags().Bool("nowordwrap", false, "Don't wrap words to the next line automatically")
967

Jeffrey Morgan's avatar
Jeffrey Morgan committed
968
969
970
971
	serveCmd := &cobra.Command{
		Use:     "serve",
		Aliases: []string{"start"},
		Short:   "Start ollama",
Michael Yang's avatar
Michael Yang committed
972
		RunE:    RunServer,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
973
974
	}

975
	pullCmd := &cobra.Command{
976
977
978
979
980
		Use:     "pull MODEL",
		Short:   "Pull a model from a registry",
		Args:    cobra.MinimumNArgs(1),
		PreRunE: checkServerHeartbeat,
		RunE:    PullHandler,
981
982
	}

983
984
	pullCmd.Flags().Bool("insecure", false, "Use an insecure registry")

985
	pushCmd := &cobra.Command{
986
987
988
989
990
		Use:     "push MODEL",
		Short:   "Push a model to a registry",
		Args:    cobra.MinimumNArgs(1),
		PreRunE: checkServerHeartbeat,
		RunE:    PushHandler,
991
992
	}

993
994
	pushCmd.Flags().Bool("insecure", false, "Use an insecure registry")

Patrick Devine's avatar
Patrick Devine committed
995
	listCmd := &cobra.Command{
996
		Use:     "list",
Patrick Devine's avatar
Patrick Devine committed
997
		Aliases: []string{"ls"},
998
		Short:   "List models",
999
		PreRunE: checkServerHeartbeat,
1000
		RunE:    ListHandler,
1001
1002
	}

Patrick Devine's avatar
Patrick Devine committed
1003
	copyCmd := &cobra.Command{
1004
1005
1006
1007
1008
		Use:     "cp",
		Short:   "Copy a model",
		Args:    cobra.MinimumNArgs(2),
		PreRunE: checkServerHeartbeat,
		RunE:    CopyHandler,
Patrick Devine's avatar
Patrick Devine committed
1009
1010
	}

1011
	deleteCmd := &cobra.Command{
1012
1013
1014
1015
1016
		Use:     "rm",
		Short:   "Remove a model",
		Args:    cobra.MinimumNArgs(1),
		PreRunE: checkServerHeartbeat,
		RunE:    DeleteHandler,
Patrick Devine's avatar
Patrick Devine committed
1017
1018
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
1019
1020
	rootCmd.AddCommand(
		serveCmd,
1021
		createCmd,
Patrick Devine's avatar
Patrick Devine committed
1022
		showCmd,
1023
		runCmd,
1024
1025
		pullCmd,
		pushCmd,
Patrick Devine's avatar
Patrick Devine committed
1026
		listCmd,
Patrick Devine's avatar
Patrick Devine committed
1027
		copyCmd,
1028
		deleteCmd,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
1029
1030
1031
1032
	)

	return rootCmd
}