cmd.go 23.4 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
	p := progress.NewProgress(os.Stderr)
	defer p.Stop()

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

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

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

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

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

96
			if fi.IsDir() {
Michael Yang's avatar
Michael Yang committed
97
98
99
				// this is likely a safetensors or pytorch directory
				// TODO make this work w/ adapters
				tempfile, err := tempZipFiles(path)
100
101
102
				if err != nil {
					return err
				}
Michael Yang's avatar
Michael Yang committed
103
				defer os.RemoveAll(tempfile)
104

Michael Yang's avatar
Michael Yang committed
105
				path = tempfile
Michael Yang's avatar
Michael Yang committed
106
107
			}

108
109
			digest, err := createBlob(cmd, client, path)
			if err != nil {
Michael Yang's avatar
Michael Yang committed
110
111
112
				return err
			}

Michael Yang's avatar
Michael Yang committed
113
			modelfile = bytes.ReplaceAll(modelfile, []byte(c.Args), []byte("@"+digest))
Michael Yang's avatar
Michael Yang committed
114
115
		}
	}
Michael Yang's avatar
Michael Yang committed
116

Michael Yang's avatar
Michael Yang committed
117
	bars := make(map[string]*progress.Bar)
118
	fn := func(resp api.ProgressResponse) error {
Michael Yang's avatar
Michael Yang committed
119
120
121
122
123
		if resp.Digest != "" {
			spinner.Stop()

			bar, ok := bars[resp.Digest]
			if !ok {
124
				bar = progress.NewBar(fmt.Sprintf("pulling %s...", resp.Digest[7:19]), resp.Total, resp.Completed)
Michael Yang's avatar
Michael Yang committed
125
126
127
128
129
130
131
132
133
134
135
136
137
				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)
		}

138
139
140
		return nil
	}

Michael Yang's avatar
Michael Yang committed
141
142
143
	quantization, _ := cmd.Flags().GetString("quantization")

	request := api.CreateRequest{Name: args[0], Modelfile: string(modelfile), Quantization: quantization}
Michael Yang's avatar
Michael Yang committed
144
	if err := client.Create(cmd.Context(), &request, fn); err != nil {
145
146
147
148
149
150
		return err
	}

	return nil
}

Michael Yang's avatar
Michael Yang committed
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
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
func tempZipFiles(path string) (string, error) {
	tempfile, err := os.CreateTemp("", "ollama-tf")
	if err != nil {
		return "", err
	}
	defer tempfile.Close()

	zipfile := zip.NewWriter(tempfile)
	defer zipfile.Close()

	tfiles, err := filepath.Glob(filepath.Join(path, "pytorch_model-*.bin"))
	if err != nil {
		return "", err
	} else if len(tfiles) == 0 {
		tfiles, err = filepath.Glob(filepath.Join(path, "model-*.safetensors"))
		if err != nil {
			return "", err
		}
	}

	files := []string{}
	files = append(files, tfiles...)

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

	// add the safetensor/torch config file + tokenizer
	files = append(files, filepath.Join(path, "config.json"))
	files = append(files, filepath.Join(path, "params.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)

		// 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
			}
		} 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 := zipfile.CreateHeader(h)
		if err != nil {
			return "", err
		}

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

	return tempfile.Name(), nil
}

233
234
235
236
237
238
239
240
241
242
243
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
	}
244
245
246
247

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

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

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

262
	name := args[0]
263

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

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

327
func PushHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
328
	client, err := api.ClientFromEnvironment()
329
330
331
	if err != nil {
		return err
	}
332

333
334
335
336
337
	insecure, err := cmd.Flags().GetBool("insecure")
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
338
339
340
341
	p := progress.NewProgress(os.Stderr)
	defer p.Stop()

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

345
	fn := func(resp api.ProgressResponse) error {
Michael Yang's avatar
Michael Yang committed
346
		if resp.Digest != "" {
347
348
349
			if spinner != nil {
				spinner.Stop()
			}
Michael Yang's avatar
Michael Yang committed
350
351
352

			bar, ok := bars[resp.Digest]
			if !ok {
353
				bar = progress.NewBar(fmt.Sprintf("pushing %s...", resp.Digest[7:19]), resp.Total, resp.Completed)
Michael Yang's avatar
Michael Yang committed
354
355
356
357
358
359
				bars[resp.Digest] = bar
				p.Add(resp.Digest, bar)
			}

			bar.Set(resp.Completed)
		} else if status != resp.Status {
360
361
362
			if spinner != nil {
				spinner.Stop()
			}
Michael Yang's avatar
Michael Yang committed
363
364
365
366
367
368

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

369
370
371
		return nil
	}

Michael Yang's avatar
Michael Yang committed
372
	request := api.PushRequest{Name: args[0], Insecure: insecure}
Michael Yang's avatar
Michael Yang committed
373
	if err := client.Push(cmd.Context(), &request, fn); err != nil {
Michael Yang's avatar
Michael Yang committed
374
375
376
		return err
	}

377
	spinner.Stop()
Michael Yang's avatar
Michael Yang committed
378
	return nil
379
380
}

