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

import (
Michael Yang's avatar
Michael Yang committed
4
	"bufio"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
5
	"context"
Michael Yang's avatar
Michael Yang committed
6
	"errors"
Bruce MacDonald's avatar
Bruce MacDonald committed
7
	"fmt"
Michael Yang's avatar
Michael Yang committed
8
	"io"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
9
10
	"log"
	"net"
Michael Yang's avatar
Michael Yang committed
11
	"net/http"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
12
	"os"
13
	"path/filepath"
Michael Yang's avatar
Michael Yang committed
14
	"strings"
Michael Yang's avatar
Michael Yang committed
15
	"time"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
16

Michael Yang's avatar
Michael Yang committed
17
	"github.com/chzyer/readline"
Patrick Devine's avatar
Patrick Devine committed
18
19
	"github.com/dustin/go-humanize"
	"github.com/olekukonko/tablewriter"
Michael Yang's avatar
Michael Yang committed
20
21
	"github.com/spf13/cobra"

Jeffrey Morgan's avatar
Jeffrey Morgan committed
22
	"github.com/jmorganca/ollama/api"
Patrick Devine's avatar
Patrick Devine committed
23
	"github.com/jmorganca/ollama/format"
24
	"github.com/jmorganca/ollama/progressbar"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
25
	"github.com/jmorganca/ollama/server"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
26
27
)

28
29
func create(cmd *cobra.Command, args []string) error {
	filename, _ := cmd.Flags().GetString("file")
30
31
32
33
34
	filename, err := filepath.Abs(filename)
	if err != nil {
		return err
	}

35
36
	client := api.NewClient()

Michael Yang's avatar
Michael Yang committed
37
38
	var spinner *Spinner

39
40
	request := api.CreateRequest{Name: args[0], Path: filename}
	fn := func(resp api.CreateProgress) error {
Michael Yang's avatar
Michael Yang committed
41
42
43
44
45
46
47
		if spinner != nil {
			spinner.Stop()
		}

		spinner = NewSpinner(resp.Status)
		go spinner.Spin(100 * time.Millisecond)

48
49
50
51
52
53
54
		return nil
	}

	if err := client.Create(context.Background(), &request, fn); err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
55
56
57
58
	if spinner != nil {
		spinner.Stop()
	}

59
60
61
	return nil
}

Michael Yang's avatar
Michael Yang committed
62
func RunRun(cmd *cobra.Command, args []string) error {
Patrick Devine's avatar
Patrick Devine committed
63
64
65
66
67
68
69
	mp := server.ParseModelPath(args[0])
	fp, err := mp.GetManifestPath(false)
	if err != nil {
		return err
	}

	_, err = os.Stat(fp)
Michael Yang's avatar
Michael Yang committed
70
71
72
	switch {
	case errors.Is(err, os.ErrNotExist):
		if err := pull(args[0]); err != nil {
Michael Yang's avatar
Michael Yang committed
73
74
75
76
77
78
79
80
			var apiStatusError api.StatusError
			if !errors.As(err, &apiStatusError) {
				return err
			}

			if apiStatusError.StatusCode != http.StatusBadGateway {
				return err
			}
Michael Yang's avatar
Michael Yang committed
81
82
		}
	case err != nil:
Michael Yang's avatar
Michael Yang committed
83
84
85
86
		return err
	}

	return RunGenerate(cmd, args)
Bruce MacDonald's avatar
Bruce MacDonald committed
87
88
}

89
90
91
92
func push(cmd *cobra.Command, args []string) error {
	client := api.NewClient()

	request := api.PushRequest{Name: args[0]}
93
	fn := func(resp api.ProgressResponse) error {
94
95
96
97
98
99
100
101
102
103
		fmt.Println(resp.Status)
		return nil
	}

	if err := client.Push(context.Background(), &request, fn); err != nil {
		return err
	}
	return nil
}

Patrick Devine's avatar
Patrick Devine committed
104
105
106
107
108
109
110
111
112
113
114
func list(cmd *cobra.Command, args []string) error {
	client := api.NewClient()

	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
115
116
117
		if len(args) == 0 || strings.HasPrefix(m.Name, args[0]) {
			data = append(data, []string{m.Name, humanize.Bytes(uint64(m.Size)), format.HumanTime(m.ModifiedAt, "Never")})
		}
Patrick Devine's avatar
Patrick Devine committed
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
	}

	table := tablewriter.NewWriter(os.Stdout)
	table.SetHeader([]string{"NAME", "SIZE", "MODIFIED"})
	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
}

134
135
136
137
func RunPull(cmd *cobra.Command, args []string) error {
	return pull(args[0])
}

