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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

128
129
130
		return nil
	}

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

	return nil
}

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

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

Michael Yang's avatar
Michael Yang committed
158
	return RunGenerate(cmd, args)
Bruce MacDonald's avatar
Bruce MacDonald committed
159
160
}

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

167
168
169
170
171
	insecure, err := cmd.Flags().GetBool("insecure")
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
172
173
174
175
	p := progress.NewProgress(os.Stderr)
	defer p.Stop()

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

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

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

			bar.Set(resp.Completed)
		} else if status != resp.Status {
194
195
196
			if spinner != nil {
				spinner.Stop()
			}
Michael Yang's avatar
Michael Yang committed
197
198
199
200
201
202

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

203
204
205
		return nil
	}

Michael Yang's avatar
Michael Yang committed
206
207
208
209
210
	request := api.PushRequest{Name: args[0], Insecure: insecure}
	if err := client.Push(context.Background(), &request, fn); err != nil {
		return err
	}

211
	spinner.Stop()
Michael Yang's avatar
Michael Yang committed
212
	return nil
213
214
}

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

	models, err := client.List(context.Background())
	if err != nil {
		return err
	}

	var data [][]string

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

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

248
func DeleteHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
249
	client, err := api.ClientFromEnvironment()
250
251
252
	if err != nil {
		return err
	}
253

254
255
256
257
258
259
	for _, name := range args {
		req := api.DeleteRequest{Name: name}
		if err := client.Delete(context.Background(), &req); err != nil {
			return err
		}
		fmt.Printf("deleted '%s'\n", name)
260
261
262
263
	}
	return nil
}

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

	req := api.ShowRequest{Name: args[0]}
	resp, err := client.Show(context.Background(), &req)
	if err != nil {
		return err
	}

	switch showType {
	case "license":
		fmt.Println(resp.License)
	case "modelfile":
		fmt.Println(resp.Modelfile)
	case "parameters":
		fmt.Println(resp.Parameters)
	case "system":
		fmt.Println(resp.System)
	case "template":
		fmt.Println(resp.Template)
	}

	return nil
}

Patrick Devine's avatar
Patrick Devine committed
342
func CopyHandler(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
343
	client, err := api.ClientFromEnvironment()
344
345
346
	if err != nil {
		return err
	}
Patrick Devine's avatar
Patrick Devine committed
347
348
349
350
351
352
353
354
355

	req := api.CopyRequest{Source: args[0], Destination: args[1]}
	if err := client.Copy(context.Background(), &req); err != nil {
		return err
	}
	fmt.Printf("copied '%s' to '%s'\n", args[0], args[1])
	return nil
}

356
func PullHandler(cmd *cobra.Command, args []string) error {
357
358
359
360
361
	insecure, err := cmd.Flags().GetBool("insecure")
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
362
	client, err := api.ClientFromEnvironment()
363
364
365
	if err != nil {
		return err
	}
366

Michael Yang's avatar
Michael Yang committed
367
368
369
370
371
	p := progress.NewProgress(os.Stderr)
	defer p.Stop()

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

372
373
	var status string
	var spinner *progress.Spinner
Michael Yang's avatar
Michael Yang committed
374

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

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

			bar.Set(resp.Completed)
		} else if status != resp.Status {
390
391
392
			if spinner != nil {
				spinner.Stop()
			}
Michael Yang's avatar
Michael Yang committed
393
394
395
396
397
398

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

399
400
		return nil
	}
401

Michael Yang's avatar
Michael Yang committed
402
403
404
405
406
407
	request := api.PullRequest{Name: args[0], Insecure: insecure}
	if err := client.Pull(context.Background(), &request, fn); err != nil {
		return err
	}

	return nil
Michael Yang's avatar
Michael Yang committed
408
409
}

