cmd.go 22.9 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
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
			// 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)

				files, err := filepath.Glob(filepath.Join(path, "model-*.safetensors"))
				if err != nil {
					return err
				}

				if len(files) == 0 {
					return fmt.Errorf("no safetensors files were found in '%s'", path)
				}

				// add the safetensor config file + tokenizer
				files = append(files, filepath.Join(path, "config.json"))
				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)
					if os.IsNotExist(err) && strings.HasSuffix(fn, "added_tokens.json") {
						continue
					} 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
163
164
			}

165
166
			digest, err := createBlob(cmd, client, path)
			if err != nil {
Michael Yang's avatar
Michael Yang committed
167
168
169
				return err
			}

Michael Yang's avatar
Michael Yang committed
170
			modelfile = bytes.ReplaceAll(modelfile, []byte(c.Args), []byte("@"+digest))
Michael Yang's avatar
Michael Yang committed
171
172
		}
	}
Michael Yang's avatar
Michael Yang committed
173

174
	fn := func(resp api.ProgressResponse) error {
Michael Yang's avatar
Michael Yang committed
175
176
177
178
179
		if resp.Digest != "" {
			spinner.Stop()

			bar, ok := bars[resp.Digest]
			if !ok {
180
				bar = progress.NewBar(fmt.Sprintf("pulling %s...", resp.Digest[7:19]), resp.Total, resp.Completed)
Michael Yang's avatar
Michael Yang committed
181
182
183
184
185
186
187
188
189
190
191
192
193
				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)
		}

194
195
196
		return nil
	}

Michael Yang's avatar
Michael Yang committed
197
198
199
	quantization, _ := cmd.Flags().GetString("quantization")

	request := api.CreateRequest{Name: args[0], Modelfile: string(modelfile), Quantization: quantization}
Michael Yang's avatar
Michael Yang committed
200
	if err := client.Create(cmd.Context(), &request, fn); err != nil {
201
202
203
204
205
206
		return err
	}

	return nil
}

207
208
209
210
211
212
213
214
215
216
217
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
	}
218
219
220
221

	if _, err := bin.Seek(0, io.SeekStart); err != nil {
		return "", err
	}
222
223
224
225
226
227
228
229

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

