cmd.go 6.93 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"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
8
9
	"log"
	"net"
Michael Yang's avatar
Michael Yang committed
10
	"net/http"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
11
	"os"
Michael Yang's avatar
Michael Yang committed
12
	"strings"
Michael Yang's avatar
Michael Yang committed
13
	"time"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
14

Patrick Devine's avatar
Patrick Devine committed
15
16
	"github.com/dustin/go-humanize"
	"github.com/olekukonko/tablewriter"
Michael Yang's avatar
Michael Yang committed
17
	"github.com/schollz/progressbar/v3"
Michael Yang's avatar
Michael Yang committed
18
19
20
	"github.com/spf13/cobra"
	"golang.org/x/term"

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

26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
func create(cmd *cobra.Command, args []string) error {
	filename, _ := cmd.Flags().GetString("file")
	client := api.NewClient()

	request := api.CreateRequest{Name: args[0], Path: filename}
	fn := func(resp api.CreateProgress) error {
		fmt.Println(resp.Status)
		return nil
	}

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

	return nil
}

Michael Yang's avatar
Michael Yang committed
43
func RunRun(cmd *cobra.Command, args []string) error {
Patrick Devine's avatar
Patrick Devine committed
44
45
46
47
48
49
50
	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
51
52
53
	switch {
	case errors.Is(err, os.ErrNotExist):
		if err := pull(args[0]); err != nil {
Michael Yang's avatar
Michael Yang committed
54
55
56
57
58
59
60
61
			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
62
63
		}
	case err != nil:
Michael Yang's avatar
Michael Yang committed
64
65
66
67
		return err
	}

	return RunGenerate(cmd, args)
Bruce MacDonald's avatar
Bruce MacDonald committed
68
69
}

70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
func push(cmd *cobra.Command, args []string) error {
	client := api.NewClient()

	request := api.PushRequest{Name: args[0]}
	fn := func(resp api.PushProgress) error {
		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
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
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 {
		data = append(data, []string{m.Name, humanize.Bytes(uint64(m.Size)), format.HumanTime(m.ModifiedAt, "Never")})
	}

	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
}

113
114
115
116
func RunPull(cmd *cobra.Command, args []string) error {
	return pull(args[0])
}

Michael Yang's avatar
Michael Yang committed
117
func pull(model string) error {
Bruce MacDonald's avatar
Bruce MacDonald committed
118
	client := api.NewClient()
119

Bruce MacDonald's avatar
Bruce MacDonald committed
120
	var bar *progressbar.ProgressBar
Michael Yang's avatar
Michael Yang committed
121

122
123
124
125
126
127
	currentLayer := ""
	request := api.PullRequest{Name: model}
	fn := func(resp api.PullProgress) error {
		if resp.Digest != currentLayer && resp.Digest != "" {
			if currentLayer != "" {
				fmt.Println()
Bruce MacDonald's avatar
Bruce MacDonald committed
128
			}
129
130
131
132
133
134
135
136
137
138
139
140
141
142
			currentLayer = resp.Digest
			layerStr := resp.Digest[7:23] + "..."
			bar = progressbar.DefaultBytes(
				int64(resp.Total),
				"pulling "+layerStr,
			)
		} else if resp.Digest == currentLayer && resp.Digest != "" {
			bar.Set(resp.Completed)
		} else {
			currentLayer = ""
			fmt.Println(resp.Status)
		}
		return nil
	}
143

144
145
146
147
	if err := client.Pull(context.Background(), &request, fn); err != nil {
		return err
	}
	return nil
Michael Yang's avatar
Michael Yang committed
148
149
}

150
func RunGenerate(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
151
	if len(args) > 1 {
Michael Yang's avatar
Michael Yang committed
152
		// join all args into a single prompt
153
		return generate(cmd, args[0], strings.Join(args[1:], " "))
Michael Yang's avatar
Michael Yang committed
154
155
156
	}

	if term.IsTerminal(int(os.Stdin.Fd())) {
157
		return generateInteractive(cmd, args[0])
Michael Yang's avatar
Michael Yang committed
158
159
	}

160
	return generateBatch(cmd, args[0])
Michael Yang's avatar
Michael Yang committed
161
162
}

Michael Yang's avatar
Michael Yang committed
163
164
var generateContextKey struct{}

165
func generate(cmd *cobra.Command, model, prompt string) error {
Michael Yang's avatar
Michael Yang committed
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
	if len(strings.TrimSpace(prompt)) > 0 {
		client := api.NewClient()

		spinner := progressbar.NewOptions(-1,
			progressbar.OptionSetWriter(os.Stderr),
			progressbar.OptionThrottle(60*time.Millisecond),
			progressbar.OptionSpinnerType(14),
			progressbar.OptionSetRenderBlankState(true),
			progressbar.OptionSetElapsedTime(false),
			progressbar.OptionClearOnFinish(),
		)

		go func() {
			for range time.Tick(60 * time.Millisecond) {
				if spinner.IsFinished() {
					break
				}

				spinner.Add(1)
Michael Yang's avatar
Michael Yang committed
185
			}
Michael Yang's avatar
Michael Yang committed
186
		}()
Michael Yang's avatar
Michael Yang committed
187

188
189
		var latest api.GenerateResponse

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

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

201
202
			latest = resp

Michael Yang's avatar
Michael Yang committed
203
			fmt.Print(resp.Response)
Michael Yang's avatar
Michael Yang committed
204
205

			cmd.SetContext(context.WithValue(cmd.Context(), generateContextKey, resp.Context))
Michael Yang's avatar
Michael Yang committed
206
			return nil
207
208
209
210
211
		}

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

Michael Yang's avatar
Michael Yang committed
213
214
		fmt.Println()
		fmt.Println()
215
216
217
218
219
220
221
222
223

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

		if verbose {
			latest.Summary()
		}
Michael Yang's avatar
Michael Yang committed
224
	}
Michael Yang's avatar
Michael Yang committed
225
226
227
228

	return nil
}

