cmd.go 23.3 KB
Newer Older
Jeffrey Morgan's avatar
Jeffrey Morgan committed
1
2
3
package cmd

import (
Michael Yang's avatar
Michael Yang committed
4
	"bytes"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
5
	"context"
6
7
	"crypto/ed25519"
	"crypto/rand"
Michael Yang's avatar
Michael Yang committed
8
	"crypto/sha256"
9
	"encoding/pem"
Michael Yang's avatar
Michael Yang committed
10
	"errors"
Bruce MacDonald's avatar
Bruce MacDonald committed
11
	"fmt"
Michael Yang's avatar
Michael Yang committed
12
	"io"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
13
14
	"log"
	"net"
15
	"net/http"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
16
	"os"
17
	"os/exec"
18
	"os/signal"
19
	"path/filepath"
20
	"runtime"
Michael Yang's avatar
Michael Yang committed
21
	"strings"
22
	"syscall"
Michael Yang's avatar
Michael Yang committed
23
	"time"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
24

Patrick Devine's avatar
Patrick Devine committed
25
	"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"
Michael Yang's avatar
Michael Yang committed
32
	"github.com/jmorganca/ollama/parser"
33
	"github.com/jmorganca/ollama/progressbar"
Patrick Devine's avatar
Patrick Devine committed
34
	"github.com/jmorganca/ollama/readline"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
35
	"github.com/jmorganca/ollama/server"
Michael Yang's avatar
Michael Yang committed
36
	"github.com/jmorganca/ollama/version"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
37
38
)

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

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

Michael Yang's avatar
Michael Yang committed
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
	modelfile, err := os.ReadFile(filename)
	if err != nil {
		return err
	}

	spinner := NewSpinner("transferring context")
	go spinner.Spin(100 * time.Millisecond)

	commands, err := parser.Parse(bytes.NewReader(modelfile))
	if err != nil {
		return err
	}

	home, err := os.UserHomeDir()
	if err != nil {
		return err
	}

	for _, c := range commands {
		switch c.Name {
		case "model", "adapter":
			path := c.Args
			if path == "~" {
				path = home
			} else if strings.HasPrefix(path, "~/") {
				path = filepath.Join(home, path[2:])
			}

			bin, err := os.Open(path)
			if errors.Is(err, os.ErrNotExist) && c.Name == "model" {
Michael Yang's avatar
Michael Yang committed
81
				continue
Michael Yang's avatar
Michael Yang committed
82
83
84
85
86
87
88
89
90
91
92
93
			} else if err != nil {
				return err
			}
			defer bin.Close()

			hash := sha256.New()
			if _, err := io.Copy(hash, bin); err != nil {
				return err
			}
			bin.Seek(0, io.SeekStart)

			digest := fmt.Sprintf("sha256:%x", hash.Sum(nil))
Michael Yang's avatar
Michael Yang committed
94
			if err = client.CreateBlob(cmd.Context(), digest, bin); err != nil {
Michael Yang's avatar
Michael Yang committed
95
96
97
				return err
			}

Michael Yang's avatar
Michael Yang committed
98
			modelfile = bytes.ReplaceAll(modelfile, []byte(c.Args), []byte("@"+digest))
Michael Yang's avatar
Michael Yang committed
99
100
		}
	}
Michael Yang's avatar
Michael Yang committed
101

102
103
	var currentDigest string
	var bar *progressbar.ProgressBar
Michael Yang's avatar
Michael Yang committed
104

Michael Yang's avatar
Michael Yang committed
105
	request := api.CreateRequest{Name: args[0], Path: filename, Modelfile: string(modelfile)}
106
107
	fn := func(resp api.ProgressResponse) error {
		if resp.Digest != currentDigest && resp.Digest != "" {
Michael Yang's avatar
Michael Yang committed
108
			spinner.Stop()
109
			currentDigest = resp.Digest
110
111
112
113
114
115
			// pulling
			bar = progressbar.DefaultBytes(
				resp.Total,
				resp.Status,
			)
			bar.Set64(resp.Completed)
116
		} else if resp.Digest == currentDigest && resp.Digest != "" {
Michael Yang's avatar
Michael Yang committed
117
			bar.Set64(resp.Completed)
118
119
		} else {
			currentDigest = ""
Michael Yang's avatar
Michael Yang committed
120
			spinner.Stop()
121
122
123
			spinner = NewSpinner(resp.Status)
			go spinner.Spin(100 * time.Millisecond)
		}
124

125
126
127
128
129
130
131
		return nil
	}

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

