cmd.go 5.94 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

Michael Yang's avatar
Michael Yang committed
15
	"github.com/schollz/progressbar/v3"
Michael Yang's avatar
Michael Yang committed
16
17
18
	"github.com/spf13/cobra"
	"golang.org/x/term"

Jeffrey Morgan's avatar
Jeffrey Morgan committed
19
20
	"github.com/jmorganca/ollama/api"
	"github.com/jmorganca/ollama/server"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
21
22
)

23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
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
40
func RunRun(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
41
42
43
44
	_, err := os.Stat(args[0])
	switch {
	case errors.Is(err, os.ErrNotExist):
		if err := pull(args[0]); err != nil {
Michael Yang's avatar
Michael Yang committed
45
46
47
48
49
50
51
52
			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
53
54
		}
	case err != nil:
Michael Yang's avatar
Michael Yang committed
55
56
57
58
		return err
	}

	return RunGenerate(cmd, args)
Bruce MacDonald's avatar
Bruce MacDonald committed
59
60
}

61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
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
}

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

Michael Yang's avatar
Michael Yang committed
80
func pull(model string) error {
Bruce MacDonald's avatar
Bruce MacDonald committed
81
	client := api.NewClient()
82

Bruce MacDonald's avatar
Bruce MacDonald committed
83
	var bar *progressbar.ProgressBar
Michael Yang's avatar
Michael Yang committed
84

85
86
87
88
89
90
	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
91
			}
92
93
94
95
96
97
98
99
100
101
102
103
104
105
			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
	}
106

107
108
109
110
	if err := client.Pull(context.Background(), &request, fn); err != nil {
		return err
	}
	return nil
Michael Yang's avatar
Michael Yang committed
111
112
}

113
func RunGenerate(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
114
	if len(args) > 1 {
Michael Yang's avatar
Michael Yang committed
115
		// join all args into a single prompt
116
		return generate(cmd, args[0], strings.Join(args[1:], " "))
Michael Yang's avatar
Michael Yang committed
117
118
119
	}

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

123
	return generateBatch(cmd, args[0])
Michael Yang's avatar
Michael Yang committed
124
125
}

Michael Yang's avatar
Michael Yang committed
126
127
var generateContextKey struct{}

128
func generate(cmd *cobra.Command, model, prompt string) error {
Michael Yang's avatar
Michael Yang committed
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
	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
148
			}
Michael Yang's avatar
Michael Yang committed
149
		}()
Michael Yang's avatar
Michael Yang committed
150

151
152
		var latest api.GenerateResponse

Michael Yang's avatar
Michael Yang committed
153
154
155
156
157
158
		generateContext, ok := cmd.Context().Value(generateContextKey).([]int)
		if !ok {
			generateContext = []int{}
		}

		request := api.GenerateRequest{Model: model, Prompt: prompt, Context: generateContext}
159
		fn := func(resp api.GenerateResponse) error {
Michael Yang's avatar
Michael Yang committed
160
161
162
			if !spinner.IsFinished() {
				spinner.Finish()
			}
Michael Yang's avatar
Michael Yang committed
163

164
165
			latest = resp

Michael Yang's avatar
Michael Yang committed
166
			fmt.Print(resp.Response)
Michael Yang's avatar
Michael Yang committed
167
168

			cmd.SetContext(context.WithValue(cmd.Context(), generateContextKey, resp.Context))
Michael Yang's avatar
Michael Yang committed
169
			return nil
170
171
172
173
174
		}

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

Michael Yang's avatar
Michael Yang committed
176
177
		fmt.Println()
		fmt.Println()
178
179
180
181
182
183
184
185
186

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

		if verbose {
			latest.Summary()
		}
Michael Yang's avatar
Michael Yang committed
187
	}
Michael Yang's avatar
Michael Yang committed
188
189
190
191

	return nil
}

192
func generateInteractive(cmd *cobra.Command, model string) error {
Michael Yang's avatar
Michael Yang committed
193
194
195
	fmt.Print(">>> ")
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
196
		if err := generate(cmd, model, scanner.Text()); err != nil {
Michael Yang's avatar
Michael Yang committed
197
198
199
200
201
202
			return err
		}

		fmt.Print(">>> ")
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
203
	return nil
Bruce MacDonald's avatar
Bruce MacDonald committed
204
205
}

206
func generateBatch(cmd *cobra.Command, model string) error {
Michael Yang's avatar
Michael Yang committed
207
208
209
210
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		prompt := scanner.Text()
		fmt.Printf(">>> %s\n", prompt)
211
		if err := generate(cmd, model, prompt); err != nil {
Michael Yang's avatar
Michael Yang committed
212
213
214
215
216
217
218
219
			return err
		}
	}

	return nil
}

func RunServer(_ *cobra.Command, _ []string) error {
220
221
222
223
224
225
226
227
228
229
230
	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
231
232
233
234
235
236
237
	if err != nil {
		return err
	}

	return server.Serve(ln)
}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
238
239
240
241
func NewCLI() *cobra.Command {
	log.SetFlags(log.LstdFlags | log.Lshortfile)

	rootCmd := &cobra.Command{
Michael Yang's avatar
Michael Yang committed
242
243
244
		Use:          "ollama",
		Short:        "Large language model runner",
		SilenceUsage: true,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
245
246
247
248
249
250
251
		CompletionOptions: cobra.CompletionOptions{
			DisableDefaultCmd: true,
		},
	}

	cobra.EnableCommandSorting = false

252
253
254
255
256
257
258
259
260
	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
261
	runCmd := &cobra.Command{
Michael Yang's avatar
Michael Yang committed
262
		Use:   "run MODEL [PROMPT]",
Jeffrey Morgan's avatar
Jeffrey Morgan committed
263
		Short: "Run a model",
Michael Yang's avatar
Michael Yang committed
264
265
		Args:  cobra.MinimumNArgs(1),
		RunE:  RunRun,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
266
267
	}

268
269
	runCmd.Flags().Bool("verbose", false, "Show timings for response")

Jeffrey Morgan's avatar
Jeffrey Morgan committed
270
271
272
273
	serveCmd := &cobra.Command{
		Use:     "serve",
		Aliases: []string{"start"},
		Short:   "Start ollama",
Michael Yang's avatar
Michael Yang committed
274
		RunE:    RunServer,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
275
276
	}

277
278
279
280
281
282
283
284
285
286
287
288
289
290
	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,
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
291
292
	rootCmd.AddCommand(
		serveCmd,
293
		createCmd,
294
		runCmd,
295
296
		pullCmd,
		pushCmd,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
297
298
299
300
	)

	return rootCmd
}