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"
Michael Yang's avatar
Michael Yang committed
33
	"github.com/jmorganca/ollama/progress"
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
	p := progress.NewProgress(os.Stderr)
	defer p.Stop()

	bars := make(map[string]*progress.Bar)

	status := "transferring context"
	spinner := progress.NewSpinner(status)
	p.Add(status, spinner)

Michael Yang's avatar
Michael Yang committed
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
	modelfile, err := os.ReadFile(filename)
	if err != nil {
		return err
	}

	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
87
				continue
Michael Yang's avatar
Michael Yang committed
88
89
90
91
92
93
94
95
96
97
98
99
			} 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
100
			if err = client.CreateBlob(cmd.Context(), digest, bin); err != nil {
Michael Yang's avatar
Michael Yang committed
101
102
103
				return err
			}

Michael Yang's avatar
Michael Yang committed
104
			modelfile = bytes.ReplaceAll(modelfile, []byte(c.Args), []byte("@"+digest))
Michael Yang's avatar
Michael Yang committed
105
106
		}
	}
Michael Yang's avatar
Michael Yang committed
107

108
	fn := func(resp api.ProgressResponse) error {
Michael Yang's avatar
Michael Yang committed
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
		if resp.Digest != "" {
			spinner.Stop()

			bar, ok := bars[resp.Digest]
			if !ok {
				bar = progress.NewBar(resp.Status, resp.Total, resp.Completed)
				bars[resp.Digest] = bar
				p.Add(resp.Digest, bar)
			}

			bar.Set(resp.Completed)
		} else if status != resp.Status {
			spinner.Stop()

			status = resp.Status
			spinner = progress.NewSpinner(status)
			p.Add(status, spinner)
		}

128
129
130
		return nil
	}

Michael Yang's avatar
Michael Yang committed
131
	request := api.CreateRequest{Name: args[0], Path: filename, Modelfile: string(modelfile)}
132
133
134
135
	if err := client.Create(context.Background(), &request, fn); err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
136
137
138
139
	if spinner != nil {
		spinner.Stop()
	}

140
141
142
	return nil
}

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

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

Michael Yang's avatar
Michael Yang committed
162
	return RunGenerate(cmd, args)
Bruce MacDonald's avatar
Bruce MacDonald committed
163
164
}

165
func PushHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
166
	client, err := api.ClientFromEnvironment()
167
168
169
	if err != nil {
		return err
	}
170

171
172
173
174
175
	insecure, err := cmd.Flags().GetBool("insecure")
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
176
177
178
179
180
181
182
183
	p := progress.NewProgress(os.Stderr)
	defer p.Stop()

	bars := make(map[string]*progress.Bar)

	var status string
	var spinner *progress.Spinner

184
	fn := func(resp api.ProgressResponse) error {
Michael Yang's avatar
Michael Yang committed
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
		if resp.Digest != "" {
			if spinner != nil {
				spinner.Stop()
				spinner = nil
			}

			bar, ok := bars[resp.Digest]
			if !ok {
				bar = progress.NewBar(resp.Status, resp.Total, resp.Completed)
				bars[resp.Digest] = bar
				p.Add(resp.Digest, bar)
			}

			bar.Set(resp.Completed)
		} else if status != resp.Status {
			if spinner != nil {
				spinner.Stop()
				spinner = nil
			}

			status = resp.Status
			spinner = progress.NewSpinner(status)
			p.Add(status, spinner)
		}

210
211
212
		return nil
	}

Michael Yang's avatar
Michael Yang committed
213
214
215
216
217
218
219
220
221
222
	request := api.PushRequest{Name: args[0], Insecure: insecure}
	if err := client.Push(context.Background(), &request, fn); err != nil {
		return err
	}

	if spinner != nil {
		spinner.Stop()
	}

	return nil
223
224
}

225
func ListHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
226
	client, err := api.ClientFromEnvironment()
227
228
229
	if err != nil {
		return err
	}
Patrick Devine's avatar
Patrick Devine committed
230
231
232
233
234
235
236
237
238

	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
239
		if len(args) == 0 || strings.HasPrefix(m.Name, args[0]) {
240
			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
241
		}
Patrick Devine's avatar
Patrick Devine committed
242
243
244
	}

	table := tablewriter.NewWriter(os.Stdout)
Patrick Devine's avatar
Patrick Devine committed
245
	table.SetHeader([]string{"NAME", "ID", "SIZE", "MODIFIED"})