Michael Yang's avatar
Michael Yang committed
132
133
134
	spinner.Stop()
	if spinner.description != "success" {
		return errors.New("unexpected end to create model")
Michael Yang's avatar
Michael Yang committed
135
136
	}

137
138
139
	return nil
}

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

146
147
148
	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
149
150
151
152
	var statusError api.StatusError
	switch {
	case errors.As(err, &statusError) && statusError.StatusCode == http.StatusNotFound:
		if err := PullHandler(cmd, args); err != nil {
153
			return err
Michael Yang's avatar
Michael Yang committed
154
		}
Michael Yang's avatar
Michael Yang committed
155
156
	case err != nil:
		return err
157
158
	}

Michael Yang's avatar
Michael Yang committed
159
	return RunGenerate(cmd, args)
Bruce MacDonald's avatar
Bruce MacDonald committed
160
161
}

162
func PushHandler(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
	}
167

168
169
170
171
172
	insecure, err := cmd.Flags().GetBool("insecure")
	if err != nil {
		return err
	}

173
174
175
	var currentDigest string
	var bar *progressbar.ProgressBar

176
	request := api.PushRequest{Name: args[0], Insecure: insecure}
177
	fn := func(resp api.ProgressResponse) error {
178
179
180
		if resp.Digest != currentDigest && resp.Digest != "" {
			currentDigest = resp.Digest
			bar = progressbar.DefaultBytes(
Michael Yang's avatar
Michael Yang committed
181
				resp.Total,
182
183
184
				fmt.Sprintf("pushing %s...", resp.Digest[7:19]),
			)

Michael Yang's avatar
Michael Yang committed
185
			bar.Set64(resp.Completed)
186
		} else if resp.Digest == currentDigest && resp.Digest != "" {
Michael Yang's avatar
Michael Yang committed
187
			bar.Set64(resp.Completed)
188
189
190
191
		} else {
			currentDigest = ""
			fmt.Println(resp.Status)
		}
192
193
194
195
196
197
		return nil
	}

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

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

203
204
205
	return nil
}

206
func ListHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
207
	client, err := api.ClientFromEnvironment()
208
209
210
	if err != nil {
		return err
	}
Patrick Devine's avatar
Patrick Devine committed
211
212
213
214
215
216
217
218
219

	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
220
		if len(args) == 0 || strings.HasPrefix(m.Name, args[0]) {
221
			data = append(data, []string{m.Name, m.Digest[:12], format.HumanBytes(m.Size), format.HumanTime(m.ModifiedAt, "Never")})
Michael Yang's avatar
Michael Yang committed
222
		}
Patrick Devine's avatar
Patrick Devine committed
223
224
225
	}

	table := tablewriter.NewWriter(os.Stdout)
Patrick Devine's avatar
Patrick Devine committed
226
	table.SetHeader([]string{"NAME", "ID", "SIZE", "MODIFIED"})
Patrick Devine's avatar
Patrick Devine committed
227
228
229
230
231
232
233
234
235
236
237
238
	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
}

239
func DeleteHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
240
	client, err := api.ClientFromEnvironment()
241
242
243
	if err != nil {
		return err
	}
244

245
246
247
248
249
250
	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)
251
252
253
254
	}
	return nil
}

Patrick Devine's avatar
Patrick Devine committed
255
func ShowHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
256
	client, err := api.ClientFromEnvironment()
Patrick Devine's avatar
Patrick Devine committed
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
	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 {
306
		return errors.New("only one of '--license', '--modelfile', '--parameters', '--system', or '--template' can be specified")
Patrick Devine's avatar
Patrick Devine committed
307
	} else if flagsSet == 0 {
308
		return errors.New("one of '--license', '--modelfile', '--parameters', '--system', or '--template' must be specified")
Patrick Devine's avatar
Patrick Devine committed
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
	}

	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
