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

import (
4
	"archive/zip"
Michael Yang's avatar
Michael Yang committed
5
	"bytes"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
6
	"context"
7
8
	"crypto/ed25519"
	"crypto/rand"
Michael Yang's avatar
Michael Yang committed
9
	"crypto/sha256"
10
	"encoding/pem"
Michael Yang's avatar
Michael Yang committed
11
	"errors"
Bruce MacDonald's avatar
Bruce MacDonald committed
12
	"fmt"
Michael Yang's avatar
Michael Yang committed
13
	"io"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
14
15
	"log"
	"net"
16
	"net/http"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
17
	"os"
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

25
26
	"github.com/containerd/console"

Patrick Devine's avatar
Patrick Devine committed
27
	"github.com/olekukonko/tablewriter"
Michael Yang's avatar
Michael Yang committed
28
	"github.com/spf13/cobra"
29
	"golang.org/x/crypto/ssh"
30
	"golang.org/x/exp/slices"
31
	"golang.org/x/term"
Michael Yang's avatar
Michael Yang committed
32

33
34
35
36
37
38
	"github.com/ollama/ollama/api"
	"github.com/ollama/ollama/format"
	"github.com/ollama/ollama/parser"
	"github.com/ollama/ollama/progress"
	"github.com/ollama/ollama/server"
	"github.com/ollama/ollama/version"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
39
40
)

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

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

Michael Yang's avatar
Michael Yang committed
53
54
55
56
57
	p := progress.NewProgress(os.Stderr)
	defer p.Stop()

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

Michael Yang's avatar
Michael Yang committed
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
	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
	}

73
74
	status := "transferring model data"
	spinner := progress.NewSpinner(status)
75
76
	p.Add(status, spinner)

Michael Yang's avatar
Michael Yang committed
77
78
79
80
81
82
83
84
85
86
	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:])
			}

87
88
89
90
			if !filepath.IsAbs(path) {
				path = filepath.Join(filepath.Dir(filename), path)
			}

91
			fi, err := os.Stat(path)
Michael Yang's avatar
Michael Yang committed
92
			if errors.Is(err, os.ErrNotExist) && c.Name == "model" {
Michael Yang's avatar
Michael Yang committed
93
				continue
Michael Yang's avatar
Michael Yang committed
94
95
96
97
			} else if err != nil {
				return err
			}

98
99
100
101
102
103
104
105
106
107
			// TODO make this work w/ adapters
			if fi.IsDir() {
				tf, err := os.CreateTemp("", "ollama-tf")
				if err != nil {
					return err
				}
				defer os.RemoveAll(tf.Name())

				zf := zip.NewWriter(tf)

108
109
110
				files := []string{}

				tfiles, err := filepath.Glob(filepath.Join(path, "pytorch_model-*.bin"))
111
112
				if err != nil {
					return err
113
114
115
116
117
				} else if len(tfiles) == 0 {
					tfiles, err = filepath.Glob(filepath.Join(path, "model-*.safetensors"))
					if err != nil {
						return err
					}
118
119
				}

120
121
				files = append(files, tfiles...)

122
				if len(files) == 0 {
123
					return fmt.Errorf("no models were found in '%s'", path)
124
125
				}

126
				// add the safetensor/torch config file + tokenizer
127
				files = append(files, filepath.Join(path, "config.json"))
128
				files = append(files, filepath.Join(path, "params.json"))
129
130
131
132
133
				files = append(files, filepath.Join(path, "added_tokens.json"))
				files = append(files, filepath.Join(path, "tokenizer.model"))

				for _, fn := range files {
					f, err := os.Open(fn)
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149

					// just skip whatever files aren't there
					if os.IsNotExist(err) {
						if strings.HasSuffix(fn, "tokenizer.model") {
							// try the parent dir before giving up
							parentDir := filepath.Dir(path)
							newFn := filepath.Join(parentDir, "tokenizer.model")
							f, err = os.Open(newFn)
							if os.IsNotExist(err) {
								continue
							} else if err != nil {
								return err
							}
						} else {
							continue
						}
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
					} else if err != nil {
						return err
					}

					fi, err := f.Stat()
					if err != nil {
						return err
					}

					h, err := zip.FileInfoHeader(fi)
					if err != nil {
						return err
					}

					h.Name = filepath.Base(fn)
					h.Method = zip.Store

					w, err := zf.CreateHeader(h)
					if err != nil {
						return err
					}

					_, err = io.Copy(w, f)
					if err != nil {
						return err
					}

				}

				if err := zf.Close(); err != nil {
					return err
				}

				if err := tf.Close(); err != nil {
					return err
				}
				path = tf.Name()
Michael Yang's avatar
Michael Yang committed
187
188
			}