Patrick Devine's avatar
Patrick Devine committed
246
247
248
249
250
251
252
253
254
255
256
257
	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
}

258
func DeleteHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
259
	client, err := api.ClientFromEnvironment()
260
261
262
	if err != nil {
		return err
	}
263

264
265
266
267
268
269
	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)
270
271
272
273
	}
	return nil
}

Patrick Devine's avatar
Patrick Devine committed
274
func ShowHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
275
	client, err := api.ClientFromEnvironment()
Patrick Devine's avatar
Patrick Devine committed
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
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
	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 {
325
		return errors.New("only one of '--license', '--modelfile', '--parameters', '--system', or '--template' can be specified")
Patrick Devine's avatar
Patrick Devine committed
326
	} else if flagsSet == 0 {
327
		return errors.New("one of '--license', '--modelfile', '--parameters', '--system', or '--template' must be specified")
Patrick Devine's avatar
Patrick Devine committed
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
	}

	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
352
func CopyHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
353
	client, err := api.ClientFromEnvironment()
354
355
356
	if err != nil {
		return err
	}
Patrick Devine's avatar
Patrick Devine committed
357
358
359
360
361
362
363
364
365

	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
}

366
func PullHandler(cmd *cobra.Command, args []string) error {
367
368
369
370
371
	insecure, err := cmd.Flags().GetBool("insecure")
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
372
	client, err := api.ClientFromEnvironment()
373
374
375
	if err != nil {
		return err
	}
376

Michael Yang's avatar
Michael Yang committed
377
378
379
380
381
382
383
384
	p := progress.NewProgress(os.Stderr)
	defer p.Stop()

	bars := make(map[string]*progress.Bar)

	var status string
	var spinner *progress.Spinner

385
	fn := func(resp api.ProgressResponse) error {
Michael Yang's avatar
Michael Yang committed
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
		if resp.Digest != "" {
			if spinner != nil {
				spinner.Stop()
				spinner = nil
			}

			bar, ok := bars[resp.Digest]
			if !ok {
				bar = progress.NewBar(resp.Status, resp.Total, resp.Completed)
				bars[resp.Digest] = bar
				p.Add(resp.Digest, bar)
			}

			bar.Set(resp.Completed)
		} else if status != resp.Status {
			if spinner != nil {
				spinner.Stop()
				spinner = nil
			}

			status = resp.Status
			spinner = progress.NewSpinner(status)
			p.Add(status, spinner)
		}

411
412
		return nil
	}
413

Michael Yang's avatar
Michael Yang committed
414
415
416
417
418
419
420
421
422
423
	request := api.PullRequest{Name: args[0], Insecure: insecure}
	if err := client.Pull(context.Background(), &request, fn); err != nil {
		return err
	}

	if spinner != nil {
		spinner.Stop()
	}

	return nil
Michael Yang's avatar
Michael Yang committed
424
425
}

426
func RunGenerate(cmd *cobra.Command, args []string) error {
427
428
429
430
	format, err := cmd.Flags().GetString("format")
	if err != nil {
		return err
	}
431

432
	prompts := args[1:]
433

434
435
436
	// 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
437
438
439
440
		if err != nil {
			return err
		}

441
		prompts = append([]string{string(in)}, prompts...)
Michael Yang's avatar
Michael Yang committed
442
443
	}

444
445
446
	// 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
447
448
	}

449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
	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
465
466
}

Michael Yang's avatar
Michael Yang committed
467
type generateContextKey string
Michael Yang's avatar
Michael Yang committed
468

Jeffrey Morgan's avatar
Jeffrey Morgan committed
469
func generate(cmd *cobra.Command, model, prompt string, wordWrap bool, format string) error {
Michael Yang's avatar
Michael Yang committed
470
	client, err := api.ClientFromEnvironment()
Patrick Devine's avatar
Patrick Devine committed
471
472
473
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
474

Patrick Devine's avatar
Patrick Devine committed
475
	var latest api.GenerateResponse
476

Patrick Devine's avatar
Patrick Devine committed
477
478
479
480
	generateContext, ok := cmd.Context().Value(generateContextKey("context")).([]int)
	if !ok {
		generateContext = []int{}
	}
Michael Yang's avatar
Michael Yang committed
481

482
	termWidth, _, err := term.GetSize(int(os.Stdout.Fd()))
483
	if err != nil {
484
		wordWrap = false
485
486
	}

487
488
489
490
491
492
493
494
495
496
497
498
499
	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
	}()

500
501
502
	var currentLineLength int
	var wordBuffer string