333
func CopyHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
334
	client, err := api.ClientFromEnvironment()
335
336
337
	if err != nil {
		return err
	}
Patrick Devine's avatar
Patrick Devine committed
338
339
340
341
342
343
344
345
346

	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
}

347
func PullHandler(cmd *cobra.Command, args []string) error {
348
349
350
351
352
353
	insecure, err := cmd.Flags().GetBool("insecure")
	if err != nil {
		return err
	}

	return pull(args[0], insecure)
354
355
}

356
func pull(model string, insecure bool) error {
Michael Yang's avatar
Michael Yang committed
357
	client, err := api.ClientFromEnvironment()
358
359
360
	if err != nil {
		return err
	}
361

362
	var currentDigest string
Bruce MacDonald's avatar
Bruce MacDonald committed
363
	var bar *progressbar.ProgressBar
Michael Yang's avatar
Michael Yang committed
364

365
	request := api.PullRequest{Name: model, Insecure: insecure}
366
367
368
	fn := func(resp api.ProgressResponse) error {
		if resp.Digest != currentDigest && resp.Digest != "" {
			currentDigest = resp.Digest
369
			bar = progressbar.DefaultBytes(
Michael Yang's avatar
Michael Yang committed
370
				resp.Total,
371
				fmt.Sprintf("pulling %s...", resp.Digest[7:19]),
372
			)
373

Michael Yang's avatar
Michael Yang committed
374
			bar.Set64(resp.Completed)
375
		} else if resp.Digest == currentDigest && resp.Digest != "" {
Michael Yang's avatar
Michael Yang committed
376
			bar.Set64(resp.Completed)
377
		} else {
378
			currentDigest = ""
379
380
			fmt.Println(resp.Status)
		}
381

382
383
		return nil
	}
384

385
386
387
	if err := client.Pull(context.Background(), &request, fn); err != nil {
		return err
	}
388
389
390
391
392

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

393
	return nil
Michael Yang's avatar
Michael Yang committed
394
395
}

396
func RunGenerate(cmd *cobra.Command, args []string) error {
397
398
399
400
	format, err := cmd.Flags().GetString("format")
	if err != nil {
		return err
	}
401

402
	prompts := args[1:]
403

404
405
406
	// prepend stdin to the prompt if provided
	if !term.IsTerminal(int(os.Stdin.Fd())) {
		in, err := io.ReadAll(os.Stdin)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
407
408
409
410
		if err != nil {
			return err
		}

411
		prompts = append([]string{string(in)}, prompts...)
Michael Yang's avatar
Michael Yang committed
412
413
	}

414
415
416
	// output is being piped
	if !term.IsTerminal(int(os.Stdout.Fd())) {
		return generate(cmd, args[0], strings.Join(prompts, " "), false, format)
Michael Yang's avatar
Michael Yang committed
417
418
	}

419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
	wordWrap := os.Getenv("TERM") == "xterm-256color"

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

	// prompts are provided via stdin or args so don't enter interactive mode
	if len(prompts) > 0 {
		return generate(cmd, args[0], strings.Join(prompts, " "), wordWrap, format)
	}

	return generateInteractive(cmd, args[0], wordWrap, format)
Michael Yang's avatar
Michael Yang committed
435
436
}

Michael Yang's avatar
Michael Yang committed
437
type generateContextKey string
Michael Yang's avatar
Michael Yang committed
438