Michael Yang's avatar
Michael Yang committed
138
func pull(model string) error {
Bruce MacDonald's avatar
Bruce MacDonald committed
139
	client := api.NewClient()
140

141
	var currentDigest string
Bruce MacDonald's avatar
Bruce MacDonald committed
142
	var bar *progressbar.ProgressBar
Michael Yang's avatar
Michael Yang committed
143

144
	request := api.PullRequest{Name: model}
145
146
147
	fn := func(resp api.ProgressResponse) error {
		if resp.Digest != currentDigest && resp.Digest != "" {
			currentDigest = resp.Digest
148
149
			bar = progressbar.DefaultBytes(
				int64(resp.Total),
150
				fmt.Sprintf("pulling %s...", resp.Digest[7:19]),
151
			)
152
153
154

			bar.Set(resp.Completed)
		} else if resp.Digest == currentDigest && resp.Digest != "" {
155
156
			bar.Set(resp.Completed)
		} else {
157
			currentDigest = ""
158
159
160
161
			fmt.Println(resp.Status)
		}
		return nil
	}
162

163
164
165
166
	if err := client.Pull(context.Background(), &request, fn); err != nil {
		return err
	}
	return nil
Michael Yang's avatar
Michael Yang committed
167
168
}

169
func RunGenerate(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
170
	if len(args) > 1 {
Michael Yang's avatar
Michael Yang committed
171
		// join all args into a single prompt
172
		return generate(cmd, args[0], strings.Join(args[1:], " "))
Michael Yang's avatar
Michael Yang committed
173
174
	}

Michael Yang's avatar
Michael Yang committed
175
	if readline.IsTerminal(int(os.Stdin.Fd())) {
176
		return generateInteractive(cmd, args[0])
Michael Yang's avatar
Michael Yang committed
177
178
	}

179
	return generateBatch(cmd, args[0])
Michael Yang's avatar
Michael Yang committed
180
181
}

Michael Yang's avatar
Michael Yang committed
182
183
var generateContextKey struct{}

184
func generate(cmd *cobra.Command, model, prompt string) error {
Michael Yang's avatar
Michael Yang committed
185
186
187
	if len(strings.TrimSpace(prompt)) > 0 {
		client := api.NewClient()

Michael Yang's avatar
Michael Yang committed
188
189
		spinner := NewSpinner("")
		go spinner.Spin(60 * time.Millisecond)
Michael Yang's avatar
Michael Yang committed
190

191
192
		var latest api.GenerateResponse

Michael Yang's avatar
Michael Yang committed
193
194
195
196
197
198
		generateContext, ok := cmd.Context().Value(generateContextKey).([]int)
		if !ok {
			generateContext = []int{}
		}

		request := api.GenerateRequest{Model: model, Prompt: prompt, Context: generateContext}
199
		fn := func(resp api.GenerateResponse) error {
Michael Yang's avatar
Michael Yang committed
200
201
202
			if !spinner.IsFinished() {
				spinner.Finish()
			}
Michael Yang's avatar
Michael Yang committed
203

204
205
			latest = resp

Michael Yang's avatar
Michael Yang committed
206
			fmt.Print(resp.Response)
Michael Yang's avatar
Michael Yang committed
207
208

			cmd.SetContext(context.WithValue(cmd.Context(), generateContextKey, resp.Context))
Michael Yang's avatar
Michael Yang committed
209
			return nil
210
211
212
213
214
		}

		if err := client.Generate(context.Background(), &request, fn); err != nil {
			return err
		}
Michael Yang's avatar
Michael Yang committed
215

Michael Yang's avatar
Michael Yang committed
216
217
		fmt.Println()
		fmt.Println()
218
219
220
221
222
223
224
225
226

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

		if verbose {
			latest.Summary()
		}
Michael Yang's avatar
Michael Yang committed
227
	}
Michael Yang's avatar
Michael Yang committed
228
229
230
231

	return nil
}

232
func generateInteractive(cmd *cobra.Command, model string) error {
Michael Yang's avatar
Michael Yang committed
233
234
235
236
237
238
239
240
241
242
243
	home, err := os.UserHomeDir()
	if err != nil {
		return err
	}

	completer := readline.NewPrefixCompleter(
		readline.PcItem("/help"),
		readline.PcItem("/list"),
		readline.PcItem("/set",
			readline.PcItem("history"),
			readline.PcItem("nohistory"),
Michael Yang's avatar
Michael Yang committed
244
245
			readline.PcItem("verbose"),
			readline.PcItem("quiet"),
Michael Yang's avatar
Michael Yang committed
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
			readline.PcItem("mode",
				readline.PcItem("vim"),
				readline.PcItem("emacs"),
				readline.PcItem("default"),
			),
		),
		readline.PcItem("/exit"),
		readline.PcItem("/bye"),
	)

	usage := func() {
		fmt.Fprintln(os.Stderr, "commands:")
		fmt.Fprintln(os.Stderr, completer.Tree("  "))
	}

	config := readline.Config{
		Prompt:       ">>> ",
		HistoryFile:  filepath.Join(home, ".ollama", "history"),
		AutoComplete: completer,
	}

	scanner, err := readline.NewEx(&config)
	if err != nil {
		return err
	}
	defer scanner.Close()

	for {
		line, err := scanner.Readline()
		switch {
		case errors.Is(err, io.EOF):
			return nil
		case errors.Is(err, readline.ErrInterrupt):
279
280
281
282
			if line == "" {
				return nil
			}

Michael Yang's avatar
Michael Yang committed
283
284
			continue
		case err != nil:
Michael Yang's avatar
Michael Yang committed
285
286
287
			return err
		}

Michael Yang's avatar
Michael Yang committed
288
		line = strings.TrimSpace(line)
Michael Yang's avatar
Michael Yang committed
289

Michael Yang's avatar
Michael Yang committed
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
		switch {
		case strings.HasPrefix(line, "/list"):
			args := strings.Fields(line)
			if err := list(cmd, args[1:]); err != nil {
				return err
			}

			continue
		case strings.HasPrefix(line, "/set"):
			args := strings.Fields(line)
			if len(args) > 1 {
				switch args[1] {
				case "history":
					scanner.HistoryEnable()
					continue
				case "nohistory":
					scanner.HistoryDisable()
					continue
Michael Yang's avatar
Michael Yang committed
308
309
310
311
312
313
				case "verbose":
					cmd.Flags().Set("verbose", "true")
					continue
				case "quiet":
					cmd.Flags().Set("verbose", "false")
					continue
Michael Yang's avatar
Michael Yang committed
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
				case "mode":
					if len(args) > 2 {
						switch args[2] {
						case "vim":
							scanner.SetVimMode(true)
							continue
						case "emacs", "default":
							scanner.SetVimMode(false)
							continue
						}
					}
				}
			}
		case line == "/help", line == "/?":
			usage()
			continue
		case line == "/exit", line == "/bye":
			return nil
		}

		if err := generate(cmd, model, line); err != nil {
			return err
		}
	}
Bruce MacDonald's avatar
Bruce MacDonald committed
338
339
}

