cmd.go 21.9 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
	"log"
	"net"
14
	"net/http"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
15
	"os"
16
	"os/exec"
17
	"os/signal"
18
	"path/filepath"
19
	"runtime"
Michael Yang's avatar
Michael Yang committed
20
	"strings"
21
	"syscall"
Michael Yang's avatar
Michael Yang committed
22
	"time"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
23

Patrick Devine's avatar
Patrick Devine committed
24
25
	"github.com/dustin/go-humanize"
	"github.com/olekukonko/tablewriter"
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"
Patrick Devine's avatar
Patrick Devine committed
33
	"github.com/jmorganca/ollama/readline"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
34
	"github.com/jmorganca/ollama/server"
Michael Yang's avatar
Michael Yang committed
35
	"github.com/jmorganca/ollama/version"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
36
37
)

38
func CreateHandler(cmd *cobra.Command, args []string) error {
39
	filename, _ := cmd.Flags().GetString("file")
40
41
42
43
44
	filename, err := filepath.Abs(filename)
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
45
	client, err := api.ClientFromEnvironment()
46
47
48
	if err != nil {
		return err
	}
49

Michael Yang's avatar
Michael Yang committed
50
51
	var spinner *Spinner

52
53
	var currentDigest string
	var bar *progressbar.ProgressBar
Michael Yang's avatar
Michael Yang committed
54

55
56
57
58
59
60
61
	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
62
63
64
65
66
67
			// pulling
			bar = progressbar.DefaultBytes(
				resp.Total,
				resp.Status,
			)
			bar.Set64(resp.Completed)
68
		} else if resp.Digest == currentDigest && resp.Digest != "" {
Michael Yang's avatar
Michael Yang committed
69
			bar.Set64(resp.Completed)
70
71
72
73
74
75
76
77
		} else {
			currentDigest = ""
			if spinner != nil {
				spinner.Stop()
			}
			spinner = NewSpinner(resp.Status)
			go spinner.Spin(100 * time.Millisecond)
		}
78

79
80
81
82
83
84
85
		return nil
	}

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

Michael Yang's avatar
Michael Yang committed
86
87
	if spinner != nil {
		spinner.Stop()
88
89
90
		if spinner.description != "success" {
			return errors.New("unexpected end to create model")
		}
Michael Yang's avatar
Michael Yang committed
91
92
	}

93
94
95
	return nil
}

96
func RunHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
97
	client, err := api.ClientFromEnvironment()
98
99
100
101
	if err != nil {
		return err
	}

102
103
104
	name := args[0]
	// check if the model exists on the server
	_, err = client.Show(context.Background(), &api.ShowRequest{Name: name})
Michael Yang's avatar
Michael Yang committed
105
106
107
108
	var statusError api.StatusError
	switch {
	case errors.As(err, &statusError) && statusError.StatusCode == http.StatusNotFound:
		if err := PullHandler(cmd, args); err != nil {
109
			return err
Michael Yang's avatar
Michael Yang committed
110
		}
Michael Yang's avatar
Michael Yang committed
111
112
	case err != nil:
		return err
113
114
	}

Michael Yang's avatar
Michael Yang committed
115
	return RunGenerate(cmd, args)
Bruce MacDonald's avatar
Bruce MacDonald committed
116
117
}

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

124
125
126
127
128
	insecure, err := cmd.Flags().GetBool("insecure")
	if err != nil {
		return err
	}

129
130
131
	var currentDigest string
	var bar *progressbar.ProgressBar

132
	request := api.PushRequest{Name: args[0], Insecure: insecure}
133
	fn := func(resp api.ProgressResponse) error {
134
135
136
		if resp.Digest != currentDigest && resp.Digest != "" {
			currentDigest = resp.Digest
			bar = progressbar.DefaultBytes(
Michael Yang's avatar
Michael Yang committed
137
				resp.Total,
138
139
140
				fmt.Sprintf("pushing %s...", resp.Digest[7:19]),
			)

Michael Yang's avatar
Michael Yang committed
141
			bar.Set64(resp.Completed)
142
		} else if resp.Digest == currentDigest && resp.Digest != "" {
Michael Yang's avatar
Michael Yang committed
143
			bar.Set64(resp.Completed)
144
145
146
147
		} else {
			currentDigest = ""
			fmt.Println(resp.Status)
		}
148
149
150
151
152
153
		return nil
	}

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

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