410
func RunGenerate(cmd *cobra.Command, args []string) error {
411
412
413
414
	format, err := cmd.Flags().GetString("format")
	if err != nil {
		return err
	}
415

416
	prompts := args[1:]
417

418
419
420
	// prepend stdin to the prompt if provided
	if !term.IsTerminal(int(os.Stdin.Fd())) {
		in, err := io.ReadAll(os.Stdin)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
421
422
423
424
		if err != nil {
			return err
		}

425
		prompts = append([]string{string(in)}, prompts...)
Michael Yang's avatar
Michael Yang committed
426
427
	}

428
429
430
	// output is being piped
	if !term.IsTerminal(int(os.Stdout.Fd())) {
		return generate(cmd, args[0], strings.Join(prompts, " "), false, format)
Michael Yang's avatar
Michael Yang committed
431
432
	}

433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
	wordWrap := os.Getenv("TERM") == "xterm-256color"

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

	// prompts are provided via stdin or args so don't enter interactive mode
	if len(prompts) > 0 {
		return generate(cmd, args[0], strings.Join(prompts, " "), wordWrap, format)
	}

	return generateInteractive(cmd, args[0], wordWrap, format)
Michael Yang's avatar
Michael Yang committed
449
450
}

Michael Yang's avatar
Michael Yang committed
451
type generateContextKey string
Michael Yang's avatar
Michael Yang committed
452

Jeffrey Morgan's avatar
Jeffrey Morgan committed
453
func generate(cmd *cobra.Command, model, prompt string, wordWrap bool, format string) error {
Michael Yang's avatar
Michael Yang committed
454
	client, err := api.ClientFromEnvironment()
Patrick Devine's avatar
Patrick Devine committed
455
456
457
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
458

Michael Yang's avatar
Michael Yang committed
459
	p := progress.NewProgress(os.Stderr)
460
	defer p.StopAndClear()
Michael Yang's avatar
Michael Yang committed
461
462
463
464

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

Patrick Devine's avatar
Patrick Devine committed
465
	var latest api.GenerateResponse
466

Patrick Devine's avatar
Patrick Devine committed
467
468
469
470
	generateContext, ok := cmd.Context().Value(generateContextKey("context")).([]int)
	if !ok {
		generateContext = []int{}
	}
Michael Yang's avatar
Michael Yang committed
471

472
	termWidth, _, err := term.GetSize(int(os.Stdout.Fd()))
473
	if err != nil {
474
		wordWrap = false
475
476
	}

477
478
479
480
481
482
483
484
485
486
487
488
489
	cancelCtx, cancel := context.WithCancel(context.Background())
	defer cancel()

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

	go func() {
		<-sigChan
		cancel()
		abort = true
	}()

490
491
492
	var currentLineLength int
	var wordBuffer string

Jeffrey Morgan's avatar
Jeffrey Morgan committed
493
	request := api.GenerateRequest{Model: model, Prompt: prompt, Context: generateContext, Format: format}
Patrick Devine's avatar
Patrick Devine committed
494
	fn := func(response api.GenerateResponse) error {
Michael Yang's avatar
Michael Yang committed
495
496
		p.StopAndClear()

Patrick Devine's avatar
Patrick Devine committed
497
		latest = response
498

499
		if wordWrap {
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
			for _, ch := range response.Response {
				if currentLineLength+1 > termWidth-5 {
					// backtrack the length of the last word and clear to the end of the line
					fmt.Printf("\x1b[%dD\x1b[K\n", len(wordBuffer))
					fmt.Printf("%s%c", wordBuffer, ch)
					currentLineLength = len(wordBuffer) + 1
				} else {
					fmt.Print(string(ch))
					currentLineLength += 1

					switch ch {
					case ' ':
						wordBuffer = ""
					case '\n':
						currentLineLength = 0
					default:
						wordBuffer += string(ch)
					}
				}
			}
		} else {
			fmt.Print(response.Response)
		}

Patrick Devine's avatar
Patrick Devine committed
524
525
		return nil
	}
526

527
	if err := client.Generate(cancelCtx, &request, fn); err != nil {
528
		if strings.Contains(err.Error(), "context canceled") && abort {
529
			return nil
530
		}
Patrick Devine's avatar
Patrick Devine committed
531
532
533
		return err
	}
	if prompt != "" {
Michael Yang's avatar
Michael Yang committed
534
535
		fmt.Println()
		fmt.Println()
Patrick Devine's avatar
Patrick Devine committed
536
	}
537

Patrick Devine's avatar
Patrick Devine committed
538
	if !latest.Done {
539
540
541
		if abort {
			return nil
		}
Patrick Devine's avatar
Patrick Devine committed
542
543
		return errors.New("unexpected end of response")
	}
544

Patrick Devine's avatar
Patrick Devine committed
545
546
547
548
	verbose, err := cmd.Flags().GetBool("verbose")
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
549

Patrick Devine's avatar
Patrick Devine committed
550
551
	if verbose {
		latest.Summary()
Michael Yang's avatar
Michael Yang committed
552
	}
Michael Yang's avatar
Michael Yang committed
553

Patrick Devine's avatar
Patrick Devine committed
554
555
556
557
	ctx := cmd.Context()
	ctx = context.WithValue(ctx, generateContextKey("context"), latest.Context)
	cmd.SetContext(ctx)

Michael Yang's avatar
Michael Yang committed
558
559
560
	return nil
}