189
190
			digest, err := createBlob(cmd, client, path)
			if err != nil {
Michael Yang's avatar
Michael Yang committed
191
192
193
				return err
			}

Michael Yang's avatar
Michael Yang committed
194
			modelfile = bytes.ReplaceAll(modelfile, []byte(c.Args), []byte("@"+digest))
Michael Yang's avatar
Michael Yang committed
195
196
		}
	}
Michael Yang's avatar
Michael Yang committed
197

198
	fn := func(resp api.ProgressResponse) error {
Michael Yang's avatar
Michael Yang committed
199
200
201
202
203
		if resp.Digest != "" {
			spinner.Stop()

			bar, ok := bars[resp.Digest]
			if !ok {
204
				bar = progress.NewBar(fmt.Sprintf("pulling %s...", resp.Digest[7:19]), resp.Total, resp.Completed)
Michael Yang's avatar
Michael Yang committed
205
206
207
208
209
210
211
212
213
214
215
216
217
				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)
		}

218
219
220
		return nil
	}

Michael Yang's avatar
Michael Yang committed
221
222
223
	quantization, _ := cmd.Flags().GetString("quantization")

	request := api.CreateRequest{Name: args[0], Modelfile: string(modelfile), Quantization: quantization}
Michael Yang's avatar
Michael Yang committed
224
	if err := client.Create(cmd.Context(), &request, fn); err != nil {
225
226
227
228
229
230
		return err
	}

	return nil
}

231
232
233
234
235
236
237
238
239
240
241
func createBlob(cmd *cobra.Command, client *api.Client, path string) (string, error) {
	bin, err := os.Open(path)
	if err != nil {
		return "", err
	}
	defer bin.Close()

	hash := sha256.New()
	if _, err := io.Copy(hash, bin); err != nil {
		return "", err
	}
242
243
244
245

	if _, err := bin.Seek(0, io.SeekStart); err != nil {
		return "", err
	}
246
247
248
249
250
251
252
253

	digest := fmt.Sprintf("sha256:%x", hash.Sum(nil))
	if err = client.CreateBlob(cmd.Context(), digest, bin); err != nil {
		return "", err
	}
	return digest, nil
}