Jeffrey Morgan's avatar
Jeffrey Morgan committed
503
	request := api.GenerateRequest{Model: model, Prompt: prompt, Context: generateContext, Format: format}
Patrick Devine's avatar
Patrick Devine committed
504
505
	fn := func(response api.GenerateResponse) error {
		latest = response
506

507
		if wordWrap {
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
			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
532
533
		return nil
	}
534

535
	if err := client.Generate(cancelCtx, &request, fn); err != nil {
536
		if strings.Contains(err.Error(), "context canceled") && abort {
537
			return nil
538
		}
Patrick Devine's avatar
Patrick Devine committed
539
540
541
		return err
	}
	if prompt != "" {
Michael Yang's avatar
Michael Yang committed
542
543
		fmt.Println()
		fmt.Println()
Patrick Devine's avatar
Patrick Devine committed
544
	}
545

Patrick Devine's avatar
Patrick Devine committed
546
	if !latest.Done {
547
548
549
		if abort {
			return nil
		}
Patrick Devine's avatar
Patrick Devine committed
550
551
		return errors.New("unexpected end of response")
	}
552

Patrick Devine's avatar
Patrick Devine committed
553
554
555
556
	verbose, err := cmd.Flags().GetBool("verbose")
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
557

Patrick Devine's avatar
Patrick Devine committed
558
559
	if verbose {
		latest.Summary()
Michael Yang's avatar
Michael Yang committed
560
	}
Michael Yang's avatar
Michael Yang committed
561

Patrick Devine's avatar
Patrick Devine committed
562
563
564
565
	ctx := cmd.Context()
	ctx = context.WithValue(ctx, generateContextKey("context"), latest.Context)
	cmd.SetContext(ctx)

Michael Yang's avatar
Michael Yang committed
566
567
568
	return nil
}

569
func generateInteractive(cmd *cobra.Command, model string, wordWrap bool, format string) error {
Patrick Devine's avatar
Patrick Devine committed
570
	// load the model
Jeffrey Morgan's avatar
Jeffrey Morgan committed
571
	if err := generate(cmd, model, "", false, ""); err != nil {
Patrick Devine's avatar
Patrick Devine committed
572
573
574
		return err
	}

Michael Yang's avatar
Michael Yang committed
575
	usage := func() {
Patrick Devine's avatar
Patrick Devine committed
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
		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
592
593
		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
594
595
596
597
598
599
600
601
602
603
604
605
606
		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
607
608
	}

Patrick Devine's avatar
Patrick Devine committed
609
610
611
612
613
	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
614
615
	}

Patrick Devine's avatar
Patrick Devine committed
616
	scanner, err := readline.New(prompt)
Michael Yang's avatar
Michael Yang committed
617
618
619
620
	if err != nil {
		return err
	}

621
622
623
	fmt.Print(readline.StartBracketedPaste)
	defer fmt.Printf(readline.EndBracketedPaste)

624
625
	var multiLineBuffer string

