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

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

Patrick Devine's avatar
Patrick Devine committed
25
	"github.com/olekukonko/tablewriter"
Michael Yang's avatar
Michael Yang committed
26
	"github.com/spf13/cobra"
27
	"golang.org/x/crypto/ssh"
28
	"golang.org/x/term"
Michael Yang's avatar
Michael Yang committed
29

Jeffrey Morgan's avatar
Jeffrey Morgan committed
30
	"github.com/jmorganca/ollama/api"
Patrick Devine's avatar
Patrick Devine committed
31
	"github.com/jmorganca/ollama/format"
Michael Yang's avatar
Michael Yang committed
32
	"github.com/jmorganca/ollama/parser"
Michael Yang's avatar
Michael Yang committed
33
	"github.com/jmorganca/ollama/progress"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
34
	"github.com/jmorganca/ollama/server"
Michael Yang's avatar
Michael Yang committed
35
	"github.com/jmorganca/ollama/version"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
36
37
)

Patrick Devine's avatar
Patrick Devine committed
38
39
type ImageData []byte

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

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

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

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

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

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

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

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

Michael Yang's avatar
Michael Yang committed
90
91
			bin, err := os.Open(path)
			if errors.Is(err, os.ErrNotExist) && c.Name == "model" {
Michael Yang's avatar
Michael Yang committed
92
				continue
Michael Yang's avatar
Michael Yang committed
93
94
95
96
97
98
99
100
101
102
103
104
			} else if err != nil {
				return err
			}
			defer bin.Close()

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

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

Michael Yang's avatar
Michael Yang committed
109
			modelfile = bytes.ReplaceAll(modelfile, []byte(c.Args), []byte("@"+digest))
Michael Yang's avatar
Michael Yang committed
110
111
		}
	}
Michael Yang's avatar
Michael Yang committed
112

113
	fn := func(resp api.ProgressResponse) error {
Michael Yang's avatar
Michael Yang committed
114
115
116
117
118
		if resp.Digest != "" {
			spinner.Stop()

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

133
134
135
		return nil
	}

136
	request := api.CreateRequest{Name: args[0], Modelfile: string(modelfile)}
Michael Yang's avatar
Michael Yang committed
137
	if err := client.Create(cmd.Context(), &request, fn); err != nil {
138
139
140
141
142
143
		return err
	}

	return nil
}

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

150
151
	name := args[0]
	// check if the model exists on the server
152
	_, err = client.Show(cmd.Context(), &api.ShowRequest{Name: name})
Michael Yang's avatar
Michael Yang committed
153
154
155
	var statusError api.StatusError
	switch {
	case errors.As(err, &statusError) && statusError.StatusCode == http.StatusNotFound:
156
		if err := PullHandler(cmd, []string{name}); err != nil {
157
			return err
Michael Yang's avatar
Michael Yang committed
158
		}
Michael Yang's avatar
Michael Yang committed
159
160
	case err != nil:
		return err
161
162
	}

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

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

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

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

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

184
	fn := func(resp api.ProgressResponse) error {
Michael Yang's avatar
Michael Yang committed
185
		if resp.Digest != "" {
186
187
188
			if spinner != nil {
				spinner.Stop()
			}
Michael Yang's avatar
Michael Yang committed
189
190
191

			bar, ok := bars[resp.Digest]
			if !ok {
192
				bar = progress.NewBar(fmt.Sprintf("pushing %s...", resp.Digest[7:19]), resp.Total, resp.Completed)
Michael Yang's avatar
Michael Yang committed
193
194
195
196
197
198
				bars[resp.Digest] = bar
				p.Add(resp.Digest, bar)
			}

			bar.Set(resp.Completed)
		} else if status != resp.Status {
199
200
201
			if spinner != nil {
				spinner.Stop()
			}
Michael Yang's avatar
Michael Yang committed
202
203
204
205
206
207

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

208
209
210
		return nil
	}

Michael Yang's avatar
Michael Yang committed
211
	request := api.PushRequest{Name: args[0], Insecure: insecure}
Michael Yang's avatar
Michael Yang committed
212
	if err := client.Push(cmd.Context(), &request, fn); err != nil {
Michael Yang's avatar
Michael Yang committed
213
214
215
		return err
	}

216
	spinner.Stop()
Michael Yang's avatar
Michael Yang committed
217
	return nil
218
219
}