254
func RunHandler(cmd *cobra.Command, args []string) error {
255
256
257
258
259
260
261
262
	if os.Getenv("OLLAMA_MODELS") != "" {
		return errors.New("OLLAMA_MODELS must only be set for 'ollama serve'")
	}

	if err := checkServerHeartbeat(cmd, args); err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
263
	client, err := api.ClientFromEnvironment()
264
265
266
267
	if err != nil {
		return err
	}

268
	name := args[0]
269

270
	// check if the model exists on the server
271
	show, err := client.Show(cmd.Context(), &api.ShowRequest{Name: name})
Michael Yang's avatar
Michael Yang committed
272
273
274
	var statusError api.StatusError
	switch {
	case errors.As(err, &statusError) && statusError.StatusCode == http.StatusNotFound:
275
		if err := PullHandler(cmd, []string{name}); err != nil {
276
			return err
Michael Yang's avatar
Michael Yang committed
277
		}
278
279
280
281
282

		show, err = client.Show(cmd.Context(), &api.ShowRequest{Name: name})
		if err != nil {
			return err
		}
Michael Yang's avatar
Michael Yang committed
283
284
	case err != nil:
		return err
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
325
326
327
328
329
330
	interactive := true

	opts := runOptions{
		Model:       args[0],
		WordWrap:    os.Getenv("TERM") == "xterm-256color",
		Options:     map[string]interface{}{},
		MultiModal:  slices.Contains(show.Details.Families, "clip"),
		ParentModel: show.Details.ParentModel,
	}

	format, err := cmd.Flags().GetString("format")
	if err != nil {
		return err
	}
	opts.Format = format

	prompts := args[1:]
	// prepend stdin to the prompt if provided
	if !term.IsTerminal(int(os.Stdin.Fd())) {
		in, err := io.ReadAll(os.Stdin)
		if err != nil {
			return err
		}

		prompts = append([]string{string(in)}, prompts...)
		opts.WordWrap = false
		interactive = false
	}
	opts.Prompt = strings.Join(prompts, " ")
	if len(prompts) > 0 {
		interactive = false
	}

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

	if !interactive {
		return generate(cmd, opts)
	}

	return generateInteractive(cmd, opts)
Bruce MacDonald's avatar
Bruce MacDonald committed
331
332
}

333
func PushHandler(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
	}
338

339
340
341
342
343
	insecure, err := cmd.Flags().GetBool("insecure")
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
344
345
346
347
	p := progress.NewProgress(os.Stderr)
	defer p.Stop()

	bars := make(map[string]*progress.Bar)
348
349
	var status string
	var spinner *progress.Spinner
Michael Yang's avatar
Michael Yang committed
350

351
	fn := func(resp api.ProgressResponse) error {
Michael Yang's avatar
Michael Yang committed
352
		if resp.Digest != "" {
353
354
355
			if spinner != nil {
				spinner.Stop()
			}
Michael Yang's avatar
Michael Yang committed
356
357
358

			bar, ok := bars[resp.Digest]
			if !ok {
359
				bar = progress.NewBar(fmt.Sprintf("pushing %s...", resp.Digest[7:19]), resp.Total, resp.Completed)
Michael Yang's avatar
Michael Yang committed
360
361
362
363
364
365
				bars[resp.Digest] = bar
				p.Add(resp.Digest, bar)
			}

			bar.Set(resp.Completed)
		} else if status != resp.Status {
366
367
368
			if spinner != nil {
				spinner.Stop()
			}
Michael Yang's avatar
Michael Yang committed
369
370
371
372
373
374

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

375
376
377
		return nil
	}

Michael Yang's avatar
Michael Yang committed
378
	request := api.PushRequest{Name: args[0], Insecure: insecure}
Michael Yang's avatar
Michael Yang committed
379
	if err := client.Push(cmd.Context(), &request, fn); err != nil {
Michael Yang's avatar
Michael Yang committed
380
381
382
		return err
	}

383
	spinner.Stop()
Michael Yang's avatar
Michael Yang committed
384
	return nil
385
386
}

387
func ListHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
388
	client, err := api.ClientFromEnvironment()
389
390
391
	if err != nil {
		return err
	}
Patrick Devine's avatar
Patrick Devine committed
392

Michael Yang's avatar
Michael Yang committed
393
	models, err := client.List(cmd.Context())
Patrick Devine's avatar
Patrick Devine committed
394
395
396
397
398
399
400
	if err != nil {
		return err
	}

	var data [][]string

	for _, m := range models.Models {
Michael Yang's avatar
Michael Yang committed
401
		if len(args) == 0 || strings.HasPrefix(m.Name, args[0]) {
402
			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
403
		}
Patrick Devine's avatar
Patrick Devine committed
404
405
406
	}

	table := tablewriter.NewWriter(os.Stdout)
Patrick Devine's avatar
Patrick Devine committed
407
	table.SetHeader([]string{"NAME", "ID", "SIZE", "MODIFIED"})
Patrick Devine's avatar
Patrick Devine committed
408
409
410
411
412
413
414
415
416
417
418
419
	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
}

420
func DeleteHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
421
	client, err := api.ClientFromEnvironment()
422
423
424
	if err != nil {
		return err
	}
425

426
427
	for _, name := range args {
		req := api.DeleteRequest{Name: name}
Michael Yang's avatar
Michael Yang committed
428
		if err := client.Delete(cmd.Context(), &req); err != nil {
429
430
431
			return err
		}
		fmt.Printf("deleted '%s'\n", name)
432
433
434
435
	}
	return nil
}

Patrick Devine's avatar
Patrick Devine committed
436
func ShowHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
437
	client, err := api.ClientFromEnvironment()
Patrick Devine's avatar
Patrick Devine committed
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
	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 {
487
		return errors.New("only one of '--license', '--modelfile', '--parameters', '--system', or '--template' can be specified")
Patrick Devine's avatar
Patrick Devine committed
488
	} else if flagsSet == 0 {
489
		return errors.New("one of '--license', '--modelfile', '--parameters', '--system', or '--template' must be specified")
Patrick Devine's avatar
Patrick Devine committed
490
491
	}