159
160
161
	return nil
}

162
func ListHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
163
	client, err := api.ClientFromEnvironment()
164
165
166
	if err != nil {
		return err
	}
Patrick Devine's avatar
Patrick Devine committed
167
168
169
170
171
172
173
174
175

	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
176
		if len(args) == 0 || strings.HasPrefix(m.Name, args[0]) {
Patrick Devine's avatar
Patrick Devine committed
177
			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
178
		}
Patrick Devine's avatar
Patrick Devine committed
179
180
181
	}

	table := tablewriter.NewWriter(os.Stdout)
Patrick Devine's avatar
Patrick Devine committed
182
	table.SetHeader([]string{"NAME", "ID", "SIZE", "MODIFIED"})
Patrick Devine's avatar
Patrick Devine committed
183
184
185
186
187
188
189
190
191
192
193
194
	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
}

195
func DeleteHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
196
	client, err := api.ClientFromEnvironment()
197
198
199
	if err != nil {
		return err
	}
200

201
202
203
204
205
206
	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)
207
208
209
210
	}
	return nil
}

Patrick Devine's avatar
Patrick Devine committed
211
func ShowHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
212
	client, err := api.ClientFromEnvironment()
Patrick Devine's avatar
Patrick Devine committed
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
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
	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 {
262
		return errors.New("only one of '--license', '--modelfile', '--parameters', '--system', or '--template' can be specified")
Patrick Devine's avatar
Patrick Devine committed
263
	} else if flagsSet == 0 {
264
		return errors.New("one of '--license', '--modelfile', '--parameters', '--system', or '--template' must be specified")
Patrick Devine's avatar
Patrick Devine committed
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
	}

	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
289
func CopyHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
290
	client, err := api.ClientFromEnvironment()
291
292
293
	if err != nil {
		return err
	}
Patrick Devine's avatar
Patrick Devine committed
294
295
296
297
298
299
300
301
302

	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
}

303
func PullHandler(cmd *cobra.Command, args []string) error {
304
305
306
307
308
309
	insecure, err := cmd.Flags().GetBool("insecure")
	if err != nil {
		return err
	}

	return pull(args[0], insecure)
310
311
}

312
func pull(model string, insecure bool) error {
Michael Yang's avatar
Michael Yang committed
313
	client, err := api.ClientFromEnvironment()
314
315
316
	if err != nil {
		return err
	}
317

318
	var currentDigest string
Bruce MacDonald's avatar
Bruce MacDonald committed
319
	var bar *progressbar.ProgressBar
Michael Yang's avatar
Michael Yang committed
320

321
	request := api.PullRequest{Name: model, Insecure: insecure}
322
323
324
	fn := func(resp api.ProgressResponse) error {
		if resp.Digest != currentDigest && resp.Digest != "" {
			currentDigest = resp.Digest
325
			bar = progressbar.DefaultBytes(
Michael Yang's avatar
Michael Yang committed
326
				resp.Total,
327
				fmt.Sprintf("pulling %s...", resp.Digest[7:19]),
328
			)
329

Michael Yang's avatar
Michael Yang committed
330
			bar.Set64(resp.Completed)
331
		} else if resp.Digest == currentDigest && resp.Digest != "" {
Michael Yang's avatar
Michael Yang committed
332
			bar.Set64(resp.Completed)
333
		} else {
334
			currentDigest = ""
335
336
			fmt.Println(resp.Status)
		}
337

338
339
		return nil
	}
340

341
342
343
	if err := client.Pull(context.Background(), &request, fn); err != nil {
		return err
	}
344
345
346
347
348

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

349
	return nil
Michael Yang's avatar
Michael Yang committed
350
351
}

352
func RunGenerate(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
353
	if len(args) > 1 {
Michael Yang's avatar
Michael Yang committed
354
		// join all args into a single prompt
355
356
357
358
359
360
361
362
363
364
365
366
367
368
		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
369
370
	}

Michael Yang's avatar
Michael Yang committed
371
	if readline.IsTerminal(int(os.Stdin.Fd())) {
372
		return generateInteractive(cmd, args[0])
Michael Yang's avatar
Michael Yang committed
373
374
	}

375
	return generateBatch(cmd, args[0])
Michael Yang's avatar
Michael Yang committed
376
377
}