220
func ListHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
221
	client, err := api.ClientFromEnvironment()
222
223
224
	if err != nil {
		return err
	}
Patrick Devine's avatar
Patrick Devine committed
225

Michael Yang's avatar
Michael Yang committed
226
	models, err := client.List(cmd.Context())
Patrick Devine's avatar
Patrick Devine committed
227
228
229
230
231
232
233
	if err != nil {
		return err
	}

	var data [][]string

	for _, m := range models.Models {
Michael Yang's avatar
Michael Yang committed
234
		if len(args) == 0 || strings.HasPrefix(m.Name, args[0]) {
235
			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
236
		}
Patrick Devine's avatar
Patrick Devine committed
237
238
239
	}

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

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

259
260
	for _, name := range args {
		req := api.DeleteRequest{Name: name}
Michael Yang's avatar
Michael Yang committed
261
		if err := client.Delete(cmd.Context(), &req); err != nil {
262
263
264
			return err
		}
		fmt.Printf("deleted '%s'\n", name)
265
266
267
268
	}
	return nil
}

Patrick Devine's avatar
Patrick Devine committed
269
func ShowHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
270
	client, err := api.ClientFromEnvironment()
Patrick Devine's avatar
Patrick Devine committed
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
307
308
309
310
311
312
313
314
315
316
317
318
319
	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 {
320
		return errors.New("only one of '--license', '--modelfile', '--parameters', '--system', or '--template' can be specified")
Patrick Devine's avatar
Patrick Devine committed
321
	} else if flagsSet == 0 {
322
		return errors.New("one of '--license', '--modelfile', '--parameters', '--system', or '--template' must be specified")
Patrick Devine's avatar
Patrick Devine committed
323
324
	}

325
	req := api.ShowRequest{Name: args[0]}
Michael Yang's avatar
Michael Yang committed
326
	resp, err := client.Show(cmd.Context(), &req)
Patrick Devine's avatar
Patrick Devine committed
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
	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
347
func CopyHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
348
	client, err := api.ClientFromEnvironment()
349
350
351
	if err != nil {
		return err
	}
Patrick Devine's avatar
Patrick Devine committed
352
353

	req := api.CopyRequest{Source: args[0], Destination: args[1]}
Michael Yang's avatar
Michael Yang committed
354
	if err := client.Copy(cmd.Context(), &req); err != nil {
Patrick Devine's avatar
Patrick Devine committed
355
356
357
358
359
360
		return err
	}
	fmt.Printf("copied '%s' to '%s'\n", args[0], args[1])
	return nil
}

361
func PullHandler(cmd *cobra.Command, args []string) error {
362
363
364
365
366
	insecure, err := cmd.Flags().GetBool("insecure")
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
367
	client, err := api.ClientFromEnvironment()
368
369
370
	if err != nil {
		return err
	}
371

Michael Yang's avatar
Michael Yang committed
372
373
374
375
376
	p := progress.NewProgress(os.Stderr)
	defer p.Stop()

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

377
378
	var status string
	var spinner *progress.Spinner
Michael Yang's avatar
Michael Yang committed
379

380
	fn := func(resp api.ProgressResponse) error {
Michael Yang's avatar
Michael Yang committed
381
		if resp.Digest != "" {
382
383
384
			if spinner != nil {
				spinner.Stop()
			}
Michael Yang's avatar
Michael Yang committed
385
386
387

			bar, ok := bars[resp.Digest]
			if !ok {
388
				bar = progress.NewBar(fmt.Sprintf("pulling %s...", resp.Digest[7:19]), resp.Total, resp.Completed)
Michael Yang's avatar
Michael Yang committed
389
390
391
392
393
394
				bars[resp.Digest] = bar
				p.Add(resp.Digest, bar)
			}

			bar.Set(resp.Completed)
		} else if status != resp.Status {
395
396
397
			if spinner != nil {
				spinner.Stop()
			}
Michael Yang's avatar
Michael Yang committed
398
399
400
401
402
403

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

404
405
		return nil
	}
406

Michael Yang's avatar
Michael Yang committed
407
	request := api.PullRequest{Name: args[0], Insecure: insecure}
Michael Yang's avatar
Michael Yang committed
408
	if err := client.Pull(cmd.Context(), &request, fn); err != nil {
Michael Yang's avatar
Michael Yang committed
409
410
411
412
		return err
	}

	return nil
Michael Yang's avatar
Michael Yang committed
413
414
}

415
416
417
418
419
420
421
func RunGenerate(cmd *cobra.Command, args []string) error {
	interactive := true

	opts := generateOptions{
		Model:    args[0],
		WordWrap: os.Getenv("TERM") == "xterm-256color",
		Options:  map[string]interface{}{},
Patrick Devine's avatar
Patrick Devine committed
422
		Images:   []ImageData{},
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
463
	}

	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)
}