492
	req := api.ShowRequest{Name: args[0]}
Michael Yang's avatar
Michael Yang committed
493
	resp, err := client.Show(cmd.Context(), &req)
Patrick Devine's avatar
Patrick Devine committed
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
	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
514
func CopyHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
515
	client, err := api.ClientFromEnvironment()
516
517
518
	if err != nil {
		return err
	}
Patrick Devine's avatar
Patrick Devine committed
519
520

	req := api.CopyRequest{Source: args[0], Destination: args[1]}
Michael Yang's avatar
Michael Yang committed
521
	if err := client.Copy(cmd.Context(), &req); err != nil {
Patrick Devine's avatar
Patrick Devine committed
522
523
524
525
526
527
		return err
	}
	fmt.Printf("copied '%s' to '%s'\n", args[0], args[1])
	return nil
}

528
func PullHandler(cmd *cobra.Command, args []string) error {
529
530
531
532
533
	insecure, err := cmd.Flags().GetBool("insecure")
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
534
	client, err := api.ClientFromEnvironment()
535
536
537
	if err != nil {
		return err
	}
538

Michael Yang's avatar
Michael Yang committed
539
540
541
542
543
	p := progress.NewProgress(os.Stderr)
	defer p.Stop()

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

544
545
	var status string
	var spinner *progress.Spinner
Michael Yang's avatar
Michael Yang committed
546

547
	fn := func(resp api.ProgressResponse) error {
Michael Yang's avatar
Michael Yang committed
548
		if resp.Digest != "" {
549
550
551
			if spinner != nil {
				spinner.Stop()
			}
Michael Yang's avatar
Michael Yang committed
552
553
554

			bar, ok := bars[resp.Digest]
			if !ok {
555
				bar = progress.NewBar(fmt.Sprintf("pulling %s...", resp.Digest[7:19]), resp.Total, resp.Completed)
Michael Yang's avatar
Michael Yang committed
556
557
558
559
560
561
				bars[resp.Digest] = bar
				p.Add(resp.Digest, bar)
			}

			bar.Set(resp.Completed)
		} else if status != resp.Status {
562
563
564
			if spinner != nil {
				spinner.Stop()
			}
Michael Yang's avatar
Michael Yang committed
565
566
567
568
569
570

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

571
572
		return nil
	}
573

Michael Yang's avatar
Michael Yang committed
574
	request := api.PullRequest{Name: args[0], Insecure: insecure}
Michael Yang's avatar
Michael Yang committed
575
	if err := client.Pull(cmd.Context(), &request, fn); err != nil {
Michael Yang's avatar
Michael Yang committed
576
577
578
579
		return err
	}

	return nil
Michael Yang's avatar
Michael Yang committed
580
581
}

582
583
type generateContextKey string

584
type runOptions struct {
585
586
587
588
589
590
591
592
593
594
595
	Model       string
	ParentModel string
	Prompt      string
	Messages    []api.Message
	WordWrap    bool
	Format      string
	System      string
	Template    string
	Images      []api.ImageData
	Options     map[string]interface{}
	MultiModal  bool
596
597
}

598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
type displayResponseState struct {
	lineLength int
	wordBuffer string
}

func displayResponse(content string, wordWrap bool, state *displayResponseState) {
	termWidth, _, _ := term.GetSize(int(os.Stdout.Fd()))
	if wordWrap && termWidth >= 10 {
		for _, ch := range content {
			if state.lineLength+1 > termWidth-5 {
				if len(state.wordBuffer) > termWidth-10 {
					fmt.Printf("%s%c", state.wordBuffer, ch)
					state.wordBuffer = ""
					state.lineLength = 0
					continue
				}

				// backtrack the length of the last word and clear to the end of the line
				fmt.Printf("\x1b[%dD\x1b[K\n", len(state.wordBuffer))
				fmt.Printf("%s%c", state.wordBuffer, ch)
				state.lineLength = len(state.wordBuffer) + 1
			} else {
				fmt.Print(string(ch))
				state.lineLength += 1

				switch ch {
				case ' ':
					state.wordBuffer = ""
				case '\n':
					state.lineLength = 0
				default:
					state.wordBuffer += string(ch)
				}
			}
		}
	} else {
		fmt.Printf("%s%s", state.wordBuffer, content)
		if len(state.wordBuffer) > 0 {
			state.wordBuffer = ""
		}
	}
}

func chat(cmd *cobra.Command, opts runOptions) (*api.Message, error) {
	client, err := api.ClientFromEnvironment()
	if err != nil {
		return nil, err
	}

	p := progress.NewProgress(os.Stderr)
	defer p.StopAndClear()

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

	cancelCtx, cancel := context.WithCancel(cmd.Context())
	defer cancel()

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT)

	go func() {
		<-sigChan
		cancel()
	}()

	var state *displayResponseState = &displayResponseState{}
	var latest api.ChatResponse
	var fullResponse strings.Builder
	var role string

	fn := func(response api.ChatResponse) error {
		p.StopAndClear()

		latest = response

		role = response.Message.Role
		content := response.Message.Content
		fullResponse.WriteString(content)

		displayResponse(content, opts.WordWrap, state)

		return nil
	}

	req := &api.ChatRequest{
		Model:    opts.Model,
		Messages: opts.Messages,
		Format:   opts.Format,
		Options:  opts.Options,
	}

	if err := client.Chat(cancelCtx, req, fn); err != nil {
		if errors.Is(err, context.Canceled) {
			return nil, nil
		}
		return nil, err
	}

	if len(opts.Messages) > 0 {
		fmt.Println()
		fmt.Println()
	}

	verbose, err := cmd.Flags().GetBool("verbose")
	if err != nil {
		return nil, err
	}

	if verbose {
		latest.Summary()
	}

	return &api.Message{Role: role, Content: fullResponse.String()}, nil
}