Michael Yang's avatar
Michael Yang committed
378
type generateContextKey string
Michael Yang's avatar
Michael Yang committed
379

380
func generate(cmd *cobra.Command, model, prompt string, wordWrap bool) error {
Michael Yang's avatar
Michael Yang committed
381
	client, err := api.ClientFromEnvironment()
Patrick Devine's avatar
Patrick Devine committed
382
383
384
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
385

Patrick Devine's avatar
Patrick Devine committed
386
387
	spinner := NewSpinner("")
	go spinner.Spin(60 * time.Millisecond)
Michael Yang's avatar
Michael Yang committed
388

Patrick Devine's avatar
Patrick Devine committed
389
	var latest api.GenerateResponse
390

Patrick Devine's avatar
Patrick Devine committed
391
392
393
394
	generateContext, ok := cmd.Context().Value(generateContextKey("context")).([]int)
	if !ok {
		generateContext = []int{}
	}
Michael Yang's avatar
Michael Yang committed
395

396
397
	termWidth, _, err := term.GetSize(int(0))
	if err != nil {
398
		wordWrap = false
399
400
	}

401
402
403
404
405
406
407
408
409
410
411
412
413
	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
	}()

414
415
416
	var currentLineLength int
	var wordBuffer string

Patrick Devine's avatar
Patrick Devine committed
417
418
419
420
421
	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
422

Patrick Devine's avatar
Patrick Devine committed
423
		latest = response
424

425
		if wordWrap {
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
			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
450
451
		return nil
	}
452

453
	if err := client.Generate(cancelCtx, &request, fn); err != nil {
454
		if strings.Contains(err.Error(), "context canceled") && abort {
455
456
			spinner.Finish()
			return nil
457
		}
Patrick Devine's avatar
Patrick Devine committed
458
459
460
		return err
	}
	if prompt != "" {
Michael Yang's avatar
Michael Yang committed
461
462
		fmt.Println()
		fmt.Println()
Patrick Devine's avatar
Patrick Devine committed
463
	}
464

Patrick Devine's avatar
Patrick Devine committed
465
	if !latest.Done {
466
467
468
		if abort {
			return nil
		}
Patrick Devine's avatar
Patrick Devine committed
469
470
		return errors.New("unexpected end of response")
	}
471

Patrick Devine's avatar
Patrick Devine committed
472
473
474
475
	verbose, err := cmd.Flags().GetBool("verbose")
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
476

Patrick Devine's avatar
Patrick Devine committed
477
478
	if verbose {
		latest.Summary()
Michael Yang's avatar
Michael Yang committed
479
	}
Michael Yang's avatar
Michael Yang committed
480

Patrick Devine's avatar
Patrick Devine committed
481
482
483
484
	ctx := cmd.Context()
	ctx = context.WithValue(ctx, generateContextKey("context"), latest.Context)
	cmd.SetContext(ctx)

Michael Yang's avatar
Michael Yang committed
485
486
487
	return nil
}

488
func generateInteractive(cmd *cobra.Command, model string) error {
Patrick Devine's avatar
Patrick Devine committed
489
	// load the model
490
	if err := generate(cmd, model, "", false); err != nil {
Patrick Devine's avatar
Patrick Devine committed
491
492
493
		return err
	}

Michael Yang's avatar
Michael Yang committed
494
	usage := func() {
Patrick Devine's avatar
Patrick Devine committed
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
		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
524
525
	}

Patrick Devine's avatar
Patrick Devine committed
526
527
528
529
530
	prompt := readline.Prompt{
		Prompt:         ">>> ",
		AltPrompt:      "... ",
		Placeholder:    "Send a message (/? for help)",
		AltPlaceholder: `Use """ to end multi-line input`,
Michael Yang's avatar
Michael Yang committed
531
532
	}

Patrick Devine's avatar
Patrick Devine committed
533
	scanner, err := readline.New(prompt)
Michael Yang's avatar
Michael Yang committed
534
535
536
537
	if err != nil {
		return err
	}

538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
	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
	}

553
554
555
	fmt.Print(readline.StartBracketedPaste)
	defer fmt.Printf(readline.EndBracketedPaste)

556
557
	var multiLineBuffer string