229
func generateInteractive(cmd *cobra.Command, model string) error {
Michael Yang's avatar
Michael Yang committed
230
231
232
	fmt.Print(">>> ")
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
233
		if err := generate(cmd, model, scanner.Text()); err != nil {
Michael Yang's avatar
Michael Yang committed
234
235
236
237
238
239
			return err
		}

		fmt.Print(">>> ")
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
240
	return nil
Bruce MacDonald's avatar
Bruce MacDonald committed
241
242
}

243
func generateBatch(cmd *cobra.Command, model string) error {
Michael Yang's avatar
Michael Yang committed
244
245
246
247
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		prompt := scanner.Text()
		fmt.Printf(">>> %s\n", prompt)
248
		if err := generate(cmd, model, prompt); err != nil {
Michael Yang's avatar
Michael Yang committed
249
250
251
252
253
254
255
256
			return err
		}
	}

	return nil
}

func RunServer(_ *cobra.Command, _ []string) error {
257
258
259
260
261
262
263
264
265
266
267
	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
268
269
270
271
272
273
274
	if err != nil {
		return err
	}

	return server.Serve(ln)
}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
275
276
277
278
func NewCLI() *cobra.Command {
	log.SetFlags(log.LstdFlags | log.Lshortfile)

	rootCmd := &cobra.Command{
Michael Yang's avatar
Michael Yang committed
279
280
281
		Use:          "ollama",
		Short:        "Large language model runner",
		SilenceUsage: true,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
282
283
284
285
286
287
288
		CompletionOptions: cobra.CompletionOptions{
			DisableDefaultCmd: true,
		},
	}

	cobra.EnableCommandSorting = false

289
290
291
292
293
294
295
296
297
	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
298
	runCmd := &cobra.Command{
Michael Yang's avatar
Michael Yang committed
299
		Use:   "run MODEL [PROMPT]",
Jeffrey Morgan's avatar
Jeffrey Morgan committed
300
		Short: "Run a model",
Michael Yang's avatar
Michael Yang committed
301
302
		Args:  cobra.MinimumNArgs(1),
		RunE:  RunRun,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
303
304
	}

305
306
	runCmd.Flags().Bool("verbose", false, "Show timings for response")

Jeffrey Morgan's avatar
Jeffrey Morgan committed
307
308
309
310
	serveCmd := &cobra.Command{
		Use:     "serve",
		Aliases: []string{"start"},
		Short:   "Start ollama",
Michael Yang's avatar
Michael Yang committed
311
		RunE:    RunServer,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
312
313
	}

314
315
316
317
318
319
320
321
322
323
324
325
326
327
	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
328
329
330
331
332
333
	listCmd := &cobra.Command{
		Use:   "list",
		Short: "List models",
		RunE:  list,
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
334
335
	rootCmd.AddCommand(
		serveCmd,
336
		createCmd,
337
		runCmd,
338
339
		pullCmd,
		pushCmd,
Patrick Devine's avatar
Patrick Devine committed
340
		listCmd,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
341
342
343
344
	)

	return rootCmd
}