func generate(cmd *cobra.Command, opts runOptions) error {
Michael Yang's avatar
Michael Yang committed
715
	client, err := api.ClientFromEnvironment()
Patrick Devine's avatar
Patrick Devine committed
716
	if err != nil {
717
		return err
Patrick Devine's avatar
Patrick Devine committed
718
	}
Michael Yang's avatar
Michael Yang committed
719

Michael Yang's avatar
Michael Yang committed
720
	p := progress.NewProgress(os.Stderr)
721
	defer p.StopAndClear()
722

Michael Yang's avatar
Michael Yang committed
723
724
725
	spinner := progress.NewSpinner("")
	p.Add("", spinner)

726
727
728
729
730
731
732
	var latest api.GenerateResponse

	generateContext, ok := cmd.Context().Value(generateContextKey("context")).([]int)
	if !ok {
		generateContext = []int{}
	}

Michael Yang's avatar
Michael Yang committed
733
	ctx, cancel := context.WithCancel(cmd.Context())
734
735
736
737
738
739
740
741
742
743
	defer cancel()

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT)

	go func() {
		<-sigChan
		cancel()
	}()

744
	var state *displayResponseState = &displayResponseState{}
745

746
	fn := func(response api.GenerateResponse) error {
Michael Yang's avatar
Michael Yang committed
747
		p.StopAndClear()
748

Patrick Devine's avatar
Patrick Devine committed
749
		latest = response
750
		content := response.Response
751

752
		displayResponse(content, opts.WordWrap, state)
753

Patrick Devine's avatar
Patrick Devine committed
754
755
		return nil
	}
756

757
758
759
760
761
762
763
	if opts.MultiModal {
		opts.Prompt, opts.Images, err = extractFileData(opts.Prompt)
		if err != nil {
			return err
		}
	}

Michael Yang's avatar
Michael Yang committed
764
765
766
767
	request := api.GenerateRequest{
		Model:    opts.Model,
		Prompt:   opts.Prompt,
		Context:  generateContext,
768
		Images:   opts.Images,
Michael Yang's avatar
Michael Yang committed
769
770
771
772
773
774
775
		Format:   opts.Format,
		System:   opts.System,
		Template: opts.Template,
		Options:  opts.Options,
	}

	if err := client.Generate(ctx, &request, fn); err != nil {
776
		if errors.Is(err, context.Canceled) {
777
			return nil
778
		}
779
		return err
Patrick Devine's avatar
Patrick Devine committed
780
	}
781

782
	if opts.Prompt != "" {
Michael Yang's avatar
Michael Yang committed
783
784
		fmt.Println()
		fmt.Println()
Patrick Devine's avatar
Patrick Devine committed
785
	}
786

787
788
789
790
	if !latest.Done {
		return nil
	}

Patrick Devine's avatar
Patrick Devine committed
791
792
	verbose, err := cmd.Flags().GetBool("verbose")
	if err != nil {
793
		return err
Patrick Devine's avatar
Patrick Devine committed
794
	}
