routes.go 7.71 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"
5
	"errors"
6
	"fmt"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
7
8
9
10
	"io"
	"log"
	"net"
	"net/http"
11
	"os"
Michael Yang's avatar
Michael Yang committed
12
	"path/filepath"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
13
	"reflect"
Michael Yang's avatar
Michael Yang committed
14
	"strings"
Michael Yang's avatar
Michael Yang committed
15
	"sync"
16
	"time"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
17

Michael Yang's avatar
Michael Yang committed
18
	"github.com/gin-contrib/cors"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
19
20
	"github.com/gin-gonic/gin"

Jeffrey Morgan's avatar
Jeffrey Morgan committed
21
	"github.com/jmorganca/ollama/api"
Michael Yang's avatar
Michael Yang committed
22
	"github.com/jmorganca/ollama/llama"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
23
24
)

Jeffrey Morgan's avatar
Jeffrey Morgan committed
25
var loaded struct {
Michael Yang's avatar
Michael Yang committed
26
27
28
29
30
31
	mu sync.Mutex

	llm *llama.LLM

	expireAt    time.Time
	expireTimer *time.Timer
Jeffrey Morgan's avatar
Jeffrey Morgan committed
32

33
34
	digest  string
	options api.Options
Michael Yang's avatar
Michael Yang committed
35
36
}

37
func GenerateHandler(c *gin.Context) {
Jeffrey Morgan's avatar
Jeffrey Morgan committed
38
39
	loaded.mu.Lock()
	defer loaded.mu.Unlock()
Michael Yang's avatar
Michael Yang committed
40

Michael Yang's avatar
Michael Yang committed
41
	checkpointStart := time.Now()
42

Michael Yang's avatar
Michael Yang committed
43
	var req api.GenerateRequest
Bruce MacDonald's avatar
Bruce MacDonald committed
44
	if err := c.ShouldBindJSON(&req); err != nil {
45
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
Bruce MacDonald's avatar
Bruce MacDonald committed
46
47
		return
	}
48

49
50
51
52
	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
53
	}
Michael Yang's avatar
Michael Yang committed
54