Michael Yang's avatar
Michael Yang committed
558
559
560
561
	for {
		line, err := scanner.Readline()
		switch {
		case errors.Is(err, io.EOF):
562
			fmt.Println()
Michael Yang's avatar
Michael Yang committed
563
564
			return nil
		case errors.Is(err, readline.ErrInterrupt):
565
			if line == "" {
Patrick Devine's avatar
Patrick Devine committed
566
				fmt.Println("\nUse Ctrl-D or /bye to exit.")
567
568
			}

Michael Yang's avatar
Michael Yang committed
569
570
			continue
		case err != nil:
Michael Yang's avatar
Michael Yang committed
571
572
573
			return err
		}

Michael Yang's avatar
Michael Yang committed
574
		line = strings.TrimSpace(line)
Michael Yang's avatar
Michael Yang committed
575

Michael Yang's avatar
Michael Yang committed
576
		switch {
Patrick Devine's avatar
Patrick Devine committed
577
		case scanner.Prompt.UseAlt:
578
			if strings.HasSuffix(line, `"""`) {
Patrick Devine's avatar
Patrick Devine committed
579
				scanner.Prompt.UseAlt = false
580
				multiLineBuffer += strings.TrimSuffix(line, `"""`)
581
				line = multiLineBuffer
582
583
584
585
586
587
				multiLineBuffer = ""
			} else {
				multiLineBuffer += line + " "
				continue
			}
		case strings.HasPrefix(line, `"""`):
Patrick Devine's avatar
Patrick Devine committed
588
			scanner.Prompt.UseAlt = true
589
590
			multiLineBuffer = strings.TrimPrefix(line, `"""`) + " "
			continue
Michael Yang's avatar
Michael Yang committed
591
592
		case strings.HasPrefix(line, "/list"):
			args := strings.Fields(line)
593
			if err := ListHandler(cmd, args[1:]); err != nil {
Michael Yang's avatar
Michael Yang committed
594
595
596
597
598
599
600
601
602
603
				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()
604
				case "wordwrap":
605
					wordWrap = true
606
607
					fmt.Println("Set 'wordwrap' mode.")
				case "nowordwrap":
608
					wordWrap = false
609
					fmt.Println("Set 'nowordwrap' mode.")
Michael Yang's avatar
Michael Yang committed
610
611
				case "verbose":
					cmd.Flags().Set("verbose", "true")
612
					fmt.Println("Set 'verbose' mode.")
Michael Yang's avatar
Michael Yang committed
613
614
				case "quiet":
					cmd.Flags().Set("verbose", "false")
615
					fmt.Println("Set 'quiet' mode.")
616
617
				default:
					fmt.Printf("Unknown command '/set %s'. Type /? for help\n", args[1])
618
619
				}
			} else {
Patrick Devine's avatar
Patrick Devine committed
620
				usageSet()
621
622
623
624
			}
		case strings.HasPrefix(line, "/show"):
			args := strings.Fields(line)
			if len(args) > 1 {
625
626
627
628
629
630
				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})
631
				if err != nil {
Patrick Devine's avatar
Patrick Devine committed
632
					fmt.Println("error: couldn't get model")
633
					return err
634
				}
Patrick Devine's avatar
Patrick Devine committed
635

636
637
				switch args[1] {
				case "license":
638
					if resp.License == "" {
639
						fmt.Print("No license was specified for this model.\n\n")
640
641
642
					} else {
						fmt.Println(resp.License)
					}
Patrick Devine's avatar
Patrick Devine committed
643
644
645
				case "modelfile":
					fmt.Println(resp.Modelfile)
				case "parameters":
646
					if resp.Parameters == "" {
647
						fmt.Print("No parameters were specified for this model.\n\n")
648
649
650
					} else {
						fmt.Println(resp.Parameters)
					}
651
				case "system":
652
					if resp.System == "" {
653
						fmt.Print("No system prompt was specified for this model.\n\n")
654
655
656
					} else {
						fmt.Println(resp.System)
					}
657
				case "template":
658
					if resp.Template == "" {
659
						fmt.Print("No prompt template was specified for this model.\n\n")
660
661
662
					} else {
						fmt.Println(resp.Template)
					}
663
				default:
664
					fmt.Printf("Unknown command '/show %s'. Type /? for help\n", args[1])
Michael Yang's avatar
Michael Yang committed
665
				}
Patrick Devine's avatar
Patrick Devine committed
666
667
668
669
670
671
672
673
674
675
676
677
			} 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()
				}