Michael Yang's avatar
Michael Yang committed
795

Patrick Devine's avatar
Patrick Devine committed
796
797
	if verbose {
		latest.Summary()
Michael Yang's avatar
Michael Yang committed
798
	}
Michael Yang's avatar
Michael Yang committed
799

Patrick Devine's avatar
Patrick Devine committed
800
801
802
	ctx = context.WithValue(cmd.Context(), generateContextKey("context"), latest.Context)
	cmd.SetContext(ctx)

803
	return nil
Michael Yang's avatar
Michael Yang committed
804
805
}

806
func RunServer(cmd *cobra.Command, _ []string) error {
807
	host, port, err := net.SplitHostPort(strings.Trim(os.Getenv("OLLAMA_HOST"), "\"'"))
Michael Yang's avatar
Michael Yang committed
808
809
	if err != nil {
		host, port = "127.0.0.1", "11434"
Michael Yang's avatar
Michael Yang committed
810
		if ip := net.ParseIP(strings.Trim(os.Getenv("OLLAMA_HOST"), "[]")); ip != nil {
Michael Yang's avatar
Michael Yang committed
811
812
			host = ip.String()
		}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
813
	}
814

Michael Yang's avatar
Michael Yang committed
815
	if err := initializeKeypair(); err != nil {
816
817
818
		return err
	}

Michael Yang's avatar
Michael Yang committed
819
	ln, err := net.Listen("tcp", net.JoinHostPort(host, port))
820
821
822
	if err != nil {
		return err
	}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
823

824
	return server.Serve(ln)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
825
826
}

827
828
829
830
831
832
833
834
835
836
837
838
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)
Michael Yang's avatar
Michael Yang committed
839
		cryptoPublicKey, cryptoPrivateKey, err := ed25519.GenerateKey(rand.Reader)
840
841
842
843
		if err != nil {
			return err
		}

Michael Yang's avatar
Michael Yang committed
844
		privateKeyBytes, err := ssh.MarshalPrivateKey(cryptoPrivateKey, "")
845
846
847
848
		if err != nil {
			return err
		}

Michael Yang's avatar
Michael Yang committed
849
		if err := os.MkdirAll(filepath.Dir(privKeyPath), 0o755); err != nil {
850
851
852
			return fmt.Errorf("could not create directory %w", err)
		}

Michael Yang's avatar
Michael Yang committed
853
		if err := os.WriteFile(privKeyPath, pem.EncodeToMemory(privateKeyBytes), 0o600); err != nil {
854
855
856
			return err
		}

Michael Yang's avatar
Michael Yang committed
857
		sshPublicKey, err := ssh.NewPublicKey(cryptoPublicKey)
858
859
860
861
		if err != nil {
			return err
		}

Michael Yang's avatar
Michael Yang committed
862
		publicKeyBytes := ssh.MarshalAuthorizedKey(sshPublicKey)
863

Michael Yang's avatar
Michael Yang committed
864
		if err := os.WriteFile(pubKeyPath, publicKeyBytes, 0o644); err != nil {
865
866
867
			return err
		}

Michael Yang's avatar
Michael Yang committed
868
		fmt.Printf("Your new public key is: \n\n%s\n", publicKeyBytes)
869
870
871
872
	}
	return nil
}

873
874
//nolint:unused
func waitForServer(ctx context.Context, client *api.Client) error {
Bruce MacDonald's avatar
Bruce MacDonald committed
875
876
877
878
879
880
881
882
	// 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:
Michael Yang's avatar
Michael Yang committed
883
			if err := client.Heartbeat(ctx); err == nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
884
885
886
887
				return nil // server has started
			}
		}
	}
888

Bruce MacDonald's avatar
Bruce MacDonald committed
889
890
}

Michael Yang's avatar
Michael Yang committed
891
func checkServerHeartbeat(cmd *cobra.Command, _ []string) error {
Michael Yang's avatar
Michael Yang committed
892
	client, err := api.ClientFromEnvironment()
893
894
895
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
896
	if err := client.Heartbeat(cmd.Context()); err != nil {
897
		if !strings.Contains(err.Error(), " refused") {
Bruce MacDonald's avatar
Bruce MacDonald committed
898
899
			return err
		}
900
901
		if err := startApp(cmd.Context(), client); err != nil {
			return fmt.Errorf("could not connect to ollama app, is it running?")
902
903
904
905
906
		}
	}
	return nil
}