type generateContextKey string

type generateOptions struct {
464
	Model    string
465
	Prompt   string
466
467
	WordWrap bool
	Format   string
468
	System   string
469
	Template string
Patrick Devine's avatar
Patrick Devine committed
470
	Images   []ImageData
471
472
473
	Options  map[string]interface{}
}

474
func generate(cmd *cobra.Command, opts generateOptions) error {
Michael Yang's avatar
Michael Yang committed
475
	client, err := api.ClientFromEnvironment()
Patrick Devine's avatar
Patrick Devine committed
476
	if err != nil {
477
		return err
Patrick Devine's avatar
Patrick Devine committed
478
	}
Michael Yang's avatar
Michael Yang committed
479

Michael Yang's avatar
Michael Yang committed
480
	p := progress.NewProgress(os.Stderr)
481
	defer p.StopAndClear()
482

Michael Yang's avatar
Michael Yang committed
483
484
485
	spinner := progress.NewSpinner("")
	p.Add("", spinner)

486
487
488
489
490
491
492
	var latest api.GenerateResponse

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

493
	termWidth, _, err := term.GetSize(int(os.Stdout.Fd()))
494
	if err != nil {
495
		opts.WordWrap = false
496
497
	}

Michael Yang's avatar
Michael Yang committed
498
	ctx, cancel := context.WithCancel(cmd.Context())
499
500
501
502
503
504
505
506
507
508
	defer cancel()

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

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

509
510
511
	var currentLineLength int
	var wordBuffer string

512
	fn := func(response api.GenerateResponse) error {
Michael Yang's avatar
Michael Yang committed
513
		p.StopAndClear()
514

Patrick Devine's avatar
Patrick Devine committed
515
		latest = response
516

517
518
		termWidth, _, _ = term.GetSize(int(os.Stdout.Fd()))
		if opts.WordWrap && termWidth >= 10 {
519
			for _, ch := range response.Response {
520
				if currentLineLength+1 > termWidth-5 {
521
522
523
524
525
526
527
					if len(wordBuffer) > termWidth-10 {
						fmt.Printf("%s%c", wordBuffer, ch)
						wordBuffer = ""
						currentLineLength = 0
						continue
					}

528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
					// backtrack the length of the last word and clear to the end of the line
					fmt.Printf("\x1b[%dD\x1b[K\n", len(wordBuffer))
					fmt.Printf("%s%c", wordBuffer, ch)
					currentLineLength = len(wordBuffer) + 1
				} else {
					fmt.Print(string(ch))
					currentLineLength += 1

					switch ch {
					case ' ':
						wordBuffer = ""
					case '\n':
						currentLineLength = 0
					default:
						wordBuffer += string(ch)
					}
				}
			}
		} else {
547
			fmt.Printf("%s%s", wordBuffer, response.Response)
548
549
550
			if len(wordBuffer) > 0 {
				wordBuffer = ""
			}
551
552
		}

Patrick Devine's avatar
Patrick Devine committed
553
554
		return nil
	}
555

Patrick Devine's avatar
Patrick Devine committed
556
557
558
559
	images := make([]api.ImageData, 0)
	for _, i := range opts.Images {
		images = append(images, api.ImageData(i))
	}
Michael Yang's avatar
Michael Yang committed
560
561
562
563
564
565
566
567
	request := api.GenerateRequest{
		Model:    opts.Model,
		Prompt:   opts.Prompt,
		Context:  generateContext,
		Format:   opts.Format,
		System:   opts.System,
		Template: opts.Template,
		Options:  opts.Options,
Patrick Devine's avatar
Patrick Devine committed
568
		Images:   images,
Michael Yang's avatar
Michael Yang committed
569
570
571
	}

	if err := client.Generate(ctx, &request, fn); err != nil {
572
		if errors.Is(err, context.Canceled) {
573
			return nil
574
		}
575
		return err
Patrick Devine's avatar
Patrick Devine committed
576
	}
577

578
	if opts.Prompt != "" {
Michael Yang's avatar
Michael Yang committed
579
580
		fmt.Println()
		fmt.Println()
Patrick Devine's avatar
Patrick Devine committed
581
	}
582

583
584
585
586
	if !latest.Done {
		return nil
	}

Patrick Devine's avatar
Patrick Devine committed
587
588
	verbose, err := cmd.Flags().GetBool("verbose")
	if err != nil {
589
		return err
Patrick Devine's avatar
Patrick Devine committed
590
	}
Michael Yang's avatar
Michael Yang committed
591

Patrick Devine's avatar
Patrick Devine committed
592
593
	if verbose {
		latest.Summary()
Michael Yang's avatar
Michael Yang committed
594
	}
Michael Yang's avatar
Michael Yang committed
595

Patrick Devine's avatar
Patrick Devine committed
596
597
598
	ctx = context.WithValue(cmd.Context(), generateContextKey("context"), latest.Context)
	cmd.SetContext(ctx)

599
	return nil
Michael Yang's avatar
Michael Yang committed
600
601
}

