cmd.go 5.73 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
func create(cmd *cobra.Command, args []string) error {
	filename, _ := cmd.Flags().GetString("file")
	client := api.NewClient()

Michael Yang's avatar
Michael Yang committed
27
28
	var spinner *Spinner

29
30
	request := api.CreateRequest{Name: args[0], Path: filename}
	fn := func(resp api.CreateProgress) error {
Michael Yang's avatar
Michael Yang committed
31
32
33
34
35
36
37
		if spinner != nil {
			spinner.Stop()
		}

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

38
39
40
41
42
43
44
		return nil
	}

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

Michael Yang's avatar
Michael Yang committed
45
46
47
48
	if spinner != nil {
		spinner.Stop()
	}

49
50
51
	return nil
}

Michael Yang's avatar
Michael Yang committed
52
func RunRun(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
53
54
55
56
	_, 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
57
58
59
60
61
62
63
64
			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
65
66
		}
	case err != nil:
Michael Yang's avatar
Michael Yang committed
67
68
69
70
		return err
	}

	return RunGenerate(cmd, args)
Bruce MacDonald's avatar
Bruce MacDonald committed
71
72
}

73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
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
92
func pull(model string) error {
Bruce MacDonald's avatar
Bruce MacDonald committed
93
	client := api.NewClient()
94

Bruce MacDonald's avatar
Bruce MacDonald committed
95
	var bar *progressbar.ProgressBar
Michael Yang's avatar
Michael Yang committed
96

97
98
99
100
101
102
	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
103
			}
104
105
106
107
108
109
110
111
112
113
114
115
116
117
			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
	}
118

119
120
121
122
	if err := client.Pull(context.Background(), &request, fn); err != nil {
		return err
	}
	return nil
Michael Yang's avatar
Michael Yang committed
123
124
}

125
func RunGenerate(cmd *cobra.Command, args []string) error {
Michael Yang's avatar
Michael Yang committed
126
	if len(args) > 1 {
Michael Yang's avatar
Michael Yang committed
127
		// join all args into a single prompt
128
		return generate(cmd, args[0], strings.Join(args[1:], " "))
Michael Yang's avatar
Michael Yang committed
129
130
131
	}

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

135
	return generateBatch(cmd, args[0])
Michael Yang's avatar
Michael Yang committed
136
137
}

Michael Yang's avatar
Michael Yang committed
138
139
var generateContextKey struct{}

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

Michael Yang's avatar
Michael Yang committed
144
145
		spinner := NewSpinner("")
		go spinner.Spin(60 * time.Millisecond)
Michael Yang's avatar
Michael Yang committed
146

147
148
		var latest api.GenerateResponse

Michael Yang's avatar
Michael Yang committed
149
150
151
152
153
154
		generateContext, ok := cmd.Context().Value(generateContextKey).([]int)
		if !ok {
			generateContext = []int{}
		}

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

160
161
			latest = resp

Michael Yang's avatar
Michael Yang committed
162
			fmt.Print(resp.Response)
Michael Yang's avatar
Michael Yang committed
163
164

			cmd.SetContext(context.WithValue(cmd.Context(), generateContextKey, resp.Context))
Michael Yang's avatar
Michael Yang committed
165
			return nil
166
167
168
169
170
		}

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

Michael Yang's avatar
Michael Yang committed
172
173
		fmt.Println()
		fmt.Println()
174
175
176
177
178
179
180
181
182

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

		if verbose {
			latest.Summary()
		}
Michael Yang's avatar
Michael Yang committed
183
	}
Michael Yang's avatar
Michael Yang committed
184
185
186
187

	return nil
}

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

		fmt.Print(">>> ")
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
199
	return nil
Bruce MacDonald's avatar
Bruce MacDonald committed
200
201
}

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

	return nil
}

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

	return server.Serve(ln)
}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
234
235
236
237
func NewCLI() *cobra.Command {
	log.SetFlags(log.LstdFlags | log.Lshortfile)

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

	cobra.EnableCommandSorting = false

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

264
265
	runCmd.Flags().Bool("verbose", false, "Show timings for response")

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

273
274
275
276
277
278
279
280
281
282
283
284
285
286
	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
287
288
	rootCmd.AddCommand(
		serveCmd,
289
		createCmd,
290
		runCmd,
291
292
		pullCmd,
		pushCmd,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
293
294
295
296
	)

	return rootCmd
}