381
func ListHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
382
	client, err := api.ClientFromEnvironment()
383
384
385
	if err != nil {
		return err
	}
Patrick Devine's avatar
Patrick Devine committed
386

Michael Yang's avatar
Michael Yang committed
387
	models, err := client.List(cmd.Context())
Patrick Devine's avatar
Patrick Devine committed
388
389
390
391
392
393
394
	if err != nil {
		return err
	}

	var data [][]string

	for _, m := range models.Models {
Michael Yang's avatar
Michael Yang committed
395
		if len(args) == 0 || strings.HasPrefix(m.Name, args[0]) {
396
			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
397
		}
Patrick Devine's avatar
Patrick Devine committed
398
399
400
	}

	table := tablewriter.NewWriter(os.Stdout)
Patrick Devine's avatar
Patrick Devine committed
401
	table.SetHeader([]string{"NAME", "ID", "SIZE", "MODIFIED"})
Patrick Devine's avatar
Patrick Devine committed
402
403
404
405
406
407
408
409
410
411
412
413
	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
}

414
func DeleteHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
415
	client, err := api.ClientFromEnvironment()
416
417
418
	if err != nil {
		return err
	}
419

420
421
	for _, name := range args {
		req := api.DeleteRequest{Name: name}
Michael Yang's avatar
Michael Yang committed
422
		if err := client.Delete(cmd.Context(), &req); err != nil {
423
424
425
			return err
		}
		fmt.Printf("deleted '%s'\n", name)
426
427
428
429
	}
	return nil
}

Patrick Devine's avatar
Patrick Devine committed
430
func ShowHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
431
	client, err := api.ClientFromEnvironment()
Patrick Devine's avatar
Patrick Devine committed
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
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
	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 {
481
		return errors.New("only one of '--license', '--modelfile', '--parameters', '--system', or '--template' can be specified")
Patrick Devine's avatar
Patrick Devine committed
482
	} else if flagsSet == 0 {
483
		return errors.New("one of '--license', '--modelfile', '--parameters', '--system', or '--template' must be specified")
Patrick Devine's avatar
Patrick Devine committed
484
485
	}

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

	req := api.CopyRequest{Source: args[0], Destination: args[1]}
Michael Yang's avatar
Michael Yang committed
515
	if err := client.Copy(cmd.Context(), &req); err != nil {
Patrick Devine's avatar
Patrick Devine committed
516
517
518
519
520
521
		return err
	}
	fmt.Printf("copied '%s' to '%s'\n", args[0], args[1])
	return nil
}