602
func RunServer(cmd *cobra.Command, _ []string) error {
Michael Yang's avatar
Michael Yang committed
603
604
605
	host, port, err := net.SplitHostPort(os.Getenv("OLLAMA_HOST"))
	if err != nil {
		host, port = "127.0.0.1", "11434"
Michael Yang's avatar
Michael Yang committed
606
		if ip := net.ParseIP(strings.Trim(os.Getenv("OLLAMA_HOST"), "[]")); ip != nil {
Michael Yang's avatar
Michael Yang committed
607
608
			host = ip.String()
		}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
609
	}
610

Michael Yang's avatar
Michael Yang committed
611
	if err := initializeKeypair(); err != nil {
612
613
614
		return err
	}

Michael Yang's avatar
Michael Yang committed
615
	ln, err := net.Listen("tcp", net.JoinHostPort(host, port))
616
617
618
	if err != nil {
		return err
	}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
619

620
	return server.Serve(ln)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
621
622
}

623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
func initializeKeypair() error {
	home, err := os.UserHomeDir()
	if err != nil {
		return err
	}

	privKeyPath := filepath.Join(home, ".ollama", "id_ed25519")
	pubKeyPath := filepath.Join(home, ".ollama", "id_ed25519.pub")

	_, err = os.Stat(privKeyPath)
	if os.IsNotExist(err) {
		fmt.Printf("Couldn't find '%s'. Generating new private key.\n", privKeyPath)
		_, privKey, err := ed25519.GenerateKey(rand.Reader)
		if err != nil {
			return err
		}

		privKeyBytes, err := format.OpenSSHPrivateKey(privKey, "")
		if err != nil {
			return err
		}

Michael Yang's avatar
Michael Yang committed
645
		err = os.MkdirAll(filepath.Dir(privKeyPath), 0o755)
646
647
648
649
		if err != nil {
			return fmt.Errorf("could not create directory %w", err)
		}

650
		err = os.WriteFile(privKeyPath, pem.EncodeToMemory(privKeyBytes), 0o600)
651
652
653
654
655
656
657
658
659
660
661
		if err != nil {
			return err
		}

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

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

662
		err = os.WriteFile(pubKeyPath, pubKeyData, 0o644)
663
664
665
666
667
668
669
670
671
		if err != nil {
			return err
		}

		fmt.Printf("Your new public key is: \n\n%s\n", string(pubKeyData))
	}
	return nil
}