230
func RunHandler(cmd *cobra.Command, args []string) error {
231
232
233
234
235
236
237
238
	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
239
	client, err := api.ClientFromEnvironment()
240
241
242
243
	if err != nil {
		return err
	}

244
	name := args[0]
245

246
	// check if the model exists on the server
247
	show, err := client.Show(cmd.Context(), &api.ShowRequest{Name: name})
Michael Yang's avatar
Michael Yang committed
248
249
250
	var statusError api.StatusError
	switch {
	case errors.As(err, &statusError) && statusError.StatusCode == http.StatusNotFound:
251
		if err := PullHandler(cmd, []string{name}); err != nil {
252
			return err
Michael Yang's avatar
Michael Yang committed
253
		}
254
255
256
257
258

		show, err = client.Show(cmd.Context(), &api.ShowRequest{Name: name})
		if err != nil {
			return err
		}
Michael Yang's avatar
Michael Yang committed
259
260
	case err != nil:
		return err
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
306
	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
307
308
}

309
func PushHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
310
	client, err := api.ClientFromEnvironment()
311
312
313
	if err != nil {
		return err
	}
314

315
316
317
318
319
	insecure, err := cmd.Flags().GetBool("insecure")
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
320
321
322
323
	p := progress.NewProgress(os.Stderr)
	defer p.Stop()

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

327
	fn := func(resp api.ProgressResponse) error {
Michael Yang's avatar
Michael Yang committed
328
		if resp.Digest != "" {
329
330
331
			if spinner != nil {
				spinner.Stop()
			}
Michael Yang's avatar
Michael Yang committed
332
333
334

			bar, ok := bars[resp.Digest]
			if !ok {
335
				bar = progress.NewBar(fmt.Sprintf("pushing %s...", resp.Digest[7:19]), resp.Total, resp.Completed)
Michael Yang's avatar
Michael Yang committed
336
337
338
339
340
341
				bars[resp.Digest] = bar
				p.Add(resp.Digest, bar)
			}

			bar.Set(resp.Completed)
		} else if status != resp.Status {
342
343
344
			if spinner != nil {
				spinner.Stop()
			}
Michael Yang's avatar
Michael Yang committed
345
346
347
348
349
350

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

351
352
353
		return nil
	}

Michael Yang's avatar
Michael Yang committed
354
	request := api.PushRequest{Name: args[0], Insecure: insecure}
Michael Yang's avatar
Michael Yang committed
355
	if err := client.Push(cmd.Context(), &request, fn); err != nil {
Michael Yang's avatar
Michael Yang committed
356
357
358
		return err
	}

359
	spinner.Stop()
Michael Yang's avatar
Michael Yang committed
360
	return nil
361
362
}

363
func ListHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
364
	client, err := api.ClientFromEnvironment()
365
366
367
	if err != nil {
		return err
	}
Patrick Devine's avatar
Patrick Devine committed
368

Michael Yang's avatar
Michael Yang committed
369
	models, err := client.List(cmd.Context())
Patrick Devine's avatar
Patrick Devine committed
370
371
372
373
374
375
376
	if err != nil {
		return err
	}

	var data [][]string

	for _, m := range models.Models {
Michael Yang's avatar
Michael Yang committed
377
		if len(args) == 0 || strings.HasPrefix(m.Name, args[0]) {
378
			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
379
		}
Patrick Devine's avatar
Patrick Devine committed
380
381
382
	}

	table := tablewriter.NewWriter(os.Stdout)
Patrick Devine's avatar
Patrick Devine committed
383
	table.SetHeader([]string{"NAME", "ID", "SIZE", "MODIFIED"})
Patrick Devine's avatar
Patrick Devine committed
384
385
386
387
388
389
390
391
392
393
394
395
	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
}

396
func DeleteHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
397
	client, err := api.ClientFromEnvironment()
398
399
400
	if err != nil {
		return err
	}
401

402
403
	for _, name := range args {
		req := api.DeleteRequest{Name: name}
Michael Yang's avatar
Michael Yang committed
404
		if err := client.Delete(cmd.Context(), &req); err != nil {
405
406
407
			return err
		}
		fmt.Printf("deleted '%s'\n", name)
408
409
410
411
	}
	return nil
}

Patrick Devine's avatar
Patrick Devine committed
412
func ShowHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
413
	client, err := api.ClientFromEnvironment()
Patrick Devine's avatar
Patrick Devine committed
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
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
	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 {
463
		return errors.New("only one of '--license', '--modelfile', '--parameters', '--system', or '--template' can be specified")
Patrick Devine's avatar
Patrick Devine committed
464
	} else if flagsSet == 0 {
465
		return errors.New("one of '--license', '--modelfile', '--parameters', '--system', or '--template' must be specified")
Patrick Devine's avatar
Patrick Devine committed
466
467
	}

468
	req := api.ShowRequest{Name: args[0]}
Michael Yang's avatar
Michael Yang committed
469
	resp, err := client.Show(cmd.Context(), &req)
Patrick Devine's avatar
Patrick Devine committed
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
	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
490
func CopyHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
491
	client, err := api.ClientFromEnvironment()
492
493
494
	if err != nil {
		return err
	}
Patrick Devine's avatar
Patrick Devine committed
495
496

	req := api.CopyRequest{Source: args[0], Destination: args[1]}
Michael Yang's avatar
Michael Yang committed
497
	if err := client.Copy(cmd.Context(), &req); err != nil {
Patrick Devine's avatar
Patrick Devine committed
498
499
500
501
502
503
		return err
	}
	fmt.Printf("copied '%s' to '%s'\n", args[0], args[1])
	return nil
}