522
func PullHandler(cmd *cobra.Command, args []string) error {
523
524
525
526
527
	insecure, err := cmd.Flags().GetBool("insecure")
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
528
	client, err := api.ClientFromEnvironment()
529
530
531
	if err != nil {
		return err
	}
532

Michael Yang's avatar
Michael Yang committed
533
534
535
536
537
	p := progress.NewProgress(os.Stderr)
	defer p.Stop()

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

538
539
	var status string
	var spinner *progress.Spinner
Michael Yang's avatar
Michael Yang committed
540

541
	fn := func(resp api.ProgressResponse) error {
Michael Yang's avatar
Michael Yang committed
542
		if resp.Digest != "" {
543
544
545
			if spinner != nil {
				spinner.Stop()
			}
Michael Yang's avatar
Michael Yang committed
546
547
548

			bar, ok := bars[resp.Digest]
			if !ok {
549
				bar = progress.NewBar(fmt.Sprintf("pulling %s...", resp.Digest[7:19]), resp.Total, resp.Completed)
Michael Yang's avatar
Michael Yang committed
550
551
552
553
554
555
				bars[resp.Digest] = bar
				p.Add(resp.Digest, bar)
			}

			bar.Set(resp.Completed)
		} else if status != resp.Status {
556
557
558
			if spinner != nil {
				spinner.Stop()
			}
Michael Yang's avatar
Michael Yang committed
559
560
561
562
563
564

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

565
566
		return nil
	}
567

Michael Yang's avatar
Michael Yang committed
568
	request := api.PullRequest{Name: args[0], Insecure: insecure}
Michael Yang's avatar
Michael Yang committed
569
	if err := client.Pull(cmd.Context(), &request, fn); err != nil {
Michael Yang's avatar
Michael Yang committed
570
571
572
573
		return err
	}

	return nil
Michael Yang's avatar
Michael Yang committed
574
575
}

576
577
type generateContextKey string

578
type runOptions struct {
579
580
581
582
583
584
585
586
587
588
589
	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
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
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
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
709
	client, err := api.ClientFromEnvironment()
Patrick Devine's avatar
Patrick Devine committed
710
	if err != nil {
711
		return err
Patrick Devine's avatar
Patrick Devine committed
712
	}
Michael Yang's avatar
Michael Yang committed
713

Michael Yang's avatar
Michael Yang committed
714
	p := progress.NewProgress(os.Stderr)
715
	defer p.StopAndClear()
716

Michael Yang's avatar
Michael Yang committed
717
718
719
	spinner := progress.NewSpinner("")
	p.Add("", spinner)

720
721
722
723
724
725
726
	var latest api.GenerateResponse

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

Michael Yang's avatar
Michael Yang committed
727
	ctx, cancel := context.WithCancel(cmd.Context())
728
729
730
731
732
733
734
735
736
737
	defer cancel()

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

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

738
	var state *displayResponseState = &displayResponseState{}
739

740
	fn := func(response api.GenerateResponse) error {
Michael Yang's avatar
Michael Yang committed
741
		p.StopAndClear()
742

Patrick Devine's avatar
Patrick Devine committed
743
		latest = response
744
		content := response.Response
745

746
		displayResponse(content, opts.WordWrap, state)
747

Patrick Devine's avatar
Patrick Devine committed
748
749
		return nil
	}
750

751
752
753
754
755
756
757
	if opts.MultiModal {
		opts.Prompt, opts.Images, err = extractFileData(opts.Prompt)
		if err != nil {
			return err
		}
	}

Michael Yang's avatar
Michael Yang committed
758
759
760
761
	request := api.GenerateRequest{
		Model:    opts.Model,
		Prompt:   opts.Prompt,
		Context:  generateContext,
762
		Images:   opts.Images,
Michael Yang's avatar
Michael Yang committed
763
764
765
766
767
768
769
		Format:   opts.Format,
		System:   opts.System,
		Template: opts.Template,
		Options:  opts.Options,
	}

	if err := client.Generate(ctx, &request, fn); err != nil {
770
		if errors.Is(err, context.Canceled) {
771
			return nil
772
		}
773
		return err
Patrick Devine's avatar
Patrick Devine committed
774
	}
775

776
	if opts.Prompt != "" {
Michael Yang's avatar
Michael Yang committed
777
778
		fmt.Println()
		fmt.Println()
Patrick Devine's avatar
Patrick Devine committed
779
	}
780

781
782
783
784
	if !latest.Done {
		return nil
	}

Patrick Devine's avatar
Patrick Devine committed
785
786
	verbose, err := cmd.Flags().GetBool("verbose")
	if err != nil {
787
		return err
Patrick Devine's avatar
Patrick Devine committed
788
	}
Michael Yang's avatar
Michael Yang committed
789

Patrick Devine's avatar
Patrick Devine committed
790
791
	if verbose {
		latest.Summary()
Michael Yang's avatar
Michael Yang committed
792
	}
Michael Yang's avatar
Michael Yang committed
793

Patrick Devine's avatar
Patrick Devine committed
794
795
796
	ctx = context.WithValue(cmd.Context(), generateContextKey("context"), latest.Context)
	cmd.SetContext(ctx)

797
	return nil
Michael Yang's avatar
Michael Yang committed
798
799
}

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

Michael Yang's avatar
Michael Yang committed
809
	if err := initializeKeypair(); err != nil {
810
811
812
		return err
	}

Michael Yang's avatar
Michael Yang committed
813
	ln, err := net.Listen("tcp", net.JoinHostPort(host, port))
814
815
816
	if err != nil {
		return err
	}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
817

818
	return server.Serve(ln)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
819
820
}

821
822
823
824
825
826
827
828
829
830
831
832
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
833
		cryptoPublicKey, cryptoPrivateKey, err := ed25519.GenerateKey(rand.Reader)
834
835
836
837
		if err != nil {
			return err
		}

Michael Yang's avatar
Michael Yang committed
838
		privateKeyBytes, err := ssh.MarshalPrivateKey(cryptoPrivateKey, "")
839
840
841
842
		if err != nil {
			return err
		}

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

Michael Yang's avatar
Michael Yang committed
847
		if err := os.WriteFile(privKeyPath, pem.EncodeToMemory(privateKeyBytes), 0o600); err != nil {
848
849
850
			return err
		}

Michael Yang's avatar
Michael Yang committed
851
		sshPublicKey, err := ssh.NewPublicKey(cryptoPublicKey)
852
853
854
855
		if err != nil {
			return err
		}

Michael Yang's avatar
Michael Yang committed
856
		publicKeyBytes := ssh.MarshalAuthorizedKey(sshPublicKey)