Michael Yang's avatar
Michael Yang committed
672
func startMacApp(ctx context.Context, client *api.Client) error {
Bruce MacDonald's avatar
Bruce MacDonald committed
673
674
675
676
677
	exe, err := os.Executable()
	if err != nil {
		return err
	}
	link, err := os.Readlink(exe)
Bruce MacDonald's avatar
Bruce MacDonald committed
678
679
680
	if err != nil {
		return err
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
681
682
683
	if !strings.Contains(link, "Ollama.app") {
		return fmt.Errorf("could not find ollama app")
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
684
685
686
687
688
689
690
691
692
693
694
695
	path := strings.Split(link, "Ollama.app")
	if err := exec.Command("/usr/bin/open", "-a", path[0]+"Ollama.app").Run(); err != nil {
		return err
	}
	// wait for the server to start
	timeout := time.After(5 * time.Second)
	tick := time.Tick(500 * time.Millisecond)
	for {
		select {
		case <-timeout:
			return errors.New("timed out waiting for server to start")
		case <-tick:
Michael Yang's avatar
Michael Yang committed
696
			if err := client.Heartbeat(ctx); err == nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
697
698
699
700
701
702
				return nil // server has started
			}
		}
	}
}

Michael Yang's avatar
Michael Yang committed
703
func checkServerHeartbeat(cmd *cobra.Command, _ []string) error {
Michael Yang's avatar
Michael Yang committed
704
	client, err := api.ClientFromEnvironment()
705
706
707
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
708
	if err := client.Heartbeat(cmd.Context()); err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
709
710
711
712
		if !strings.Contains(err.Error(), "connection refused") {
			return err
		}
		if runtime.GOOS == "darwin" {
Michael Yang's avatar
Michael Yang committed
713
			if err := startMacApp(cmd.Context(), client); err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
714
				return fmt.Errorf("could not connect to ollama app, is it running?")
715
			}
Bruce MacDonald's avatar
Bruce MacDonald committed
716
		} else {
717
718
719
720
721
722
			return fmt.Errorf("could not connect to ollama server, run 'ollama serve' to start it")
		}
	}
	return nil
}

Michael Yang's avatar
Michael Yang committed
723
724
725
726
727
728
729
730
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
731
732
733
734
735
		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
736
737
	}

738
	if serverVersion != version.Version {
Michael Yang's avatar
Michael Yang committed
739
		fmt.Printf("Warning: client version is %s\n", version.Version)
740
	}