Michael Yang's avatar
Michael Yang committed
626
627
628
629
	for {
		line, err := scanner.Readline()
		switch {
		case errors.Is(err, io.EOF):
630
			fmt.Println()
Michael Yang's avatar
Michael Yang committed
631
632
			return nil
		case errors.Is(err, readline.ErrInterrupt):
633
			if line == "" {
Patrick Devine's avatar
Patrick Devine committed
634
				fmt.Println("\nUse Ctrl-D or /bye to exit.")
635
636
			}

Michael Yang's avatar
Michael Yang committed
637
638
			continue
		case err != nil:
Michael Yang's avatar
Michael Yang committed
639
640
641
			return err
		}

Michael Yang's avatar
Michael Yang committed
642
		line = strings.TrimSpace(line)
Michael Yang's avatar
Michael Yang committed
643

Michael Yang's avatar
Michael Yang committed
644
		switch {
Patrick Devine's avatar
Patrick Devine committed
645
		case scanner.Prompt.UseAlt:
646
			if strings.HasSuffix(line, `"""`) {
Patrick Devine's avatar
Patrick Devine committed
647
				scanner.Prompt.UseAlt = false
648
				multiLineBuffer += strings.TrimSuffix(line, `"""`)
649
				line = multiLineBuffer
650
651
652
653
654
655
				multiLineBuffer = ""
			} else {
				multiLineBuffer += line + " "
				continue
			}
		case strings.HasPrefix(line, `"""`):
Patrick Devine's avatar
Patrick Devine committed
656
			scanner.Prompt.UseAlt = true
657
658
			multiLineBuffer = strings.TrimPrefix(line, `"""`) + " "
			continue
Michael Yang's avatar
Michael Yang committed
659
660
		case strings.HasPrefix(line, "/list"):
			args := strings.Fields(line)
661
			if err := ListHandler(cmd, args[1:]); err != nil {
Michael Yang's avatar
Michael Yang committed
662
663
664
665
666
667
668
669
670
671
				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()
672
				case "wordwrap":
673
					wordWrap = true
674
675
					fmt.Println("Set 'wordwrap' mode.")
				case "nowordwrap":
676
					wordWrap = false
677
					fmt.Println("Set 'nowordwrap' mode.")
Michael Yang's avatar
Michael Yang committed
678
679
				case "verbose":
					cmd.Flags().Set("verbose", "true")
680
					fmt.Println("Set 'verbose' mode.")
Michael Yang's avatar
Michael Yang committed
681
682
				case "quiet":
					cmd.Flags().Set("verbose", "false")
683
					fmt.Println("Set 'quiet' mode.")
Jeffrey Morgan's avatar
Jeffrey Morgan committed
684
685
686
687
688
689
690
691
692
693
				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.")
694
695
				default:
					fmt.Printf("Unknown command '/set %s'. Type /? for help\n", args[1])
696
697
				}
			} else {
Patrick Devine's avatar
Patrick Devine committed
698
				usageSet()
699
700
701
702
			}
		case strings.HasPrefix(line, "/show"):
			args := strings.Fields(line)
			if len(args) > 1 {
703
704
705
706
707
708
				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})
709
				if err != nil {
Patrick Devine's avatar
Patrick Devine committed
710
					fmt.Println("error: couldn't get model")
711
					return err
712
				}
Patrick Devine's avatar
Patrick Devine committed
713

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

Patrick Devine's avatar
Patrick Devine committed
766
		if len(line) > 0 && line[0] != '/' {
Jeffrey Morgan's avatar
Jeffrey Morgan committed
767
			if err := generate(cmd, model, line, wordWrap, format); err != nil {
Patrick Devine's avatar
Patrick Devine committed
768
769
				return err
			}
Michael Yang's avatar
Michael Yang committed
770
771
		}
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
772
773
}

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

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

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

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

Jeffrey Morgan's avatar
Jeffrey Morgan committed
797
	return server.Serve(ln, origins)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
798
799
}

800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
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
822
		err = os.MkdirAll(filepath.Dir(privKeyPath), 0o755)
823
824
825
826
		if err != nil {
			return fmt.Errorf("could not create directory %w", err)
		}

827
		err = os.WriteFile(privKeyPath, pem.EncodeToMemory(privKeyBytes), 0o600)
828
829
830
831
832
833
834
835
836
837
838
		if err != nil {
			return err
		}

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

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

839
		err = os.WriteFile(pubKeyPath, pubKeyData, 0o644)
840
841
842
843
844
845
846
847
848
		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
849
func startMacApp(client *api.Client) error {
Bruce MacDonald's avatar
Bruce MacDonald committed
850
851
852
853
854
	exe, err := os.Executable()
	if err != nil {
		return err
	}
	link, err := os.Readlink(exe)
Bruce MacDonald's avatar
Bruce MacDonald committed
855
856
857
	if err != nil {
		return err
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
858
859
860
	if !strings.Contains(link, "Ollama.app") {
		return fmt.Errorf("could not find ollama app")
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
	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
			}
		}
	}
}

880
func checkServerHeartbeat(_ *cobra.Command, _ []string) error {
Michael Yang's avatar
Michael Yang committed
881
	client, err := api.ClientFromEnvironment()
882
883
884
	if err != nil {
		return err
	}
885
	if err := client.Heartbeat(context.Background()); err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
886
887
888
889
890
		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
891
				return fmt.Errorf("could not connect to ollama app, is it running?")
892
			}