340
func generateBatch(cmd *cobra.Command, model string) error {
Michael Yang's avatar
Michael Yang committed
341
342
343
344
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		prompt := scanner.Text()
		fmt.Printf(">>> %s\n", prompt)
345
		if err := generate(cmd, model, prompt); err != nil {
Michael Yang's avatar
Michael Yang committed
346
347
348
349
350
351
352
353
			return err
		}
	}

	return nil
}

func RunServer(_ *cobra.Command, _ []string) error {
354
355
356
357
358
359
360
361
362
363
364
	host := os.Getenv("OLLAMA_HOST")
	if host == "" {
		host = "127.0.0.1"
	}

	port := os.Getenv("OLLAMA_PORT")
	if port == "" {
		port = "11434"
	}

	ln, err := net.Listen("tcp", fmt.Sprintf("%s:%s", host, port))
Jeffrey Morgan's avatar
Jeffrey Morgan committed
365
366
367
368
369
370
371
	if err != nil {
		return err
	}

	return server.Serve(ln)
}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
372
373
374
375
func NewCLI() *cobra.Command {
	log.SetFlags(log.LstdFlags | log.Lshortfile)

	rootCmd := &cobra.Command{
Michael Yang's avatar
Michael Yang committed
376
377
378
		Use:          "ollama",
		Short:        "Large language model runner",
		SilenceUsage: true,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
379
380
381
382
383
384
385
		CompletionOptions: cobra.CompletionOptions{
			DisableDefaultCmd: true,
		},
	}

	cobra.EnableCommandSorting = false

386
387
388
389
390
391
392
393
394
	createCmd := &cobra.Command{
		Use:   "create MODEL",
		Short: "Create a model from a Modelfile",
		Args:  cobra.MinimumNArgs(1),
		RunE:  create,
	}

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

Jeffrey Morgan's avatar
Jeffrey Morgan committed
395
	runCmd := &cobra.Command{
Michael Yang's avatar
Michael Yang committed
396
		Use:   "run MODEL [PROMPT]",
Jeffrey Morgan's avatar
Jeffrey Morgan committed
397
		Short: "Run a model",
Michael Yang's avatar
Michael Yang committed
398
399
		Args:  cobra.MinimumNArgs(1),
		RunE:  RunRun,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
400
401
	}

402
403
	runCmd.Flags().Bool("verbose", false, "Show timings for response")

Jeffrey Morgan's avatar
Jeffrey Morgan committed
404
405
406
407
	serveCmd := &cobra.Command{
		Use:     "serve",
		Aliases: []string{"start"},
		Short:   "Start ollama",
Michael Yang's avatar
Michael Yang committed
408
		RunE:    RunServer,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
409
410
	}

411
412
413
414
415
416
417
418
419
420
421
422
423
424
	pullCmd := &cobra.Command{
		Use:   "pull MODEL",
		Short: "Pull a model from a registry",
		Args:  cobra.MinimumNArgs(1),
		RunE:  RunPull,
	}

	pushCmd := &cobra.Command{
		Use:   "push MODEL",
		Short: "Push a model to a registry",
		Args:  cobra.MinimumNArgs(1),
		RunE:  push,
	}

Patrick Devine's avatar
Patrick Devine committed
425
426
427
428
429
430
	listCmd := &cobra.Command{
		Use:   "list",
		Short: "List models",
		RunE:  list,
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
431
432
	rootCmd.AddCommand(
		serveCmd,
433
		createCmd,
434
		runCmd,
435
436
		pullCmd,
		pushCmd,
Patrick Devine's avatar
Patrick Devine committed
437
		listCmd,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
438
439
440
441
	)

	return rootCmd
}