678
679
			} else {
				usage()
Michael Yang's avatar
Michael Yang committed
680
681
682
			}
		case line == "/exit", line == "/bye":
			return nil
683
684
685
		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
686
687
		}

Patrick Devine's avatar
Patrick Devine committed
688
		if len(line) > 0 && line[0] != '/' {
689
			if err := generate(cmd, model, line, wordWrap); err != nil {
Patrick Devine's avatar
Patrick Devine committed
690
691
				return err
			}
Michael Yang's avatar
Michael Yang committed
692
693
		}
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
694
695
}

696
func generateBatch(cmd *cobra.Command, model string) error {
Michael Yang's avatar
Michael Yang committed
697
698
699
700
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		prompt := scanner.Text()
		fmt.Printf(">>> %s\n", prompt)
701
		if err := generate(cmd, model, prompt, false); err != nil {
Michael Yang's avatar
Michael Yang committed
702
703
704
705
706
707
708
			return err
		}
	}

	return nil
}

709
func RunServer(cmd *cobra.Command, _ []string) error {
Michael Yang's avatar
Michael Yang committed
710
711
712
	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
713
		if ip := net.ParseIP(strings.Trim(os.Getenv("OLLAMA_HOST"), "[]")); ip != nil {
Michael Yang's avatar
Michael Yang committed
714
715
			host = ip.String()
		}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
716
	}
717

Michael Yang's avatar
Michael Yang committed
718
	if err := initializeKeypair(); err != nil {
719
720
721
		return err
	}

Michael Yang's avatar
Michael Yang committed
722
	ln, err := net.Listen("tcp", net.JoinHostPort(host, port))
723
724
725
	if err != nil {
		return err
	}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
726

727
728
729
730
731
	var origins []string
	if o := os.Getenv("OLLAMA_ORIGINS"); o != "" {
		origins = strings.Split(o, ",")
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
732
	return server.Serve(ln, origins)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
733
734
}

735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
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
757
		err = os.MkdirAll(filepath.Dir(privKeyPath), 0o755)
758
759
760
761
		if err != nil {
			return fmt.Errorf("could not create directory %w", err)
		}

762
		err = os.WriteFile(privKeyPath, pem.EncodeToMemory(privKeyBytes), 0o600)
763
764
765
766
767
768
769
770
771
772
773
		if err != nil {
			return err
		}

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

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

774
		err = os.WriteFile(pubKeyPath, pubKeyData, 0o644)
775
776
777
778
779
780
781
782
783
		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
784
func startMacApp(client *api.Client) error {
Bruce MacDonald's avatar
Bruce MacDonald committed
785
786
787
788
789
	exe, err := os.Executable()
	if err != nil {
		return err
	}
	link, err := os.Readlink(exe)
Bruce MacDonald's avatar
Bruce MacDonald committed
790
791
792
	if err != nil {
		return err
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
793
794
795
	if !strings.Contains(link, "Ollama.app") {
		return fmt.Errorf("could not find ollama app")
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
	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
			}
		}
	}
}

815
func checkServerHeartbeat(_ *cobra.Command, _ []string) error {
Michael Yang's avatar
Michael Yang committed
816
	client, err := api.ClientFromEnvironment()
817
818
819
	if err != nil {
		return err
	}
820
	if err := client.Heartbeat(context.Background()); err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
821
822
823
824
825
		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
826
				return fmt.Errorf("could not connect to ollama app, is it running?")
827
			}