55
56
57
58
59
60
61
62
63
64
65
66
67
68
	opts := api.DefaultOptions()
	if err := opts.FromMap(model.Options); err != nil {
		log.Printf("could not load model options: %v", err)
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	if err := opts.FromMap(req.Options); err != nil {
		log.Printf("could not merge model options: %v", err)
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	if model.Digest != loaded.digest || !reflect.DeepEqual(loaded.options, opts) {
Jeffrey Morgan's avatar
Jeffrey Morgan committed
69
70
71
		if loaded.llm != nil {
			loaded.llm.Close()
			loaded.llm = nil
72
			loaded.digest = ""
Michael Yang's avatar
Michael Yang committed
73
		}
Michael Yang's avatar
Michael Yang committed
74

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

Jeffrey Morgan's avatar
Jeffrey Morgan committed
81
82
		loaded.llm = llm
		loaded.digest = model.Digest
83
		loaded.options = opts
Michael Yang's avatar
Michael Yang committed
84
85
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
86
	sessionDuration := 5 * time.Minute
Michael Yang's avatar
Michael Yang committed
87

Jeffrey Morgan's avatar
Jeffrey Morgan committed
88
89
90
91
92
	loaded.expireAt = time.Now().Add(sessionDuration)
	if loaded.expireTimer == nil {
		loaded.expireTimer = time.AfterFunc(sessionDuration, func() {
			loaded.mu.Lock()
			defer loaded.mu.Unlock()
Michael Yang's avatar
Michael Yang committed
93

Jeffrey Morgan's avatar
Jeffrey Morgan committed
94
			if time.Now().Before(loaded.expireAt) {
Michael Yang's avatar
Michael Yang committed
95
96
97
				return
			}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
98
			if loaded.llm == nil {
Michael Yang's avatar
Michael Yang committed
99
100
101
				return
			}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
102
103
			loaded.llm.Close()
			loaded.llm = nil
104
			loaded.digest = ""
Michael Yang's avatar
Michael Yang committed
105
		})
Michael Yang's avatar
Michael Yang committed
106
	}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
107
	loaded.expireTimer.Reset(sessionDuration)
Michael Yang's avatar
Michael Yang committed
108

Michael Yang's avatar
Michael Yang committed
109
110
	checkpointLoaded := time.Now()

Michael Yang's avatar
Michael Yang committed
111
	prompt, err := model.Prompt(req)
Michael Yang's avatar
Michael Yang committed
112
113
114
115
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
116

Michael Yang's avatar
Michael Yang committed
117
118
119
	ch := make(chan any)
	go func() {
		defer close(ch)
Michael Yang's avatar
Michael Yang committed
120
		fn := func(r api.GenerateResponse) {
Jeffrey Morgan's avatar
Jeffrey Morgan committed
121
122
			loaded.expireAt = time.Now().Add(sessionDuration)
			loaded.expireTimer.Reset(sessionDuration)
Michael Yang's avatar
Michael Yang committed
123

Michael Yang's avatar
Michael Yang committed
124
125
126
			r.Model = req.Model
			r.CreatedAt = time.Now().UTC()
			if r.Done {
Michael Yang's avatar
Michael Yang committed
127
128
				r.TotalDuration = time.Since(checkpointStart)
				r.LoadDuration = checkpointLoaded.Sub(checkpointStart)
Michael Yang's avatar
Michael Yang committed
129
130
131
			}

			ch <- r
Michael Yang's avatar
Michael Yang committed
132
133
		}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
134
		if err := loaded.llm.Predict(req.Context, prompt, fn); err != nil {
Michael Yang's avatar
Michael Yang committed
135
136
			ch <- gin.H{"error": err.Error()}
		}
Michael Yang's avatar
Michael Yang committed
137
	}()
Michael Yang's avatar
Michael Yang committed
138

Michael Yang's avatar
Michael Yang committed
139
	streamResponse(c, ch)
Michael Yang's avatar
Michael Yang committed
140
}
Michael Yang's avatar
Michael Yang committed
141

142
func PullModelHandler(c *gin.Context) {
Michael Yang's avatar
Michael Yang committed
143
144
145
146
147
148
	var req api.PullRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

149
150
151
	ch := make(chan any)
	go func() {
		defer close(ch)
152
153
		fn := func(r api.ProgressResponse) {
			ch <- r
154
		}
155

156
157
158
159
160
161
162
		regOpts := &RegistryOptions{
			Insecure: req.Insecure,
			Username: req.Username,
			Password: req.Password,
		}

		if err := PullModel(req.Name, regOpts, fn); err != nil {
Michael Yang's avatar
Michael Yang committed
163
			ch <- gin.H{"error": err.Error()}
164
165
166
167
168
169
		}
	}()

	streamResponse(c, ch)
}

170
func PushModelHandler(c *gin.Context) {
171
172
173
	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
174
175
		return
	}
Michael Yang's avatar
Michael Yang committed
176

177
178
179
	ch := make(chan any)
	go func() {
		defer close(ch)
180
181
		fn := func(r api.ProgressResponse) {
			ch <- r
182
		}
183

184
185
186
187
188
189
190
		regOpts := &RegistryOptions{
			Insecure: req.Insecure,
			Username: req.Username,
			Password: req.Password,
		}

		if err := PushModel(req.Name, regOpts, fn); err != nil {
Michael Yang's avatar
Michael Yang committed
191
			ch <- gin.H{"error": err.Error()}
192
193
194
195
196
197
		}
	}()

	streamResponse(c, ch)
}

198
func CreateModelHandler(c *gin.Context) {
199
200
201
	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
202
		return
203
204
	}

Michael Yang's avatar
Michael Yang committed
205
	ch := make(chan any)
Michael Yang's avatar
Michael Yang committed
206
207
	go func() {
		defer close(ch)
208
209
		fn := func(resp api.ProgressResponse) {
			ch <- resp
210
211
		}

212
		if err := CreateModel(req.Name, req.Path, fn); err != nil {
Michael Yang's avatar
Michael Yang committed
213
			ch <- gin.H{"error": err.Error()}
214
		}
Michael Yang's avatar
Michael Yang committed
215
	}()
Michael Yang's avatar
Michael Yang committed
216

Michael Yang's avatar
Michael Yang committed
217
	streamResponse(c, ch)
Bruce MacDonald's avatar
Bruce MacDonald committed
218
219
}

220
221
222
223
224
225
226
func DeleteModelHandler(c *gin.Context) {
	var req api.DeleteRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

227
228
229
230
	if err := DeleteModel(req.Name); err != nil {
		if os.IsNotExist(err) {
			c.JSON(http.StatusNotFound, gin.H{"error": fmt.Sprintf("model '%s' not found", req.Name)})
		} else {
231
232
			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		}
233
234
		return
	}
235
236
237
}