504
func PullHandler(cmd *cobra.Command, args []string) error {
505
506
507
508
509
	insecure, err := cmd.Flags().GetBool("insecure")
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
510
	client, err := api.ClientFromEnvironment()
511
512
513
	if err != nil {
		return err
	}
514

Michael Yang's avatar
Michael Yang committed
515
516
517
518
519
	p := progress.NewProgress(os.Stderr)
	defer p.Stop()

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

520
521
	var status string
	var spinner *progress.Spinner
Michael Yang's avatar
Michael Yang committed
522

523
	fn := func(resp api.ProgressResponse) error {
Michael Yang's avatar
Michael Yang committed
524
		if resp.Digest != "" {
525
526
527
			if spinner != nil {
				spinner.Stop()
			}
Michael Yang's avatar
Michael Yang committed
528
529
530

			bar, ok := bars[resp.Digest]
			if !ok {
531
				bar = progress.NewBar(fmt.Sprintf("pulling %s...", resp.Digest[7:19]), resp.Total, resp.Completed)
Michael Yang's avatar
Michael Yang committed
532
533
534
535
536
537
				bars[resp.Digest] = bar
				p.Add(resp.Digest, bar)
			}

			bar.Set(resp.Completed)
		} else if status != resp.Status {
538
539
540
			if spinner != nil {
				spinner.Stop()
			}
Michael Yang's avatar
Michael Yang committed
541
542
543
544
545
546

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

547
548
		return nil
	}
549

Michael Yang's avatar
Michael Yang committed
550
	request := api.PullRequest{Name: args[0], Insecure: insecure}
Michael Yang's avatar
Michael Yang committed
551
	if err := client.Pull(cmd.Context(), &request, fn); err != nil {
Michael Yang's avatar
Michael Yang committed
552
553
554
555
		return err
	}

	return nil
Michael Yang's avatar
Michael Yang committed
556
557
}

558
559
type generateContextKey string

560
type runOptions struct {
561
562
563
564
565
566
567
568
569
570
571
	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
572
573
}

574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
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
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
691
	client, err := api.ClientFromEnvironment()
Patrick Devine's avatar
Patrick Devine committed
692
	if err != nil {
693
		return err
Patrick Devine's avatar
Patrick Devine committed
694
	}
Michael Yang's avatar
Michael Yang committed
695

Michael Yang's avatar
Michael Yang committed
696
	p := progress.NewProgress(os.Stderr)
697
	defer p.StopAndClear()
698

Michael Yang's avatar
Michael Yang committed
699
700
701
	spinner := progress.NewSpinner("")
	p.Add("", spinner)

702
703
704
705
706
707
708
	var latest api.GenerateResponse

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

Michael Yang's avatar
Michael Yang committed
709
	ctx, cancel := context.WithCancel(cmd.Context())
710
711
712
713
714
715
716
717
718
719
	defer cancel()

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

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

720
	var state *displayResponseState = &displayResponseState{}
721

722
	fn := func(response api.GenerateResponse) error {
Michael Yang's avatar
Michael Yang committed
723
		p.StopAndClear()
724

Patrick Devine's avatar
Patrick Devine committed
725
		latest = response
726
		content := response.Response
727

728
		displayResponse(content, opts.WordWrap, state)
729

Patrick Devine's avatar
Patrick Devine committed
730
731
		return nil
	}
732

733
734
735
736
737
738
739
	if opts.MultiModal {
		opts.Prompt, opts.Images, err = extractFileData(opts.Prompt)
		if err != nil {
			return err
		}
	}

Michael Yang's avatar
Michael Yang committed
740
741
742
743
	request := api.GenerateRequest{
		Model:    opts.Model,
		Prompt:   opts.Prompt,
		Context:  generateContext,
744
		Images:   opts.Images,
Michael Yang's avatar
Michael Yang committed
745
746
747
748
749
750
751
		Format:   opts.Format,
		System:   opts.System,
		Template: opts.Template,
		Options:  opts.Options,
	}

	if err := client.Generate(ctx, &request, fn); err != nil {
752
		if errors.Is(err, context.Canceled) {
753
			return nil
754
		}
755
		return err
Patrick Devine's avatar
Patrick Devine committed
756
	}
757

758
	if opts.Prompt != "" {
Michael Yang's avatar
Michael Yang committed
759
760
		fmt.Println()
		fmt.Println()
Patrick Devine's avatar
Patrick Devine committed
761
	}
762

763
764
765
766
	if !latest.Done {
		return nil
	}

Patrick Devine's avatar
Patrick Devine committed
767
768
	verbose, err := cmd.Flags().GetBool("verbose")
	if err != nil {
769
		return err
Patrick Devine's avatar
Patrick Devine committed
770
	}
Michael Yang's avatar
Michael Yang committed
771

Patrick Devine's avatar
Patrick Devine committed
772
773
	if verbose {
		latest.Summary()
Michael Yang's avatar
Michael Yang committed
774
	}
Michael Yang's avatar
Michael Yang committed
775

Patrick Devine's avatar
Patrick Devine committed
776
777
778
	ctx = context.WithValue(cmd.Context(), generateContextKey("context"), latest.Context)
	cmd.SetContext(ctx)

779
	return nil
Michael Yang's avatar
Michael Yang committed
780
781
}