Bruce MacDonald's avatar
Bruce MacDonald committed
828
		} else {
829
830
831
832
833
834
			return fmt.Errorf("could not connect to ollama server, run 'ollama serve' to start it")
		}
	}
	return nil
}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
835
836
837
838
func NewCLI() *cobra.Command {
	log.SetFlags(log.LstdFlags | log.Lshortfile)

	rootCmd := &cobra.Command{
839
840
841
842
		Use:           "ollama",
		Short:         "Large language model runner",
		SilenceUsage:  true,
		SilenceErrors: true,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
843
844
845
		CompletionOptions: cobra.CompletionOptions{
			DisableDefaultCmd: true,
		},
Michael Yang's avatar
Michael Yang committed
846
		Version: version.Version,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
847
848
849
850
	}

	cobra.EnableCommandSorting = false

851
	createCmd := &cobra.Command{
852
853
		Use:     "create MODEL",
		Short:   "Create a model from a Modelfile",
Michael Yang's avatar
Michael Yang committed
854
		Args:    cobra.ExactArgs(1),
855
856
		PreRunE: checkServerHeartbeat,
		RunE:    CreateHandler,
857
858
859
860
	}

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

Patrick Devine's avatar
Patrick Devine committed
861
862
863
	showCmd := &cobra.Command{
		Use:     "show MODEL",
		Short:   "Show information for a model",
Michael Yang's avatar
Michael Yang committed
864
		Args:    cobra.ExactArgs(1),
Patrick Devine's avatar
Patrick Devine committed
865
866
867
868
869
870
871
872
873
874
		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
875
	runCmd := &cobra.Command{
876
877
878
879
880
		Use:     "run MODEL [PROMPT]",
		Short:   "Run a model",
		Args:    cobra.MinimumNArgs(1),
		PreRunE: checkServerHeartbeat,
		RunE:    RunHandler,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
881
882
	}

883
	runCmd.Flags().Bool("verbose", false, "Show timings for response")
884
	runCmd.Flags().Bool("insecure", false, "Use an insecure registry")
885
	runCmd.Flags().Bool("nowordwrap", false, "Don't wrap words to the next line automatically")
886

Jeffrey Morgan's avatar
Jeffrey Morgan committed
887
888
889
890
	serveCmd := &cobra.Command{
		Use:     "serve",
		Aliases: []string{"start"},
		Short:   "Start ollama",
Michael Yang's avatar
Michael Yang committed
891
		Args:    cobra.ExactArgs(0),
Michael Yang's avatar
Michael Yang committed
892
		RunE:    RunServer,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
893
894
	}

895
	pullCmd := &cobra.Command{
896
897
		Use:     "pull MODEL",
		Short:   "Pull a model from a registry",
Michael Yang's avatar
Michael Yang committed
898
		Args:    cobra.ExactArgs(1),
899
900
		PreRunE: checkServerHeartbeat,
		RunE:    PullHandler,
901
902
	}

903
904
	pullCmd.Flags().Bool("insecure", false, "Use an insecure registry")

905
	pushCmd := &cobra.Command{
906
907
		Use:     "push MODEL",
		Short:   "Push a model to a registry",
Michael Yang's avatar
Michael Yang committed
908
		Args:    cobra.ExactArgs(1),
909
910
		PreRunE: checkServerHeartbeat,
		RunE:    PushHandler,
911
912
	}

913
914
	pushCmd.Flags().Bool("insecure", false, "Use an insecure registry")

Patrick Devine's avatar
Patrick Devine committed
915
	listCmd := &cobra.Command{
916
		Use:     "list",
Patrick Devine's avatar
Patrick Devine committed
917
		Aliases: []string{"ls"},
918
		Short:   "List models",
919
		PreRunE: checkServerHeartbeat,
920
		RunE:    ListHandler,
921
922
	}

Patrick Devine's avatar
Patrick Devine committed
923
	copyCmd := &cobra.Command{
Michael Yang's avatar
Michael Yang committed
924
		Use:     "cp SOURCE TARGET",
925
		Short:   "Copy a model",
Michael Yang's avatar
Michael Yang committed
926
		Args:    cobra.ExactArgs(2),
927
928
		PreRunE: checkServerHeartbeat,
		RunE:    CopyHandler,
Patrick Devine's avatar
Patrick Devine committed
929
930
	}

931
	deleteCmd := &cobra.Command{
Michael Yang's avatar
Michael Yang committed
932
		Use:     "rm MODEL [MODEL...]",
933
934
935
936
		Short:   "Remove a model",
		Args:    cobra.MinimumNArgs(1),
		PreRunE: checkServerHeartbeat,
		RunE:    DeleteHandler,
Patrick Devine's avatar
Patrick Devine committed
937
938
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
939
940
	rootCmd.AddCommand(
		serveCmd,
941
		createCmd,
Patrick Devine's avatar
Patrick Devine committed
942
		showCmd,
943
		runCmd,
944
945
		pullCmd,
		pushCmd,
Patrick Devine's avatar
Patrick Devine committed
946
		listCmd,
Patrick Devine's avatar
Patrick Devine committed
947
		copyCmd,
948
		deleteCmd,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
949
950
951
952
	)

	return rootCmd
}