Michael Yang's avatar
Michael Yang committed
741
742
}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
743
744
func NewCLI() *cobra.Command {
	log.SetFlags(log.LstdFlags | log.Lshortfile)
Michael Yang's avatar
Michael Yang committed
745
	cobra.EnableCommandSorting = false
Jeffrey Morgan's avatar
Jeffrey Morgan committed
746
747

	rootCmd := &cobra.Command{
748
749
750
751
		Use:           "ollama",
		Short:         "Large language model runner",
		SilenceUsage:  true,
		SilenceErrors: true,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
752
753
754
		CompletionOptions: cobra.CompletionOptions{
			DisableDefaultCmd: true,
		},
Michael Yang's avatar
Michael Yang committed
755
756
757
758
759
760
761
762
		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
763
764
	}

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

767
	createCmd := &cobra.Command{
768
769
		Use:     "create MODEL",
		Short:   "Create a model from a Modelfile",
Michael Yang's avatar
Michael Yang committed
770
		Args:    cobra.ExactArgs(1),
771
772
		PreRunE: checkServerHeartbeat,
		RunE:    CreateHandler,
773
774
775
776
	}

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

Patrick Devine's avatar
Patrick Devine committed
777
778
779
	showCmd := &cobra.Command{
		Use:     "show MODEL",
		Short:   "Show information for a model",
Michael Yang's avatar
Michael Yang committed
780
		Args:    cobra.ExactArgs(1),
Patrick Devine's avatar
Patrick Devine committed
781
782
783
784
785
786
787
788
		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")
789
	showCmd.Flags().Bool("system", false, "Show system message of a model")
Patrick Devine's avatar
Patrick Devine committed
790

Jeffrey Morgan's avatar
Jeffrey Morgan committed
791
	runCmd := &cobra.Command{
792
793
794
795
796
		Use:     "run MODEL [PROMPT]",
		Short:   "Run a model",
		Args:    cobra.MinimumNArgs(1),
		PreRunE: checkServerHeartbeat,
		RunE:    RunHandler,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
797
798
	}

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

Jeffrey Morgan's avatar
Jeffrey Morgan committed
804
805
806
807
	serveCmd := &cobra.Command{
		Use:     "serve",
		Aliases: []string{"start"},
		Short:   "Start ollama",
Michael Yang's avatar
Michael Yang committed
808
		Args:    cobra.ExactArgs(0),
Michael Yang's avatar
Michael Yang committed
809
		RunE:    RunServer,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
810
811
	}

812
	pullCmd := &cobra.Command{
813
814
		Use:     "pull MODEL",
		Short:   "Pull a model from a registry",
Michael Yang's avatar
Michael Yang committed
815
		Args:    cobra.ExactArgs(1),
816
817
		PreRunE: checkServerHeartbeat,
		RunE:    PullHandler,
818
819
	}

820
821
	pullCmd.Flags().Bool("insecure", false, "Use an insecure registry")

822
	pushCmd := &cobra.Command{
823
824
		Use:     "push MODEL",
		Short:   "Push a model to a registry",
Michael Yang's avatar
Michael Yang committed
825
		Args:    cobra.ExactArgs(1),
826
827
		PreRunE: checkServerHeartbeat,
		RunE:    PushHandler,
828
829
	}

830
831
	pushCmd.Flags().Bool("insecure", false, "Use an insecure registry")

Patrick Devine's avatar
Patrick Devine committed
832
	listCmd := &cobra.Command{
833
		Use:     "list",
Patrick Devine's avatar
Patrick Devine committed
834
		Aliases: []string{"ls"},
835
		Short:   "List models",
836
		PreRunE: checkServerHeartbeat,
837
		RunE:    ListHandler,
838
839
	}

Patrick Devine's avatar
Patrick Devine committed
840
	copyCmd := &cobra.Command{
Michael Yang's avatar
Michael Yang committed
841
		Use:     "cp SOURCE TARGET",
842
		Short:   "Copy a model",
Michael Yang's avatar
Michael Yang committed
843
		Args:    cobra.ExactArgs(2),
844
845
		PreRunE: checkServerHeartbeat,
		RunE:    CopyHandler,
Patrick Devine's avatar
Patrick Devine committed
846
847
	}

848
	deleteCmd := &cobra.Command{
Michael Yang's avatar
Michael Yang committed
849
		Use:     "rm MODEL [MODEL...]",
850
851
852
853
		Short:   "Remove a model",
		Args:    cobra.MinimumNArgs(1),
		PreRunE: checkServerHeartbeat,
		RunE:    DeleteHandler,
Patrick Devine's avatar
Patrick Devine committed
854
855
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
856
857
	rootCmd.AddCommand(
		serveCmd,
858
		createCmd,
Patrick Devine's avatar
Patrick Devine committed
859
		showCmd,
860
		runCmd,
861
862
		pullCmd,
		pushCmd,
Patrick Devine's avatar
Patrick Devine committed
863
		listCmd,
Patrick Devine's avatar
Patrick Devine committed
864
		copyCmd,
865
		deleteCmd,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
866
867
868
869
	)

	return rootCmd
}