782
func RunServer(cmd *cobra.Command, _ []string) error {
783
	host, port, err := net.SplitHostPort(strings.Trim(os.Getenv("OLLAMA_HOST"), "\"'"))
Michael Yang's avatar
Michael Yang committed
784
785
	if err != nil {
		host, port = "127.0.0.1", "11434"
Michael Yang's avatar
Michael Yang committed
786
		if ip := net.ParseIP(strings.Trim(os.Getenv("OLLAMA_HOST"), "[]")); ip != nil {
Michael Yang's avatar
Michael Yang committed
787
788
			host = ip.String()
		}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
789
	}
790

Michael Yang's avatar
Michael Yang committed
791
	if err := initializeKeypair(); err != nil {
792
793
794
		return err
	}

Michael Yang's avatar
Michael Yang committed
795
	ln, err := net.Listen("tcp", net.JoinHostPort(host, port))
796
797
798
	if err != nil {
		return err
	}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
799

800
	return server.Serve(ln)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
801
802
}

803
804
805
806
807
808
809
810
811
812
813
814
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
815
		cryptoPublicKey, cryptoPrivateKey, err := ed25519.GenerateKey(rand.Reader)
816
817
818
819
		if err != nil {
			return err
		}

Michael Yang's avatar
Michael Yang committed
820
		privateKeyBytes, err := ssh.MarshalPrivateKey(cryptoPrivateKey, "")
821
822
823
824
		if err != nil {
			return err
		}

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

Michael Yang's avatar
Michael Yang committed
829
		if err := os.WriteFile(privKeyPath, pem.EncodeToMemory(privateKeyBytes), 0o600); err != nil {
830
831
832
			return err
		}

Michael Yang's avatar
Michael Yang committed
833
		sshPublicKey, err := ssh.NewPublicKey(cryptoPublicKey)
834
835
836
837
		if err != nil {
			return err
		}

Michael Yang's avatar
Michael Yang committed
838
		publicKeyBytes := ssh.MarshalAuthorizedKey(sshPublicKey)
839

Michael Yang's avatar
Michael Yang committed
840
		if err := os.WriteFile(pubKeyPath, publicKeyBytes, 0o644); err != nil {
841
842
843
			return err
		}

Michael Yang's avatar
Michael Yang committed
844
		fmt.Printf("Your new public key is: \n\n%s\n", publicKeyBytes)