Michael Yang's avatar
Michael Yang committed
907
908
909
910
911
912
913
914
func versionHandler(cmd *cobra.Command, _ []string) {
	client, err := api.ClientFromEnvironment()
	if err != nil {
		return
	}

	serverVersion, err := client.Version(cmd.Context())
	if err != nil {
Michael Yang's avatar
Michael Yang committed
915
916
917
918
919
		fmt.Println("Warning: could not connect to a running Ollama instance")
	}

	if serverVersion != "" {
		fmt.Printf("ollama version is %s\n", serverVersion)
Michael Yang's avatar
Michael Yang committed
920
921
	}

922
	if serverVersion != version.Version {
Michael Yang's avatar
Michael Yang committed
923
		fmt.Printf("Warning: client version is %s\n", version.Version)
924
	}
Michael Yang's avatar
Michael Yang committed
925
926
}

927
928
929
930
931
932
933
934
func appendHostEnvDocs(cmd *cobra.Command) {
	const hostEnvDocs = `
Environment Variables:
      OLLAMA_HOST        The host:port or base URL of the Ollama server (e.g. http://localhost:11434)
`
	cmd.SetUsageTemplate(cmd.UsageTemplate() + hostEnvDocs)
}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
935
936
func NewCLI() *cobra.Command {
	log.SetFlags(log.LstdFlags | log.Lshortfile)
Michael Yang's avatar
Michael Yang committed
937
	cobra.EnableCommandSorting = false
Jeffrey Morgan's avatar
Jeffrey Morgan committed
938

939
	if runtime.GOOS == "windows" {
940
		console.ConsoleFromFile(os.Stdin) //nolint:errcheck
941
942
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
943
	rootCmd := &cobra.Command{
944
945
946
947
		Use:           "ollama",
		Short:         "Large language model runner",
		SilenceUsage:  true,
		SilenceErrors: true,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
948
949
950
		CompletionOptions: cobra.CompletionOptions{
			DisableDefaultCmd: true,
		},
Michael Yang's avatar
Michael Yang committed
951
952
953
954
955
956
957
958
		Run: func(cmd *cobra.Command, args []string) {
			if version, _ := cmd.Flags().GetBool("version"); version {
				versionHandler(cmd, args)
				return
			}

			cmd.Print(cmd.UsageString())
		},
Jeffrey Morgan's avatar
Jeffrey Morgan committed
959
960
	}

Michael Yang's avatar
Michael Yang committed
961
	rootCmd.Flags().BoolP("version", "v", false, "Show version information")
Jeffrey Morgan's avatar
Jeffrey Morgan committed
962

963
	createCmd := &cobra.Command{
964
965
		Use:     "create MODEL",
		Short:   "Create a model from a Modelfile",
Michael Yang's avatar
Michael Yang committed
966
		Args:    cobra.ExactArgs(1),
967
968
		PreRunE: checkServerHeartbeat,
		RunE:    CreateHandler,
969
970
971
	}

	createCmd.Flags().StringP("file", "f", "Modelfile", "Name of the Modelfile (default \"Modelfile\")")
Michael Yang's avatar
Michael Yang committed
972
	createCmd.Flags().StringP("quantization", "q", "", "Quantization level.")
973

Patrick Devine's avatar
Patrick Devine committed
974
975
976
	showCmd := &cobra.Command{
		Use:     "show MODEL",
		Short:   "Show information for a model",
Michael Yang's avatar
Michael Yang committed
977
		Args:    cobra.ExactArgs(1),
Patrick Devine's avatar
Patrick Devine committed
978
979
980
981
982
983
984
985
		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")
986
	showCmd.Flags().Bool("system", false, "Show system message of a model")
Patrick Devine's avatar
Patrick Devine committed
987

Jeffrey Morgan's avatar
Jeffrey Morgan committed
988
	runCmd := &cobra.Command{
989
990
991
992
		Use:   "run MODEL [PROMPT]",
		Short: "Run a model",
		Args:  cobra.MinimumNArgs(1),
		RunE:  RunHandler,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
993
994
	}

995
	runCmd.Flags().Bool("verbose", false, "Show timings for response")
996
	runCmd.Flags().Bool("insecure", false, "Use an insecure registry")
997
	runCmd.Flags().Bool("nowordwrap", false, "Don't wrap words to the next line automatically")
Jeffrey Morgan's avatar
Jeffrey Morgan committed
998
	runCmd.Flags().String("format", "", "Response format (e.g. json)")
Jeffrey Morgan's avatar
Jeffrey Morgan committed
999
1000
1001
1002
	serveCmd := &cobra.Command{
		Use:     "serve",
		Aliases: []string{"start"},
		Short:   "Start ollama",
Michael Yang's avatar
Michael Yang committed
1003
		Args:    cobra.ExactArgs(0),
Michael Yang's avatar
Michael Yang committed
1004
		RunE:    RunServer,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
1005
	}
1006
1007
1008
	serveCmd.SetUsageTemplate(serveCmd.UsageTemplate() + `
Environment Variables:

1009
1010
1011
1012
    OLLAMA_HOST         The host:port to bind to (default "127.0.0.1:11434")
    OLLAMA_ORIGINS      A comma separated list of allowed origins.
    OLLAMA_MODELS       The path to the models directory (default is "~/.ollama/models")
    OLLAMA_KEEP_ALIVE   The duration that models stay loaded in memory (default is "5m")
1013
    OLLAMA_DEBUG        Set to 1 to enable additional debug logging
1014
`)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
1015

1016
	pullCmd := &cobra.Command{
1017
1018
		Use:     "pull MODEL",
		Short:   "Pull a model from a registry",
Michael Yang's avatar
Michael Yang committed
1019
		Args:    cobra.ExactArgs(1),
1020
1021
		PreRunE: checkServerHeartbeat,
		RunE:    PullHandler,
1022
1023
	}

1024
1025
	pullCmd.Flags().Bool("insecure", false, "Use an insecure registry")

1026
	pushCmd := &cobra.Command{
1027
1028
		Use:     "push MODEL",
		Short:   "Push a model to a registry",
Michael Yang's avatar
Michael Yang committed
1029
		Args:    cobra.ExactArgs(1),
1030
1031
		PreRunE: checkServerHeartbeat,
		RunE:    PushHandler,
1032
1033
	}

1034
1035
	pushCmd.Flags().Bool("insecure", false, "Use an insecure registry")

Patrick Devine's avatar
Patrick Devine committed
1036
	listCmd := &cobra.Command{
1037
		Use:     "list",
Patrick Devine's avatar
Patrick Devine committed
1038
		Aliases: []string{"ls"},
1039
		Short:   "List models",
1040
		PreRunE: checkServerHeartbeat,
1041
		RunE:    ListHandler,
1042
	}
Patrick Devine's avatar
Patrick Devine committed
1043
	copyCmd := &cobra.Command{
Michael Yang's avatar
Michael Yang committed
1044
		Use:     "cp SOURCE TARGET",
1045
		Short:   "Copy a model",
Michael Yang's avatar
Michael Yang committed
1046
		Args:    cobra.ExactArgs(2),
1047
1048
		PreRunE: checkServerHeartbeat,
		RunE:    CopyHandler,
Patrick Devine's avatar
Patrick Devine committed
1049
1050
	}

1051
	deleteCmd := &cobra.Command{
Michael Yang's avatar
Michael Yang committed
1052
		Use:     "rm MODEL [MODEL...]",
1053
1054
1055
1056
		Short:   "Remove a model",
		Args:    cobra.MinimumNArgs(1),
		PreRunE: checkServerHeartbeat,
		RunE:    DeleteHandler,
Patrick Devine's avatar
Patrick Devine committed
1057
1058
	}

1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
	for _, cmd := range []*cobra.Command{
		createCmd,
		showCmd,
		runCmd,
		pullCmd,
		pushCmd,
		listCmd,
		copyCmd,
		deleteCmd,
	} {
		appendHostEnvDocs(cmd)
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
1072
1073
	rootCmd.AddCommand(
		serveCmd,
1074
		createCmd,
Patrick Devine's avatar
Patrick Devine committed
1075
		showCmd,
1076
		runCmd,
1077
1078
		pullCmd,
		pushCmd,
Patrick Devine's avatar
Patrick Devine committed
1079
		listCmd,
Patrick Devine's avatar
Patrick Devine committed
1080
		copyCmd,
1081
		deleteCmd,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
1082
1083
1084
1085
	)

	return rootCmd
}