func ListModelsHandler(c *gin.Context) {
Patrick Devine's avatar
Patrick Devine committed
238
239
240
241
242
243
244
245
	var models []api.ListResponseModel
	fp, err := GetManifestPath()
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
	err = filepath.Walk(fp, func(path string, info os.FileInfo, err error) error {
		if err != nil {
246
247
248
249
			if errors.Is(err, os.ErrNotExist) {
				log.Printf("manifest file does not exist: %s", fp)
				return nil
			}
Patrick Devine's avatar
Patrick Devine committed
250
251
252
253
254
			return err
		}
		if !info.IsDir() {
			fi, err := os.Stat(path)
			if err != nil {
255
256
				log.Printf("skipping file: %s", fp)
				return nil
Patrick Devine's avatar
Patrick Devine committed
257
258
259
260
261
262
263
264
265
266
			}
			path := path[len(fp)+1:]
			slashIndex := strings.LastIndex(path, "/")
			if slashIndex == -1 {
				return nil
			}
			tag := path[:slashIndex] + ":" + path[slashIndex+1:]
			mp := ParseModelPath(tag)
			manifest, err := GetManifest(mp)
			if err != nil {
267
268
				log.Printf("skipping file: %s", fp)
				return nil
Patrick Devine's avatar
Patrick Devine committed
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
			}
			model := api.ListResponseModel{
				Name:       mp.GetShortTagname(),
				Size:       manifest.GetTotalSize(),
				ModifiedAt: fi.ModTime(),
			}
			models = append(models, model)
		}
		return nil
	})
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

Michael Yang's avatar
Michael Yang committed
284
	c.JSON(http.StatusOK, api.ListResponse{Models: models})
Patrick Devine's avatar
Patrick Devine committed
285
286
}

Patrick Devine's avatar
Patrick Devine committed
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
func CopyModelHandler(c *gin.Context) {
	var req api.CopyRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	if err := CopyModel(req.Source, req.Destination); err != nil {
		if os.IsNotExist(err) {
			c.JSON(http.StatusNotFound, gin.H{"error": fmt.Sprintf("model '%s' not found", req.Source)})
		} else {
			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		}
		return
	}
}

Bruce MacDonald's avatar
Bruce MacDonald committed
304
func Serve(ln net.Listener) error {
Michael Yang's avatar
Michael Yang committed
305
306
307
308
309
310
311
312
313
314
315
316
317
318
	config := cors.DefaultConfig()
	config.AllowWildcard = true
	// only allow http/https from localhost
	config.AllowOrigins = []string{
		"http://localhost",
		"http://localhost:*",
		"https://localhost",
		"https://localhost:*",
		"http://127.0.0.1",
		"http://127.0.0.1:*",
		"https://127.0.0.1",
		"https://127.0.0.1:*",
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
319
	r := gin.Default()
Michael Yang's avatar
Michael Yang committed
320
	r.Use(cors.New(config))
Bruce MacDonald's avatar
Bruce MacDonald committed
321

322
323
324
	r.GET("/", func(c *gin.Context) {
		c.String(http.StatusOK, "Ollama is running")
	})
Bruce MacDonald's avatar
Bruce MacDonald committed
325
326
327
	r.HEAD("/", func(c *gin.Context) {
		c.Status(http.StatusOK)
	})
328

329
330
331
332
	r.POST("/api/pull", PullModelHandler)
	r.POST("/api/generate", GenerateHandler)
	r.POST("/api/create", CreateModelHandler)
	r.POST("/api/push", PushModelHandler)
Patrick Devine's avatar
Patrick Devine committed
333
	r.POST("/api/copy", CopyModelHandler)
334
335
	r.GET("/api/tags", ListModelsHandler)
	r.DELETE("/api/delete", DeleteModelHandler)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
336
337
338
339
340
341
342
343

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

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

Michael Yang's avatar
Michael Yang committed
345
func streamResponse(c *gin.Context, ch chan any) {
Michael Yang's avatar
Michael Yang committed
346
347
348
349
350
351
352
353
	c.Stream(func(w io.Writer) bool {
		val, ok := <-ch
		if !ok {
			return false
		}

		bts, err := json.Marshal(val)
		if err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
354
			log.Printf("streamResponse: json.Marshal failed with %s", err)
Michael Yang's avatar
Michael Yang committed
355
356
357
358
359
			return false
		}

		bts = append(bts, '\n')
		if _, err := w.Write(bts); err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
360
			log.Printf("streamResponse: w.Write failed with %s", err)
Michael Yang's avatar
Michael Yang committed
361
362
363
364
365
366
			return false
		}

		return true
	})
}