845
846
847
848
	}
	return nil
}

849
850
//nolint:unused
func waitForServer(ctx context.Context, client *api.Client) error {
Bruce MacDonald's avatar
Bruce MacDonald committed
851
852
853
854
855
856
857
858
	// 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
859
			if err := client.Heartbeat(ctx); err == nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
860
861
862
863
				return nil // server has started
			}
		}
	}
864

Bruce MacDonald's avatar
Bruce MacDonald committed
865
866
}

Michael Yang's avatar
Michael Yang committed
867
func checkServerHeartbeat(cmd *cobra.Command, _ []string) error {
Michael Yang's avatar
Michael Yang committed
868
	client, err := api.ClientFromEnvironment()
869
870
871
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
872
	if err := client.Heartbeat(cmd.Context()); err != nil {
873
		if !strings.Contains(err.Error(), " refused") {
Bruce MacDonald's avatar
Bruce MacDonald committed
874
875
			return err
		}
876
877
		if err := startApp(cmd.Context(), client); err != nil {
			return fmt.Errorf("could not connect to ollama app, is it running?")
878
879
880
881
882
		}
	}
	return nil
}

Michael Yang's avatar
Michael Yang committed
883
884
885
886
887
888
889
890
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
891
892
893
894
895
		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
896
897
	}

898
	if serverVersion != version.Version {
Michael Yang's avatar
Michael Yang committed
899
		fmt.Printf("Warning: client version is %s\n", version.Version)
900
	}
Michael Yang's avatar
Michael Yang committed
901
902
}