Jeffrey Morgan's avatar
Jeffrey Morgan committed
439
func generate(cmd *cobra.Command, model, prompt string, wordWrap bool, format string) error {
Michael Yang's avatar
Michael Yang committed
440
	client, err := api.ClientFromEnvironment()
Patrick Devine's avatar
Patrick Devine committed
441
442
443
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
444

Patrick Devine's avatar
Patrick Devine committed
445
446
	spinner := NewSpinner("")
	go spinner.Spin(60 * time.Millisecond)
Michael Yang's avatar
Michael Yang committed
447

Patrick Devine's avatar
Patrick Devine committed
448
	var latest api.GenerateResponse
449

Patrick Devine's avatar
Patrick Devine committed
450
451
452
453
	generateContext, ok := cmd.Context().Value(generateContextKey("context")).([]int)
	if !ok {
		generateContext = []int{}
	}
Michael Yang's avatar
Michael Yang committed
454

455
	termWidth, _, err := term.GetSize(int(os.Stdout.Fd()))
456
	if err != nil {
457
		wordWrap = false
458
459
	}

460
461
462
463
464
465
466
467
468
469
470
471
472
	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
	}()

473
474
475
	var currentLineLength int
	var wordBuffer string

Jeffrey Morgan's avatar
Jeffrey Morgan committed
476
	request := api.GenerateRequest{Model: model, Prompt: prompt, Context: generateContext, Format: format}
Patrick Devine's avatar
Patrick Devine committed
477
478
479
480
	fn := func(response api.GenerateResponse) error {
		if !spinner.IsFinished() {
			spinner.Finish()
		}
Michael Yang's avatar
Michael Yang committed
481

Patrick Devine's avatar
Patrick Devine committed
482
		latest = response
483

484
		if wordWrap {
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
			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
509
510
		return nil
	}
511

512
	if err := client.Generate(cancelCtx, &request, fn); err != nil {
513
		if strings.Contains(err.Error(), "context canceled") && abort {
514
515
			spinner.Finish()
			return nil
516
		}
Patrick Devine's avatar
Patrick Devine committed
517
518
519
		return err
	}
	if prompt != "" {
Michael Yang's avatar
Michael Yang committed
520
521
		fmt.Println()
		fmt.Println()
Patrick Devine's avatar
Patrick Devine committed
522
	}
523

Patrick Devine's avatar
Patrick Devine committed
524
	if !latest.Done {
525
526
527
		if abort {
			return nil
		}
Patrick Devine's avatar
Patrick Devine committed
528
529
		return errors.New("unexpected end of response")
	}
530

Patrick Devine's avatar
Patrick Devine committed
531
532
533
534
	verbose, err := cmd.Flags().GetBool("verbose")
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
535

Patrick Devine's avatar
Patrick Devine committed
536
537
	if verbose {
		latest.Summary()
Michael Yang's avatar
Michael Yang committed
538
	}
Michael Yang's avatar
Michael Yang committed
539

Patrick Devine's avatar
Patrick Devine committed
540
541
542
543
	ctx := cmd.Context()
	ctx = context.WithValue(ctx, generateContextKey("context"), latest.Context)
	cmd.SetContext(ctx)

Michael Yang's avatar
Michael Yang committed
544
545
546
	return nil
}

547
func generateInteractive(cmd *cobra.Command, model string, wordWrap bool, format string) error {
Patrick Devine's avatar
Patrick Devine committed
548
	// load the model
Jeffrey Morgan's avatar
Jeffrey Morgan committed
549
	if err := generate(cmd, model, "", false, ""); err != nil {
Patrick Devine's avatar
Patrick Devine committed
550
551
552
		return err
	}

Michael Yang's avatar
Michael Yang committed
553
	usage := func() {
Patrick Devine's avatar
Patrick Devine committed
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
		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")
Jeffrey Morgan's avatar
Jeffrey Morgan committed
570
571
		fmt.Fprintln(os.Stderr, "  /set format json  Enable JSON mode")
		fmt.Fprintln(os.Stderr, "  /set noformat     Disable formatting")
Patrick Devine's avatar
Patrick Devine committed
572
573
574
575
576
577
578
579
580
581
582
583
584
		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
585
586
	}

Patrick Devine's avatar
Patrick Devine committed
587
588
589
590
591
	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
592
593
	}

Patrick Devine's avatar
Patrick Devine committed
594
	scanner, err := readline.New(prompt)
Michael Yang's avatar
Michael Yang committed
595
596
597
598
	if err != nil {
		return err
	}

599
600
601
	fmt.Print(readline.StartBracketedPaste)
	defer fmt.Printf(readline.EndBracketedPaste)

