cmd.go 2.89 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"
Bruce MacDonald's avatar
Bruce MacDonald committed
6
	"fmt"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
7
8
9
10
11
	"log"
	"net"
	"os"
	"path"

Michael Yang's avatar
Michael Yang committed
12
	"github.com/schollz/progressbar/v3"
Michael Yang's avatar
Michael Yang committed
13
14
15
	"github.com/spf13/cobra"
	"golang.org/x/term"

Jeffrey Morgan's avatar
Jeffrey Morgan committed
16
17
	"github.com/jmorganca/ollama/api"
	"github.com/jmorganca/ollama/server"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
18
19
)

Bruce MacDonald's avatar
Bruce MacDonald committed
20
func cacheDir() string {
Jeffrey Morgan's avatar
Jeffrey Morgan committed
21
22
23
24
25
	home, err := os.UserHomeDir()
	if err != nil {
		panic(err)
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
26
	return path.Join(home, ".ollama")
Jeffrey Morgan's avatar
Jeffrey Morgan committed
27
28
}

Michael Yang's avatar
Michael Yang committed
29
30
31
32
33
34
35
func RunRun(cmd *cobra.Command, args []string) error {
	if err := pull(args[0]); err != nil {
		return err
	}

	fmt.Println("Up to date.")
	return RunGenerate(cmd, args)
Bruce MacDonald's avatar
Bruce MacDonald committed
36
37
}

Michael Yang's avatar
Michael Yang committed
38
func pull(model string) error {
Michael Yang's avatar
Michael Yang committed
39
	client := api.NewClient()
Michael Yang's avatar
Michael Yang committed
40
41
42
43
44
45
46
47
48
49
50
51
52

	var bar *progressbar.ProgressBar
	return client.Pull(
		context.Background(),
		&api.PullRequest{Model: model},
		func(progress api.PullProgress) error {
			if bar == nil {
				bar = progressbar.DefaultBytes(progress.Total)
			}

			return bar.Set64(progress.Completed)
		},
	)
Michael Yang's avatar
Michael Yang committed
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
}

func RunGenerate(_ *cobra.Command, args []string) error {
	if len(args) > 1 {
		return generate(args[0], args[1:]...)
	}

	if term.IsTerminal(int(os.Stdin.Fd())) {
		return generateInteractive(args[0])
	}

	return generateBatch(args[0])
}

func generate(model string, prompts ...string) error {
Michael Yang's avatar
Michael Yang committed
68
	client := api.NewClient()
Michael Yang's avatar
Michael Yang committed
69
70

	for _, prompt := range prompts {
Michael Yang's avatar
Michael Yang committed
71
		client.Generate(context.Background(), &api.GenerateRequest{Model: model, Prompt: prompt}, func(resp api.GenerateResponse) error {
Michael Yang's avatar
Michael Yang committed
72
			fmt.Print(resp.Response)
Michael Yang's avatar
Michael Yang committed
73
			return nil
Michael Yang's avatar
Michael Yang committed
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
		})
	}

	fmt.Println()
	fmt.Println()
	return nil
}

func generateInteractive(model string) error {
	fmt.Print(">>> ")
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		if err := generate(model, scanner.Text()); err != nil {
			return err
		}

		fmt.Print(">>> ")
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
93
	return nil
Bruce MacDonald's avatar
Bruce MacDonald committed
94
95
}

Michael Yang's avatar
Michael Yang committed
96
97
98
99
100
101
102
103
104
105
106
107
108
109
func generateBatch(model string) error {
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		prompt := scanner.Text()
		fmt.Printf(">>> %s\n", prompt)
		if err := generate(model, prompt); err != nil {
			return err
		}
	}

	return nil
}

func RunServer(_ *cobra.Command, _ []string) error {
Michael Yang's avatar
Michael Yang committed
110
	ln, err := net.Listen("tcp", "127.0.0.1:11434")
Jeffrey Morgan's avatar
Jeffrey Morgan committed
111
112
113
114
115
116
117
	if err != nil {
		return err
	}

	return server.Serve(ln)
}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
118
119
120
121
func NewCLI() *cobra.Command {
	log.SetFlags(log.LstdFlags | log.Lshortfile)

	rootCmd := &cobra.Command{
Michael Yang's avatar
Michael Yang committed
122
123
124
		Use:          "ollama",
		Short:        "Large language model runner",
		SilenceUsage: true,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
125
126
127
		CompletionOptions: cobra.CompletionOptions{
			DisableDefaultCmd: true,
		},
Michael Yang's avatar
Michael Yang committed
128
		PersistentPreRunE: func(_ *cobra.Command, args []string) error {
Bruce MacDonald's avatar
Bruce MacDonald committed
129
			// create the models directory and it's parent
Michael Yang's avatar
Michael Yang committed
130
			return os.MkdirAll(path.Join(cacheDir(), "models"), 0o700)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
131
132
133
134
135
136
		},
	}

	cobra.EnableCommandSorting = false

	runCmd := &cobra.Command{
Michael Yang's avatar
Michael Yang committed
137
		Use:   "run MODEL [PROMPT]",
Jeffrey Morgan's avatar
Jeffrey Morgan committed
138
		Short: "Run a model",
Michael Yang's avatar
Michael Yang committed
139
140
		Args:  cobra.MinimumNArgs(1),
		RunE:  RunRun,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
141
142
143
144
145
146
	}

	serveCmd := &cobra.Command{
		Use:     "serve",
		Aliases: []string{"start"},
		Short:   "Start ollama",
Michael Yang's avatar
Michael Yang committed
147
		RunE:    RunServer,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
148
149
150
151
	}

	rootCmd.AddCommand(
		serveCmd,
152
		runCmd,
Jeffrey Morgan's avatar
Jeffrey Morgan committed
153
154
155
156
	)

	return rootCmd
}