903
904
905
906
907
908
909
910
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
911
912
func NewCLI() *cobra.Command {
	log.SetFlags(log.LstdFlags | log.Lshortfile)
Michael Yang's avatar
Michael Yang committed
913
	cobra.EnableCommandSorting = false
Jeffrey Morgan's avatar
Jeffrey Morgan committed
914

915
	if runtime.GOOS == "windows" {
916
		console.ConsoleFromFile(os.Stdin) //nolint:errcheck
917
918
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
919
	rootCmd := &cobra.Command{
920
921
922
923
		Use:           "ollama",
		Short:         "Large language model runner",
		SilenceUsage:  true,
		SilenceErrors: true,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
924
925
926
		CompletionOptions: cobra.CompletionOptions{
			DisableDefaultCmd: true,
		},
Michael Yang's avatar
Michael Yang committed
927
928
929
930
931
932
933
934
		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
935
936
	}

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

939
	createCmd := &cobra.Command{
940
941
		Use:     "create MODEL",
		Short:   "Create a model from a Modelfile",
Michael Yang's avatar
Michael Yang committed
942
		Args:    cobra.ExactArgs(1),
943
944
		PreRunE: checkServerHeartbeat,
		RunE:    CreateHandler,
945
946
947
	}

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

Patrick Devine's avatar
Patrick Devine committed
950
951
952
	showCmd := &cobra.Command{
		Use:     "show MODEL",
		Short:   "Show information for a model",
Michael Yang's avatar
Michael Yang committed
953
		Args:    cobra.ExactArgs(1),
Patrick Devine's avatar
Patrick Devine committed
954
955
956
957
958
959
960
961
		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")
962
	showCmd.Flags().Bool("system", false, "Show system message of a model")
Patrick Devine's avatar
Patrick Devine committed
963

Jeffrey Morgan's avatar
Jeffrey Morgan committed
964
	runCmd := &cobra.Command{
965
966
967
968
		Use:   "run MODEL [PROMPT]",
		Short: "Run a model",
		Args:  cobra.MinimumNArgs(1),
		RunE:  RunHandler,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
969
970
	}

971
	runCmd.Flags().Bool("verbose", false, "Show timings for response")
972
	runCmd.Flags().Bool("insecure", false, "Use an insecure registry")
973
	runCmd.Flags().Bool("nowordwrap", false, "Don't wrap words to the next line automatically")
Jeffrey Morgan's avatar
Jeffrey Morgan committed
974
	runCmd.Flags().String("format", "", "Response format (e.g. json)")
Jeffrey Morgan's avatar
Jeffrey Morgan committed
975
976
977
978
	serveCmd := &cobra.Command{
		Use:     "serve",
		Aliases: []string{"start"},
		Short:   "Start ollama",
Michael Yang's avatar
Michael Yang committed
979
		Args:    cobra.ExactArgs(0),
Michael Yang's avatar
Michael Yang committed
980
		RunE:    RunServer,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
981
	}
982
983
984
	serveCmd.SetUsageTemplate(serveCmd.UsageTemplate() + `
Environment Variables:

985
986
987
988
    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")
989
    OLLAMA_DEBUG        Set to 1 to enable additional debug logging
990
`)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
991

992
	pullCmd := &cobra.Command{
993
994
		Use:     "pull MODEL",
		Short:   "Pull a model from a registry",
Michael Yang's avatar
Michael Yang committed
995
		Args:    cobra.ExactArgs(1),
996
997
		PreRunE: checkServerHeartbeat,
		RunE:    PullHandler,
998
999
	}

1000
1001
	pullCmd.Flags().Bool("insecure", false, "Use an insecure registry")

1002
	pushCmd := &cobra.Command{
1003
1004
		Use:     "push MODEL",
		Short:   "Push a model to a registry",
Michael Yang's avatar
Michael Yang committed
1005
		Args:    cobra.ExactArgs(1),
1006
1007
		PreRunE: checkServerHeartbeat,
		RunE:    PushHandler,
1008
1009
	}

1010
1011
	pushCmd.Flags().Bool("insecure", false, "Use an insecure registry")

Patrick Devine's avatar
Patrick Devine committed
1012
	listCmd := &cobra.Command{
1013
		Use:     "list",
Patrick Devine's avatar
Patrick Devine committed
1014
		Aliases: []string{"ls"},
1015
		Short:   "List models",
1016
		PreRunE: checkServerHeartbeat,
1017
		RunE:    ListHandler,
1018
	}
Patrick Devine's avatar
Patrick Devine committed
1019
	copyCmd := &cobra.Command{
Michael Yang's avatar
Michael Yang committed
1020
		Use:     "cp SOURCE TARGET",
1021
		Short:   "Copy a model",
Michael Yang's avatar
Michael Yang committed
1022
		Args:    cobra.ExactArgs(2),
1023
1024
		PreRunE: checkServerHeartbeat,
		RunE:    CopyHandler,
Patrick Devine's avatar
Patrick Devine committed
1025
1026
	}

1027
	deleteCmd := &cobra.Command{
Michael Yang's avatar
Michael Yang committed
1028
		Use:     "rm MODEL [MODEL...]",
1029
1030
1031
1032
		Short:   "Remove a model",
		Args:    cobra.MinimumNArgs(1),
		PreRunE: checkServerHeartbeat,
		RunE:    DeleteHandler,
Patrick Devine's avatar
Patrick Devine committed
1033
1034
	}

1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
	for _, cmd := range []*cobra.Command{
		createCmd,
		showCmd,
		runCmd,
		pullCmd,
		pushCmd,
		listCmd,
		copyCmd,
		deleteCmd,
	} {
		appendHostEnvDocs(cmd)
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
1048
1049
	rootCmd.AddCommand(
		serveCmd,
1050
		createCmd,
Patrick Devine's avatar
Patrick Devine committed
1051
		showCmd,
1052
		runCmd,
1053
1054
		pullCmd,
		pushCmd,
Patrick Devine's avatar
Patrick Devine committed
1055
		listCmd,
Patrick Devine's avatar
Patrick Devine committed
1056
		copyCmd,
1057
		deleteCmd,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
1058
1059
1060
1061
	)

	return rootCmd
}