602
603
	var multiLineBuffer string

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

Michael Yang's avatar
Michael Yang committed
615
616
			continue
		case err != nil:
Michael Yang's avatar
Michael Yang committed
617
618
619
			return err
		}

Michael Yang's avatar
Michael Yang committed
620
		line = strings.TrimSpace(line)
Michael Yang's avatar
Michael Yang committed
621

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

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

Patrick Devine's avatar
Patrick Devine committed
744
		if len(line) > 0 && line[0] != '/' {
Jeffrey Morgan's avatar
Jeffrey Morgan committed
745
			if err := generate(cmd, model, line, wordWrap, format); err != nil {
Patrick Devine's avatar
Patrick Devine committed
746
747
				return err
			}
Michael Yang's avatar
Michael Yang committed
748
749
		}
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
750
751
}

752
func RunServer(cmd *cobra.Command, _ []string) error {
Michael Yang's avatar
Michael Yang committed
753
754
755
	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
756
		if ip := net.ParseIP(strings.Trim(os.Getenv("OLLAMA_HOST"), "[]")); ip != nil {
Michael Yang's avatar
Michael Yang committed
757
758
			host = ip.String()
		}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
759
	}
760

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

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

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

Jeffrey Morgan's avatar
Jeffrey Morgan committed
775
	return server.Serve(ln, origins)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
776
777
}

778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
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
800
		err = os.MkdirAll(filepath.Dir(privKeyPath), 0o755)
801
802
803
804
		if err != nil {
			return fmt.Errorf("could not create directory %w", err)
		}

805
		err = os.WriteFile(privKeyPath, pem.EncodeToMemory(privKeyBytes), 0o600)
806
807
808
809
810
811
812
813
814
815
816
		if err != nil {
			return err
		}

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

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

817
		err = os.WriteFile(pubKeyPath, pubKeyData, 0o644)
818
819
820
821
822
823
824
825
826
		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
827
func startMacApp(client *api.Client) error {
Bruce MacDonald's avatar
Bruce MacDonald committed
828
829
830
831
832
	exe, err := os.Executable()
	if err != nil {
		return err
	}
	link, err := os.Readlink(exe)
Bruce MacDonald's avatar
Bruce MacDonald committed
833
834
835
	if err != nil {
		return err
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
836
837
838
	if !strings.Contains(link, "Ollama.app") {
		return fmt.Errorf("could not find ollama app")
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
	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
			}
		}
	}
}

858
func checkServerHeartbeat(_ *cobra.Command, _ []string) error {
Michael Yang's avatar
Michael Yang committed
859
	client, err := api.ClientFromEnvironment()
860
861
862
	if err != nil {
		return err
	}
863
	if err := client.Heartbeat(context.Background()); err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
864
865
866
867
868
		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
869
				return fmt.Errorf("could not connect to ollama app, is it running?")
870
			}