857

Michael Yang's avatar
Michael Yang committed
858
		if err := os.WriteFile(pubKeyPath, publicKeyBytes, 0o644); err != nil {
859
860
861
			return err
		}

Michael Yang's avatar
Michael Yang committed
862
		fmt.Printf("Your new public key is: \n\n%s\n", publicKeyBytes)
863
864
865
866
	}
	return nil
}

867
868
//nolint:unused
func waitForServer(ctx context.Context, client *api.Client) error {
Bruce MacDonald's avatar
Bruce MacDonald committed
869
870
871
872
873
874
875
876
	// 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
877
			if err := client.Heartbeat(ctx); err == nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
878
879
880
881
				return nil // server has started
			}
		}
	}
882

Bruce MacDonald's avatar
Bruce MacDonald committed
883
884
}

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

Michael Yang's avatar
Michael Yang committed
901
902
903
904
905
906
907
908
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
909
910
911
912
913
		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
914
915
	}

916
	if serverVersion != version.Version {
Michael Yang's avatar
Michael Yang committed
917
		fmt.Printf("Warning: client version is %s\n", version.Version)
918
	}
Michael Yang's avatar
Michael Yang committed
919
920
}

921
922
923
924
925
926
927
928
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
929
930
func NewCLI() *cobra.Command {
	log.SetFlags(log.LstdFlags | log.Lshortfile)
Michael Yang's avatar
Michael Yang committed
931
	cobra.EnableCommandSorting = false
Jeffrey Morgan's avatar
Jeffrey Morgan committed
932

933
	if runtime.GOOS == "windows" {
934
		console.ConsoleFromFile(os.Stdin) //nolint:errcheck
935
936
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
937
	rootCmd := &cobra.Command{
938
939
940
941
		Use:           "ollama",
		Short:         "Large language model runner",
		SilenceUsage:  true,
		SilenceErrors: true,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
942
943
944
		CompletionOptions: cobra.CompletionOptions{
			DisableDefaultCmd: true,
		},
Michael Yang's avatar
Michael Yang committed
945
946
947
948
949
950
951
952
		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
953
954
	}

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

957
	createCmd := &cobra.Command{
958
959
		Use:     "create MODEL",
		Short:   "Create a model from a Modelfile",
Michael Yang's avatar
Michael Yang committed
960
		Args:    cobra.ExactArgs(1),
961
962
		PreRunE: checkServerHeartbeat,
		RunE:    CreateHandler,
963
964
965
	}

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

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

Jeffrey Morgan's avatar
Jeffrey Morgan committed
982
	runCmd := &cobra.Command{
983
984
985
986
987
		Use:     "run MODEL [PROMPT]",
		Short:   "Run a model",
		Args:    cobra.MinimumNArgs(1),
		PreRunE: checkServerHeartbeat,
		RunE:    RunHandler,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
988
989
	}

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

1004
1005
1006
1007
    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")
1008
    OLLAMA_DEBUG        Set to 1 to enable additional debug logging
1009
`)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
1010

1011
	pullCmd := &cobra.Command{
1012
1013
		Use:     "pull MODEL",
		Short:   "Pull a model from a registry",
Michael Yang's avatar
Michael Yang committed
1014
		Args:    cobra.ExactArgs(1),
1015
1016
		PreRunE: checkServerHeartbeat,
		RunE:    PullHandler,
1017
1018
	}

1019
1020
	pullCmd.Flags().Bool("insecure", false, "Use an insecure registry")

1021
	pushCmd := &cobra.Command{
1022
1023
		Use:     "push MODEL",
		Short:   "Push a model to a registry",
Michael Yang's avatar
Michael Yang committed
1024
		Args:    cobra.ExactArgs(1),
1025
1026
		PreRunE: checkServerHeartbeat,
		RunE:    PushHandler,
1027
1028
	}

1029
1030
	pushCmd.Flags().Bool("insecure", false, "Use an insecure registry")

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

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

1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
	for _, cmd := range []*cobra.Command{
		createCmd,
		showCmd,
		runCmd,
		pullCmd,
		pushCmd,
		listCmd,
		copyCmd,
		deleteCmd,
	} {
		appendHostEnvDocs(cmd)
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
1067
1068
	rootCmd.AddCommand(
		serveCmd,
1069
		createCmd,
Patrick Devine's avatar
Patrick Devine committed
1070
		showCmd,
1071
		runCmd,
1072
1073
		pullCmd,
		pushCmd,
Patrick Devine's avatar
Patrick Devine committed
1074
		listCmd,
Patrick Devine's avatar
Patrick Devine committed
1075
		copyCmd,
1076
		deleteCmd,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
1077
1078
1079
1080
	)

	return rootCmd
}