Bruce MacDonald's avatar
Bruce MacDonald committed
893
		} else {
894
895
896
897
898
899
			return fmt.Errorf("could not connect to ollama server, run 'ollama serve' to start it")
		}
	}
	return nil
}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
900
901
902
903
func NewCLI() *cobra.Command {
	log.SetFlags(log.LstdFlags | log.Lshortfile)

	rootCmd := &cobra.Command{
904
905
906
907
		Use:           "ollama",
		Short:         "Large language model runner",
		SilenceUsage:  true,
		SilenceErrors: true,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
908
909
910
		CompletionOptions: cobra.CompletionOptions{
			DisableDefaultCmd: true,
		},
Michael Yang's avatar
Michael Yang committed
911
		Version: version.Version,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
912
913
914
915
	}

	cobra.EnableCommandSorting = false

916
	createCmd := &cobra.Command{
917
918
		Use:     "create MODEL",
		Short:   "Create a model from a Modelfile",
Michael Yang's avatar
Michael Yang committed
919
		Args:    cobra.ExactArgs(1),
920
921
		PreRunE: checkServerHeartbeat,
		RunE:    CreateHandler,
922
923
924
925
	}

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

Patrick Devine's avatar
Patrick Devine committed
926
927
928
	showCmd := &cobra.Command{
		Use:     "show MODEL",
		Short:   "Show information for a model",
Michael Yang's avatar
Michael Yang committed
929
		Args:    cobra.ExactArgs(1),
Patrick Devine's avatar
Patrick Devine committed
930
931
932
933
934
935
936
937
938
939
		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
940
	runCmd := &cobra.Command{
941
942
943
944
945
		Use:     "run MODEL [PROMPT]",
		Short:   "Run a model",
		Args:    cobra.MinimumNArgs(1),
		PreRunE: checkServerHeartbeat,
		RunE:    RunHandler,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
946
947
	}

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

Jeffrey Morgan's avatar
Jeffrey Morgan committed
953
954
955
956
	serveCmd := &cobra.Command{
		Use:     "serve",
		Aliases: []string{"start"},
		Short:   "Start ollama",
Michael Yang's avatar
Michael Yang committed
957
		Args:    cobra.ExactArgs(0),
Michael Yang's avatar
Michael Yang committed
958
		RunE:    RunServer,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
959
960
	}

961
	pullCmd := &cobra.Command{
962
963
		Use:     "pull MODEL",
		Short:   "Pull a model from a registry",
Michael Yang's avatar
Michael Yang committed
964
		Args:    cobra.ExactArgs(1),
965
966
		PreRunE: checkServerHeartbeat,
		RunE:    PullHandler,
967
968
	}

969
970
	pullCmd.Flags().Bool("insecure", false, "Use an insecure registry")

971
	pushCmd := &cobra.Command{
972
973
		Use:     "push MODEL",
		Short:   "Push a model to a registry",
Michael Yang's avatar
Michael Yang committed
974
		Args:    cobra.ExactArgs(1),
975
976
		PreRunE: checkServerHeartbeat,
		RunE:    PushHandler,
977
978
	}

979
980
	pushCmd.Flags().Bool("insecure", false, "Use an insecure registry")

Patrick Devine's avatar
Patrick Devine committed
981
	listCmd := &cobra.Command{
982
		Use:     "list",
Patrick Devine's avatar
Patrick Devine committed
983
		Aliases: []string{"ls"},
984
		Short:   "List models",
985
		PreRunE: checkServerHeartbeat,
986
		RunE:    ListHandler,
987
988
	}

Patrick Devine's avatar
Patrick Devine committed
989
	copyCmd := &cobra.Command{
Michael Yang's avatar
Michael Yang committed
990
		Use:     "cp SOURCE TARGET",
991
		Short:   "Copy a model",
Michael Yang's avatar
Michael Yang committed
992
		Args:    cobra.ExactArgs(2),
993
994
		PreRunE: checkServerHeartbeat,
		RunE:    CopyHandler,
Patrick Devine's avatar
Patrick Devine committed
995
996
	}

997
	deleteCmd := &cobra.Command{
Michael Yang's avatar
Michael Yang committed
998
		Use:     "rm MODEL [MODEL...]",
999
1000
1001
1002
		Short:   "Remove a model",
		Args:    cobra.MinimumNArgs(1),
		PreRunE: checkServerHeartbeat,
		RunE:    DeleteHandler,
Patrick Devine's avatar
Patrick Devine committed
1003
1004
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
1005
1006
	rootCmd.AddCommand(
		serveCmd,
1007
		createCmd,
Patrick Devine's avatar
Patrick Devine committed
1008
		showCmd,
1009
		runCmd,
1010
1011
		pullCmd,
		pushCmd,
Patrick Devine's avatar
Patrick Devine committed
1012
		listCmd,
Patrick Devine's avatar
Patrick Devine committed
1013
		copyCmd,
1014
		deleteCmd,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
1015
1016
1017
1018
	)

	return rootCmd
}