Bruce MacDonald's avatar
Bruce MacDonald committed
871
		} else {
872
873
874
875
876
877
			return fmt.Errorf("could not connect to ollama server, run 'ollama serve' to start it")
		}
	}
	return nil
}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
878
879
880
881
func NewCLI() *cobra.Command {
	log.SetFlags(log.LstdFlags | log.Lshortfile)

	rootCmd := &cobra.Command{
882
883
884
885
		Use:           "ollama",
		Short:         "Large language model runner",
		SilenceUsage:  true,
		SilenceErrors: true,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
886
887
888
		CompletionOptions: cobra.CompletionOptions{
			DisableDefaultCmd: true,
		},
Michael Yang's avatar
Michael Yang committed
889
		Version: version.Version,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
890
891
892
893
	}

	cobra.EnableCommandSorting = false

894
	createCmd := &cobra.Command{
895
896
		Use:     "create MODEL",
		Short:   "Create a model from a Modelfile",
Michael Yang's avatar
Michael Yang committed
897
		Args:    cobra.ExactArgs(1),
898
899
		PreRunE: checkServerHeartbeat,
		RunE:    CreateHandler,
900
901
902
903
	}

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

Patrick Devine's avatar
Patrick Devine committed
904
905
906
	showCmd := &cobra.Command{
		Use:     "show MODEL",
		Short:   "Show information for a model",
Michael Yang's avatar
Michael Yang committed
907
		Args:    cobra.ExactArgs(1),
Patrick Devine's avatar
Patrick Devine committed
908
909
910
911
912
913
914
915
916
917
		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
918
	runCmd := &cobra.Command{
919
920
921
922
923
		Use:     "run MODEL [PROMPT]",
		Short:   "Run a model",
		Args:    cobra.MinimumNArgs(1),
		PreRunE: checkServerHeartbeat,
		RunE:    RunHandler,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
924
925
	}

926
	runCmd.Flags().Bool("verbose", false, "Show timings for response")
927
	runCmd.Flags().Bool("insecure", false, "Use an insecure registry")
928
	runCmd.Flags().Bool("nowordwrap", false, "Don't wrap words to the next line automatically")
Jeffrey Morgan's avatar
Jeffrey Morgan committed
929
	runCmd.Flags().String("format", "", "Response format (e.g. json)")
930

Jeffrey Morgan's avatar
Jeffrey Morgan committed
931
932
933
934
	serveCmd := &cobra.Command{
		Use:     "serve",
		Aliases: []string{"start"},
		Short:   "Start ollama",
Michael Yang's avatar
Michael Yang committed
935
		Args:    cobra.ExactArgs(0),
Michael Yang's avatar
Michael Yang committed
936
		RunE:    RunServer,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
937
938
	}

939
	pullCmd := &cobra.Command{
940
941
		Use:     "pull MODEL",
		Short:   "Pull a model from a registry",
Michael Yang's avatar
Michael Yang committed
942
		Args:    cobra.ExactArgs(1),
943
944
		PreRunE: checkServerHeartbeat,
		RunE:    PullHandler,
945
946
	}

947
948
	pullCmd.Flags().Bool("insecure", false, "Use an insecure registry")

949
	pushCmd := &cobra.Command{
950
951
		Use:     "push MODEL",
		Short:   "Push a model to a registry",
Michael Yang's avatar
Michael Yang committed
952
		Args:    cobra.ExactArgs(1),
953
954
		PreRunE: checkServerHeartbeat,
		RunE:    PushHandler,
955
956
	}

957
958
	pushCmd.Flags().Bool("insecure", false, "Use an insecure registry")

Patrick Devine's avatar
Patrick Devine committed
959
	listCmd := &cobra.Command{
960
		Use:     "list",
Patrick Devine's avatar
Patrick Devine committed
961
		Aliases: []string{"ls"},
962
		Short:   "List models",
963
		PreRunE: checkServerHeartbeat,
964
		RunE:    ListHandler,
965
966
	}

Patrick Devine's avatar
Patrick Devine committed
967
	copyCmd := &cobra.Command{
Michael Yang's avatar
Michael Yang committed
968
		Use:     "cp SOURCE TARGET",
969
		Short:   "Copy a model",
Michael Yang's avatar
Michael Yang committed
970
		Args:    cobra.ExactArgs(2),
971
972
		PreRunE: checkServerHeartbeat,
		RunE:    CopyHandler,
Patrick Devine's avatar
Patrick Devine committed
973
974
	}

975
	deleteCmd := &cobra.Command{
Michael Yang's avatar
Michael Yang committed
976
		Use:     "rm MODEL [MODEL...]",
977
978
979
980
		Short:   "Remove a model",
		Args:    cobra.MinimumNArgs(1),
		PreRunE: checkServerHeartbeat,
		RunE:    DeleteHandler,
Patrick Devine's avatar
Patrick Devine committed
981
982
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
983
984
	rootCmd.AddCommand(
		serveCmd,
985
		createCmd,
Patrick Devine's avatar
Patrick Devine committed
986
		showCmd,
987
		runCmd,
988
989
		pullCmd,
		pushCmd,
Patrick Devine's avatar
Patrick Devine committed
990
		listCmd,
Patrick Devine's avatar
Patrick Devine committed
991
		copyCmd,
992
		deleteCmd,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
993
994
995
996
	)

	return rootCmd
}