cmd.go 6.81 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"
12
	"path/filepath"
Michael Yang's avatar
Michael Yang committed
13
	"strings"
Michael Yang's avatar
Michael Yang committed
14
	"time"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
15

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

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"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
24
	"github.com/jmorganca/ollama/server"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
25
26
)

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

34
35
	client := api.NewClient()

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

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

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

47
48
49
50
51
52
53
		return nil
	}

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

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

58
59
60
	return nil
}

Michael Yang's avatar
Michael Yang committed
61
func RunRun(cmd *cobra.Command, args []string) error {
Patrick Devine's avatar
Patrick Devine committed
62
63
64
65
66
67
68
	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
69
70
71
	switch {
	case errors.Is(err, os.ErrNotExist):
		if err := pull(args[0]); err != nil {
Michael Yang's avatar
Michael Yang committed
72
73
74
75
76
77
78
79
			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
80
81
		}
	case err != nil:
Michael Yang's avatar
Michael Yang committed
82
83
84
85
		return err
	}

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

88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
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
}

131
132
133
134
func RunPull(cmd *cobra.Command, args []string) error {
	return pull(args[0])
}

Michael Yang's avatar
Michael Yang committed
135
func pull(model string) error {
Bruce MacDonald's avatar
Bruce MacDonald committed
136
	client := api.NewClient()
137

Bruce MacDonald's avatar
Bruce MacDonald committed
138
	var bar *progressbar.ProgressBar
Michael Yang's avatar
Michael Yang committed
139

140
141
142
143
144
145
	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
146
			}
147
148
149
150
151
152
153
154
155
156
157
158
159
160
			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
	}
161

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

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

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

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

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

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

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

190
191
		var latest api.GenerateResponse

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

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

203
204
			latest = resp

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

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

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

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

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

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

	return nil
}

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

		fmt.Print(">>> ")
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
242
	return nil
Bruce MacDonald's avatar
Bruce MacDonald committed
243
244
}

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

	return nil
}

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

	return server.Serve(ln)
}

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

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

	cobra.EnableCommandSorting = false

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

307
308
	runCmd.Flags().Bool("verbose", false, "Show timings for response")

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

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

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

	return rootCmd
}