561
func generateInteractive(cmd *cobra.Command, model string, wordWrap bool, format string) error {
Patrick Devine's avatar
Patrick Devine committed
562
	// load the model
Jeffrey Morgan's avatar
Jeffrey Morgan committed
563
	if err := generate(cmd, model, "", false, ""); err != nil {
Patrick Devine's avatar
Patrick Devine committed
564
565
566
		return err
	}

Michael Yang's avatar
Michael Yang committed
567
	usage := func() {
Patrick Devine's avatar
Patrick Devine committed
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
		fmt.Fprintln(os.Stderr, "Available Commands:")
		fmt.Fprintln(os.Stderr, "  /set         Set session variables")
		fmt.Fprintln(os.Stderr, "  /show        Show model information")
		fmt.Fprintln(os.Stderr, "  /bye         Exit")
		fmt.Fprintln(os.Stderr, "  /?, /help    Help for a command")
		fmt.Fprintln(os.Stderr, "")
		fmt.Fprintln(os.Stderr, "Use \"\"\" to begin a multi-line message.")
		fmt.Fprintln(os.Stderr, "")
	}

	usageSet := func() {
		fmt.Fprintln(os.Stderr, "Available Commands:")
		fmt.Fprintln(os.Stderr, "  /set history      Enable history")
		fmt.Fprintln(os.Stderr, "  /set nohistory    Disable history")
		fmt.Fprintln(os.Stderr, "  /set wordwrap     Enable wordwrap")
		fmt.Fprintln(os.Stderr, "  /set nowordwrap   Disable wordwrap")
Jeffrey Morgan's avatar
Jeffrey Morgan committed
584
585
		fmt.Fprintln(os.Stderr, "  /set format json  Enable JSON mode")
		fmt.Fprintln(os.Stderr, "  /set noformat     Disable formatting")
Patrick Devine's avatar
Patrick Devine committed
586
587
588
589
590
591
592
593
594
595
596
597
598
		fmt.Fprintln(os.Stderr, "  /set verbose      Show LLM stats")
		fmt.Fprintln(os.Stderr, "  /set quiet        Disable LLM stats")
		fmt.Fprintln(os.Stderr, "")
	}

	usageShow := func() {
		fmt.Fprintln(os.Stderr, "Available Commands:")
		fmt.Fprintln(os.Stderr, "  /show license      Show model license")
		fmt.Fprintln(os.Stderr, "  /show modelfile    Show Modelfile for this model")
		fmt.Fprintln(os.Stderr, "  /show parameters   Show parameters for this model")
		fmt.Fprintln(os.Stderr, "  /show system       Show system prompt")
		fmt.Fprintln(os.Stderr, "  /show template     Show prompt template")
		fmt.Fprintln(os.Stderr, "")
Michael Yang's avatar
Michael Yang committed
599
600
	}

Patrick Devine's avatar
Patrick Devine committed
601
602
603
604
605
	prompt := readline.Prompt{
		Prompt:         ">>> ",
		AltPrompt:      "... ",
		Placeholder:    "Send a message (/? for help)",
		AltPlaceholder: `Use """ to end multi-line input`,
Michael Yang's avatar
Michael Yang committed
606
607
	}

Patrick Devine's avatar
Patrick Devine committed
608
	scanner, err := readline.New(prompt)
Michael Yang's avatar
Michael Yang committed
609
610
611
612
	if err != nil {
		return err
	}

613
614
615
	fmt.Print(readline.StartBracketedPaste)
	defer fmt.Printf(readline.EndBracketedPaste)

616
617
	var multiLineBuffer string

Michael Yang's avatar
Michael Yang committed
618
619
620
621
	for {
		line, err := scanner.Readline()
		switch {
		case errors.Is(err, io.EOF):
622
			fmt.Println()
Michael Yang's avatar
Michael Yang committed
623
624
			return nil
		case errors.Is(err, readline.ErrInterrupt):
625
			if line == "" {
Patrick Devine's avatar
Patrick Devine committed
626
				fmt.Println("\nUse Ctrl-D or /bye to exit.")
627
628
			}

Michael Yang's avatar
Michael Yang committed
629
630
			continue
		case err != nil:
Michael Yang's avatar
Michael Yang committed
631
632
633
			return err
		}

Michael Yang's avatar
Michael Yang committed
634
		line = strings.TrimSpace(line)
Michael Yang's avatar
Michael Yang committed
635

Michael Yang's avatar
Michael Yang committed
636
		switch {
Patrick Devine's avatar
Patrick Devine committed
637
		case scanner.Prompt.UseAlt:
638
			if strings.HasSuffix(line, `"""`) {
Patrick Devine's avatar
Patrick Devine committed
639
				scanner.Prompt.UseAlt = false
640
				multiLineBuffer += strings.TrimSuffix(line, `"""`)
641
				line = multiLineBuffer
642
643
644
645
646
647
				multiLineBuffer = ""
			} else {
				multiLineBuffer += line + " "
				continue
			}
		case strings.HasPrefix(line, `"""`):
Patrick Devine's avatar
Patrick Devine committed
648
			scanner.Prompt.UseAlt = true
649
650
			multiLineBuffer = strings.TrimPrefix(line, `"""`) + " "
			continue
Michael Yang's avatar
Michael Yang committed
651
652
		case strings.HasPrefix(line, "/list"):
			args := strings.Fields(line)
653
			if err := ListHandler(cmd, args[1:]); err != nil {
Michael Yang's avatar
Michael Yang committed
654
655
656
657
658
659
660
661
662
663
				return err
			}
		case strings.HasPrefix(line, "/set"):
			args := strings.Fields(line)
			if len(args) > 1 {
				switch args[1] {
				case "history":
					scanner.HistoryEnable()
				case "nohistory":
					scanner.HistoryDisable()
664
				case "wordwrap":
665
					wordWrap = true
666
667
					fmt.Println("Set 'wordwrap' mode.")
				case "nowordwrap":
668
					wordWrap = false
669
					fmt.Println("Set 'nowordwrap' mode.")
Michael Yang's avatar
Michael Yang committed
670
671
				case "verbose":
					cmd.Flags().Set("verbose", "true")
672
					fmt.Println("Set 'verbose' mode.")
Michael Yang's avatar
Michael Yang committed
673
674
				case "quiet":
					cmd.Flags().Set("verbose", "false")
675
					fmt.Println("Set 'quiet' mode.")
Jeffrey Morgan's avatar
Jeffrey Morgan committed
676
677
678
679
680
681
682
683
684
685
				case "format":
					if len(args) < 3 || args[2] != "json" {
						fmt.Println("Invalid or missing format. For 'json' mode use '/set format json'")
					} else {
						format = args[2]
						fmt.Printf("Set format to '%s' mode.\n", args[2])
					}
				case "noformat":
					format = ""
					fmt.Println("Disabled format.")
686
687
				default:
					fmt.Printf("Unknown command '/set %s'. Type /? for help\n", args[1])
688
689
				}
			} else {
Patrick Devine's avatar
Patrick Devine committed
690
				usageSet()
691
692
693
694
			}
		case strings.HasPrefix(line, "/show"):
			args := strings.Fields(line)
			if len(args) > 1 {
695
696
697
698
699
700
				client, err := api.ClientFromEnvironment()
				if err != nil {
					fmt.Println("error: couldn't connect to ollama server")
					return err
				}
				resp, err := client.Show(cmd.Context(), &api.ShowRequest{Name: model})
701
				if err != nil {
Patrick Devine's avatar
Patrick Devine committed
702
					fmt.Println("error: couldn't get model")
703
					return err
704
				}
Patrick Devine's avatar
Patrick Devine committed
705

706
707
				switch args[1] {
				case "license":
708
					if resp.License == "" {
709
						fmt.Print("No license was specified for this model.\n\n")
710
711
712
					} else {
						fmt.Println(resp.License)
					}
Patrick Devine's avatar
Patrick Devine committed
713
714
715
				case "modelfile":
					fmt.Println(resp.Modelfile)
				case "parameters":
716
					if resp.Parameters == "" {
717
						fmt.Print("No parameters were specified for this model.\n\n")
718
719
720
					} else {
						fmt.Println(resp.Parameters)
					}
721
				case "system":
722
					if resp.System == "" {
723
						fmt.Print("No system prompt was specified for this model.\n\n")
724
725
726
					} else {
						fmt.Println(resp.System)
					}
727
				case "template":
728
					if resp.Template == "" {
729
						fmt.Print("No prompt template was specified for this model.\n\n")
730
731
732
					} else {
						fmt.Println(resp.Template)
					}
733
				default:
734
					fmt.Printf("Unknown command '/show %s'. Type /? for help\n", args[1])
Michael Yang's avatar
Michael Yang committed
735
				}
Patrick Devine's avatar
Patrick Devine committed
736
737
738
739
740
741
742
743
744
745
746
747
			} else {
				usageShow()
			}
		case strings.HasPrefix(line, "/help"), strings.HasPrefix(line, "/?"):
			args := strings.Fields(line)
			if len(args) > 1 {
				switch args[1] {
				case "set", "/set":
					usageSet()
				case "show", "/show":
					usageShow()
				}
748
749
			} else {
				usage()
Michael Yang's avatar
Michael Yang committed
750
751
752
			}
		case line == "/exit", line == "/bye":
			return nil
753
754
755
		case strings.HasPrefix(line, "/"):
			args := strings.Fields(line)
			fmt.Printf("Unknown command '%s'. Type /? for help\n", args[0])
Michael Yang's avatar
Michael Yang committed
756
757
		}

Patrick Devine's avatar
Patrick Devine committed
758
		if len(line) > 0 && line[0] != '/' {
Jeffrey Morgan's avatar
Jeffrey Morgan committed
759
			if err := generate(cmd, model, line, wordWrap, format); err != nil {
Patrick Devine's avatar
Patrick Devine committed
760
761
				return err
			}
Michael Yang's avatar
Michael Yang committed
762
763
		}
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
764
765
}

766
func RunServer(cmd *cobra.Command, _ []string) error {
Michael Yang's avatar
Michael Yang committed
767
768
769
	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
770
		if ip := net.ParseIP(strings.Trim(os.Getenv("OLLAMA_HOST"), "[]")); ip != nil {
Michael Yang's avatar
Michael Yang committed
771
772
			host = ip.String()
		}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
773
	}
774

Michael Yang's avatar
Michael Yang committed
775
	if err := initializeKeypair(); err != nil {
776
777
778
		return err
	}

Michael Yang's avatar
Michael Yang committed
779
	ln, err := net.Listen("tcp", net.JoinHostPort(host, port))
780
781
782
	if err != nil {
		return err
	}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
783

784
785
786
787
788
	var origins []string
	if o := os.Getenv("OLLAMA_ORIGINS"); o != "" {
		origins = strings.Split(o, ",")
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
789
	return server.Serve(ln, origins)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
790
791
}

792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
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
814
		err = os.MkdirAll(filepath.Dir(privKeyPath), 0o755)
815
816
817
818
		if err != nil {
			return fmt.Errorf("could not create directory %w", err)
		}

819
		err = os.WriteFile(privKeyPath, pem.EncodeToMemory(privKeyBytes), 0o600)
820
821
822
823
824
825
826
827
828
829
830
		if err != nil {
			return err
		}

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

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

831
		err = os.WriteFile(pubKeyPath, pubKeyData, 0o644)
832
833
834
835
836
837
838
839
840
		if err != nil {
			return err
		}

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

Bruce MacDonald's avatar
Bruce MacDonald committed
841
func startMacApp(client *api.Client) error {
Bruce MacDonald's avatar
Bruce MacDonald committed
842
843
844
845
846
	exe, err := os.Executable()
	if err != nil {
		return err
	}
	link, err := os.Readlink(exe)
Bruce MacDonald's avatar
Bruce MacDonald committed
847
848
849
	if err != nil {
		return err
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
850
851
852
	if !strings.Contains(link, "Ollama.app") {
		return fmt.Errorf("could not find ollama app")
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
	path := strings.Split(link, "Ollama.app")
	if err := exec.Command("/usr/bin/open", "-a", path[0]+"Ollama.app").Run(); err != nil {
		return err
	}
	// wait for the server to start
	timeout := time.After(5 * time.Second)
	tick := time.Tick(500 * time.Millisecond)
	for {
		select {
		case <-timeout:
			return errors.New("timed out waiting for server to start")
		case <-tick:
			if err := client.Heartbeat(context.Background()); err == nil {
				return nil // server has started
			}
		}
	}
}

872
func checkServerHeartbeat(_ *cobra.Command, _ []string) error {
Michael Yang's avatar
Michael Yang committed
873
	client, err := api.ClientFromEnvironment()
874
875
876
	if err != nil {
		return err
	}
877
	if err := client.Heartbeat(context.Background()); err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
878
879
880
881
882
		if !strings.Contains(err.Error(), "connection refused") {
			return err
		}
		if runtime.GOOS == "darwin" {
			if err := startMacApp(client); err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
883
				return fmt.Errorf("could not connect to ollama app, is it running?")
884
			}
Bruce MacDonald's avatar
Bruce MacDonald committed
885
		} else {
886
887
888
889
890
891
			return fmt.Errorf("could not connect to ollama server, run 'ollama serve' to start it")
		}
	}
	return nil
}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
892
893
894
895
func NewCLI() *cobra.Command {
	log.SetFlags(log.LstdFlags | log.Lshortfile)

	rootCmd := &cobra.Command{
896
897
898
899
		Use:           "ollama",
		Short:         "Large language model runner",
		SilenceUsage:  true,
		SilenceErrors: true,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
900
901
902
		CompletionOptions: cobra.CompletionOptions{
			DisableDefaultCmd: true,
		},
Michael Yang's avatar
Michael Yang committed
903
		Version: version.Version,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
904
905
906
907
	}

	cobra.EnableCommandSorting = false

908
	createCmd := &cobra.Command{
909
910
		Use:     "create MODEL",
		Short:   "Create a model from a Modelfile",
Michael Yang's avatar
Michael Yang committed
911
		Args:    cobra.ExactArgs(1),
912
913
		PreRunE: checkServerHeartbeat,
		RunE:    CreateHandler,
914
915
916
917
	}

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

Patrick Devine's avatar
Patrick Devine committed
918
919
920
	showCmd := &cobra.Command{
		Use:     "show MODEL",
		Short:   "Show information for a model",
Michael Yang's avatar
Michael Yang committed
921
		Args:    cobra.ExactArgs(1),
Patrick Devine's avatar
Patrick Devine committed
922
923
924
925
926
927
928
929
930
931
		PreRunE: checkServerHeartbeat,
		RunE:    ShowHandler,
	}

	showCmd.Flags().Bool("license", false, "Show license of a model")
	showCmd.Flags().Bool("modelfile", false, "Show Modelfile of a model")
	showCmd.Flags().Bool("parameters", false, "Show parameters of a model")
	showCmd.Flags().Bool("template", false, "Show template of a model")
	showCmd.Flags().Bool("system", false, "Show system prompt of a model")

Jeffrey Morgan's avatar
Jeffrey Morgan committed
932
	runCmd := &cobra.Command{
933
934
935
936
937
		Use:     "run MODEL [PROMPT]",
		Short:   "Run a model",
		Args:    cobra.MinimumNArgs(1),
		PreRunE: checkServerHeartbeat,
		RunE:    RunHandler,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
938
939
	}

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

Jeffrey Morgan's avatar
Jeffrey Morgan committed
945
946
947
948
	serveCmd := &cobra.Command{
		Use:     "serve",
		Aliases: []string{"start"},
		Short:   "Start ollama",
Michael Yang's avatar
Michael Yang committed
949
		Args:    cobra.ExactArgs(0),
Michael Yang's avatar
Michael Yang committed
950
		RunE:    RunServer,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
951
952
	}

953
	pullCmd := &cobra.Command{
954
955
		Use:     "pull MODEL",
		Short:   "Pull a model from a registry",
Michael Yang's avatar
Michael Yang committed
956
		Args:    cobra.ExactArgs(1),
957
958
		PreRunE: checkServerHeartbeat,
		RunE:    PullHandler,
959
960
	}

961
962
	pullCmd.Flags().Bool("insecure", false, "Use an insecure registry")

963
	pushCmd := &cobra.Command{
964
965
		Use:     "push MODEL",
		Short:   "Push a model to a registry",
Michael Yang's avatar
Michael Yang committed
966
		Args:    cobra.ExactArgs(1),
967
968
		PreRunE: checkServerHeartbeat,
		RunE:    PushHandler,
969
970
	}

971
972
	pushCmd.Flags().Bool("insecure", false, "Use an insecure registry")

Patrick Devine's avatar
Patrick Devine committed
973
	listCmd := &cobra.Command{
974
		Use:     "list",
Patrick Devine's avatar
Patrick Devine committed
975
		Aliases: []string{"ls"},
976
		Short:   "List models",
977
		PreRunE: checkServerHeartbeat,
978
		RunE:    ListHandler,
979
980
	}

Patrick Devine's avatar
Patrick Devine committed
981
	copyCmd := &cobra.Command{
Michael Yang's avatar
Michael Yang committed
982
		Use:     "cp SOURCE TARGET",
983
		Short:   "Copy a model",
Michael Yang's avatar
Michael Yang committed
984
		Args:    cobra.ExactArgs(2),
985
986
		PreRunE: checkServerHeartbeat,
		RunE:    CopyHandler,
Patrick Devine's avatar
Patrick Devine committed
987
988
	}

989
	deleteCmd := &cobra.Command{
Michael Yang's avatar
Michael Yang committed
990
		Use:     "rm MODEL [MODEL...]",
991
992
993
994
		Short:   "Remove a model",
		Args:    cobra.MinimumNArgs(1),
		PreRunE: checkServerHeartbeat,
		RunE:    DeleteHandler,
Patrick Devine's avatar
Patrick Devine committed
995
996
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
997
998
	rootCmd.AddCommand(
		serveCmd,
999
		createCmd,
Patrick Devine's avatar
Patrick Devine committed
1000
		showCmd,
1001
		runCmd,
1002
1003
		pullCmd,
		pushCmd,
Patrick Devine's avatar
Patrick Devine committed
1004
		listCmd,
Patrick Devine's avatar
Patrick Devine committed
1005
		copyCmd,
1006
		deleteCmd,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
1007
1008
1009
1010
	)

	return rootCmd
}