routes.go 4.51 KB
Newer Older
Jeffrey Morgan's avatar
Jeffrey Morgan committed
1
2
3
package server

import (
Michael Yang's avatar
Michael Yang committed
4
	"encoding/json"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
5
6
7
8
	"io"
	"log"
	"net"
	"net/http"
9
	"os"
Michael Yang's avatar
Michael Yang committed
10
	"path/filepath"
Michael Yang's avatar
Michael Yang committed
11
12
	"strings"
	"text/template"
13
	"time"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
14

Michael Yang's avatar
Michael Yang committed
15
	"dario.cat/mergo"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
16
17
	"github.com/gin-gonic/gin"

Jeffrey Morgan's avatar
Jeffrey Morgan committed
18
	"github.com/jmorganca/ollama/api"
Michael Yang's avatar
Michael Yang committed
19
	"github.com/jmorganca/ollama/llama"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
20
21
)

22
23
24
25
26
27
func cacheDir() string {
	home, err := os.UserHomeDir()
	if err != nil {
		panic(err)
	}

Michael Yang's avatar
Michael Yang committed
28
	return filepath.Join(home, ".ollama")
29
30
}

Bruce MacDonald's avatar
Bruce MacDonald committed
31
func generate(c *gin.Context) {
32
33
	start := time.Now()

Michael Yang's avatar
Michael Yang committed
34
	var req api.GenerateRequest
Bruce MacDonald's avatar
Bruce MacDonald committed
35
	if err := c.ShouldBindJSON(&req); err != nil {
36
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
Bruce MacDonald's avatar
Bruce MacDonald committed
37
38
		return
	}
39

40
41
42
43
	model, err := GetModel(req.Model)
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
Bruce MacDonald's avatar
Bruce MacDonald committed
44
	}
Michael Yang's avatar
Michael Yang committed
45

Michael Yang's avatar
Michael Yang committed
46
47
48
49
50
51
52
53
54
55
56
	opts := api.DefaultOptions()
	if err := mergo.Merge(&opts, model.Options, mergo.WithOverride); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	if err := mergo.Merge(&opts, req.Options, mergo.WithOverride); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

57
58
59
60
	templ, err := template.New("").Parse(model.Prompt)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
Michael Yang's avatar
Michael Yang committed
61
62
	}

63
64
65
66
	var sb strings.Builder
	if err = templ.Execute(&sb, req); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
Michael Yang's avatar
Michael Yang committed
67
	}
68
	req.Prompt = sb.String()
Michael Yang's avatar
Michael Yang committed
69

Michael Yang's avatar
Michael Yang committed
70
	log.Printf("prompt: \n%s", req.Prompt)
71

Michael Yang's avatar
Michael Yang committed
72
	llm, err := llama.New(model.ModelPath, opts)
Michael Yang's avatar
Michael Yang committed
73
74
75
76
77
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	defer llm.Close()
Jeffrey Morgan's avatar
Jeffrey Morgan committed
78

Michael Yang's avatar
Michael Yang committed
79
80
81
82
83
84
85
86
87
88
89
90
91
	ch := make(chan any)
	go func() {
		defer close(ch)
		llm.Predict(req.Context, req.Prompt, func(r api.GenerateResponse) {
			r.Model = req.Model
			r.CreatedAt = time.Now().UTC()
			if r.Done {
				r.TotalDuration = time.Since(start)
			}

			ch <- r
		})
	}()
Michael Yang's avatar
Michael Yang committed
92

Michael Yang's avatar
Michael Yang committed
93
	streamResponse(c, ch)
Michael Yang's avatar
Michael Yang committed
94
}
Michael Yang's avatar
Michael Yang committed
95

Michael Yang's avatar
Michael Yang committed
96
97
98
99
100
101
102
func pull(c *gin.Context) {
	var req api.PullRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

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
	ch := make(chan any)
	go func() {
		defer close(ch)
		fn := func(status, digest string, total, completed int, percent float64) {
			ch <- api.PullProgress{
				Status:    status,
				Digest:    digest,
				Total:     total,
				Completed: completed,
				Percent:   percent,
			}
		}
		if err := PullModel(req.Name, req.Username, req.Password, fn); err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
			return
		}
	}()

	streamResponse(c, ch)
}

func push(c *gin.Context) {
	var req api.PushRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
Michael Yang's avatar
Michael Yang committed
128
129
		return
	}
Michael Yang's avatar
Michael Yang committed
130

131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
	ch := make(chan any)
	go func() {
		defer close(ch)
		fn := func(status, digest string, total, completed int, percent float64) {
			ch <- api.PushProgress{
				Status:    status,
				Digest:    digest,
				Total:     total,
				Completed: completed,
				Percent:   percent,
			}
		}
		if err := PushModel(req.Name, req.Username, req.Password, fn); err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
			return
		}
	}()

	streamResponse(c, ch)
}

func create(c *gin.Context) {
	var req api.CreateRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"message": err.Error()})
Michael Yang's avatar
Michael Yang committed
156
		return
157
158
159
	}

	// NOTE consider passing the entire Modelfile in the json instead of the path to it
Michael Yang's avatar
Michael Yang committed
160

161
162
163
	file, err := os.Open(req.Path)
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"message": err.Error()})
Michael Yang's avatar
Michael Yang committed
164
165
		return
	}
166
	defer file.Close()
Michael Yang's avatar
Michael Yang committed
167

Michael Yang's avatar
Michael Yang committed
168
	ch := make(chan any)
Michael Yang's avatar
Michael Yang committed
169
170
	go func() {
		defer close(ch)
171
172
173
		fn := func(status string) {
			ch <- api.CreateProgress{
				Status: status,
Michael Yang's avatar
Michael Yang committed
174
			}
175
176
177
178
179
180
		}

		if err := CreateModel(req.Name, file, fn); err != nil {
			c.JSON(http.StatusBadRequest, gin.H{"message": err.Error()})
			return
		}
Michael Yang's avatar
Michael Yang committed
181
	}()
Michael Yang's avatar
Michael Yang committed
182

Michael Yang's avatar
Michael Yang committed
183
	streamResponse(c, ch)
Bruce MacDonald's avatar
Bruce MacDonald committed
184
185
186
187
188
}

func Serve(ln net.Listener) error {
	r := gin.Default()

189
190
191
192
	r.GET("/", func(c *gin.Context) {
		c.String(http.StatusOK, "Ollama is running")
	})

Michael Yang's avatar
Michael Yang committed
193
	r.POST("/api/pull", pull)
Bruce MacDonald's avatar
Bruce MacDonald committed
194
	r.POST("/api/generate", generate)
195
196
	r.POST("/api/create", create)
	r.POST("/api/push", push)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
197
198
199
200
201
202
203
204

	log.Printf("Listening on %s", ln.Addr())
	s := &http.Server{
		Handler: r,
	}

	return s.Serve(ln)
}
Michael Yang's avatar
Michael Yang committed
205

Michael Yang's avatar
Michael Yang committed
206
func streamResponse(c *gin.Context, ch chan any) {
Michael Yang's avatar
Michael Yang committed
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
	c.Stream(func(w io.Writer) bool {
		val, ok := <-ch
		if !ok {
			return false
		}

		bts, err := json.Marshal(val)
		if err != nil {
			return false
		}

		bts = append(bts, '\n')
		if _, err := w.Write(bts); err != nil {
			return false
		}

		return true
	})
}