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

import (
Michael Yang's avatar
Michael Yang committed
4
	"cmp"
5
	"context"
Michael Yang's avatar
Michael Yang committed
6
	"encoding/json"
7
	"errors"
8
	"fmt"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
9
	"io"
10
	"io/fs"
11
	"log/slog"
12
	"math"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
13
14
	"net"
	"net/http"
15
	"net/netip"
16
	"os"
17
	"os/signal"
Michael Yang's avatar
Michael Yang committed
18
	"path/filepath"
19
	"strconv"
Michael Yang's avatar
Michael Yang committed
20
	"strings"
21
	"syscall"
22
	"time"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
23

Michael Yang's avatar
Michael Yang committed
24
	"github.com/gin-contrib/cors"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
25
	"github.com/gin-gonic/gin"
26
	"golang.org/x/exp/slices"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
27

28
29
30
31
	"github.com/ollama/ollama/api"
	"github.com/ollama/ollama/gpu"
	"github.com/ollama/ollama/llm"
	"github.com/ollama/ollama/openai"
Michael Yang's avatar
Michael Yang committed
32
	"github.com/ollama/ollama/types/model"
33
	"github.com/ollama/ollama/version"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
34
35
)

Michael Yang's avatar
Michael Yang committed
36
37
var mode string = gin.DebugMode

38
type Server struct {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
39
40
	addr  net.Addr
	sched *Scheduler
41
42
}

Michael Yang's avatar
Michael Yang committed
43
44
45
46
47
48
49
50
51
52
53
54
func init() {
	switch mode {
	case gin.DebugMode:
	case gin.ReleaseMode:
	case gin.TestMode:
	default:
		mode = gin.DebugMode
	}

	gin.SetMode(mode)
}

55
56
var defaultSessionDuration = 5 * time.Minute

57
58
59
60
61
62
63
64
65
66
67
func modelOptions(model *Model, requestOpts map[string]interface{}) (api.Options, error) {
	opts := api.DefaultOptions()
	if err := opts.FromMap(model.Options); err != nil {
		return api.Options{}, err
	}

	if err := opts.FromMap(requestOpts); err != nil {
		return api.Options{}, err
	}

	return opts, nil
Bruce MacDonald's avatar
Bruce MacDonald committed
68
69
}

70
71
72
73
74
75
func isSupportedImageType(image []byte) bool {
	contentType := http.DetectContentType(image)
	allowedTypes := []string{"image/jpeg", "image/jpg", "image/png"}
	return slices.Contains(allowedTypes, contentType)
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
76
func (s *Server) GenerateHandler(c *gin.Context) {
Bruce MacDonald's avatar
Bruce MacDonald committed
77
78
79

	checkpointStart := time.Now()
	var req api.GenerateRequest
Michael Yang's avatar
Michael Yang committed
80
	err := c.ShouldBindJSON(&req)
Patrick Devine's avatar
Patrick Devine committed
81

Michael Yang's avatar
Michael Yang committed
82
83
84
85
86
87
	switch {
	case errors.Is(err, io.EOF):
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "missing request body"})
		return
	case err != nil:
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
Bruce MacDonald's avatar
Bruce MacDonald committed
88
89
90
		return
	}

91
92
93
	// validate the request
	switch {
	case req.Model == "":
94
95
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "model is required"})
		return
96
97
98
	case len(req.Format) > 0 && req.Format != "json":
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "format must be json"})
		return
99
100
101
	case req.Raw && (req.Template != "" || req.System != "" || len(req.Context) > 0):
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "raw mode does not support template, system, or context"})
		return
102
103
	}

104
105
106
107
108
109
110
	for _, img := range req.Images {
		if !isSupportedImageType(img) {
			c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "unsupported image format"})
			return
		}
	}

111
	model, err := GetModel(req.Model)
Bruce MacDonald's avatar
Bruce MacDonald committed
112
	if err != nil {
113
		var pErr *fs.PathError
114
		if errors.As(err, &pErr) {
115
			c.JSON(http.StatusNotFound, gin.H{"error": fmt.Sprintf("model '%s' not found, try pulling it first", req.Model)})
116
117
118
119
120
121
			return
		}
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

122
	if model.IsEmbedding() {
123
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "embedding models do not support generate"})
124
125
126
		return
	}

127
128
129
	opts, err := modelOptions(model, req.Options)
	if err != nil {
		if errors.Is(err, api.ErrInvalidOpts) {
Bruce MacDonald's avatar
Bruce MacDonald committed
130
			c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
131
			return
132
		}
133
134
135
136
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

137
138
	var sessionDuration time.Duration
	if req.KeepAlive == nil {
139
		sessionDuration = getDefaultSessionDuration()
140
141
142
143
	} else {
		sessionDuration = req.KeepAlive.Duration
	}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
144
145
146
147
148
	rCh, eCh := s.sched.GetRunner(c.Request.Context(), model, opts, sessionDuration)
	var runner *runnerRef
	select {
	case runner = <-rCh:
	case err = <-eCh:
149
		handleErrorResponse(c, err)
Bruce MacDonald's avatar
Bruce MacDonald committed
150
151
152
		return
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
153
	// an empty request loads the model
154
155
	// note: for a short while template was used in lieu
	// of `raw` mode so we need to check for it too
Bruce MacDonald's avatar
Bruce MacDonald committed
156
	if req.Prompt == "" && req.Template == "" && req.System == "" {
157
		c.JSON(http.StatusOK, api.GenerateResponse{
158
159
			CreatedAt: time.Now().UTC(),
			Model:     req.Model,
Michael Yang's avatar
Michael Yang committed
160
161
			Done:      true,
		})
Bruce MacDonald's avatar
Bruce MacDonald committed
162
163
164
165
166
		return
	}

	checkpointLoaded := time.Now()

Bruce MacDonald's avatar
Bruce MacDonald committed
167
168
169
170
171
	var prompt string
	switch {
	case req.Raw:
		prompt = req.Prompt
	case req.Prompt != "":
172
173
		if req.Template == "" {
			req.Template = model.Template
Bruce MacDonald's avatar
Bruce MacDonald committed
174
175
		}

176
177
178
179
180
181
182
183
184
		if req.System == "" {
			req.System = model.System
		}

		slog.Debug("generate handler", "prompt", req.Prompt)
		slog.Debug("generate handler", "template", req.Template)
		slog.Debug("generate handler", "system", req.System)

		var sb strings.Builder
Michael Yang's avatar
Michael Yang committed
185
186
187
188
189
190
191
192
193
194
195
196
197
		for i := range req.Images {
			fmt.Fprintf(&sb, "[img-%d] ", i)
		}

		sb.WriteString(req.Prompt)

		p, err := Prompt(req.Template, req.System, sb.String(), "", true)
		if err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
			return
		}

		sb.Reset()
Bruce MacDonald's avatar
Bruce MacDonald committed
198
		if req.Context != nil {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
199
			prev, err := runner.llama.Detokenize(c.Request.Context(), req.Context)
Bruce MacDonald's avatar
Bruce MacDonald committed
200
201
202
203
204
			if err != nil {
				c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
				return
			}

205
			sb.WriteString(prev)
206
207
		}

208
209
210
		sb.WriteString(p)

		prompt = sb.String()
Bruce MacDonald's avatar
Bruce MacDonald committed
211
212
	}

Michael Yang's avatar
Michael Yang committed
213
	slog.Debug("generate handler", "prompt", prompt)
214

Bruce MacDonald's avatar
Bruce MacDonald committed
215
	ch := make(chan any)
Bruce MacDonald's avatar
Bruce MacDonald committed
216
	var generated strings.Builder
Bruce MacDonald's avatar
Bruce MacDonald committed
217
218
219
	go func() {
		defer close(ch)

220
		fn := func(r llm.CompletionResponse) {
Bruce MacDonald's avatar
Bruce MacDonald committed
221
222
223
224
			// Build up the full response
			if _, err := generated.WriteString(r.Content); err != nil {
				ch <- gin.H{"error": err.Error()}
				return
Bruce MacDonald's avatar
Bruce MacDonald committed
225
226
			}

Bruce MacDonald's avatar
Bruce MacDonald committed
227
			resp := api.GenerateResponse{
228
				Model:     req.Model,
229
				CreatedAt: time.Now().UTC(),
230
231
				Done:      r.Done,
				Response:  r.Content,
Bruce MacDonald's avatar
Bruce MacDonald committed
232
233
234
235
236
237
				Metrics: api.Metrics{
					PromptEvalCount:    r.PromptEvalCount,
					PromptEvalDuration: r.PromptEvalDuration,
					EvalCount:          r.EvalCount,
					EvalDuration:       r.EvalDuration,
				},
Bruce MacDonald's avatar
Bruce MacDonald committed
238
239
			}

240
241
242
243
244
			if r.Done {
				resp.TotalDuration = time.Since(checkpointStart)
				resp.LoadDuration = checkpointLoaded.Sub(checkpointStart)

				if !req.Raw {
245
					p, err := Prompt(req.Template, req.System, req.Prompt, generated.String(), false)
246
					if err != nil {
247
						c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
248
249
						return
					}
250
251

					// TODO (jmorganca): encode() should not strip special tokens
Daniel Hiltgen's avatar
Daniel Hiltgen committed
252
					tokens, err := runner.llama.Tokenize(c.Request.Context(), p)
253
254
255
256
					if err != nil {
						ch <- gin.H{"error": err.Error()}
						return
					}
257
258

					resp.Context = append(req.Context, tokens...)
Bruce MacDonald's avatar
Bruce MacDonald committed
259
260
261
262
				}
			}

			ch <- resp
Bruce MacDonald's avatar
Bruce MacDonald committed
263
264
		}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
265
		var images []llm.ImageData
Michael Yang's avatar
Michael Yang committed
266
		for i := range req.Images {
Jeffrey Morgan's avatar
Jeffrey Morgan committed
267
268
269
270
			images = append(images, llm.ImageData{
				ID:   i,
				Data: req.Images[i],
			})
Michael Yang's avatar
Michael Yang committed
271
272
		}

Bruce MacDonald's avatar
Bruce MacDonald committed
273
		// Start prediction
274
		req := llm.CompletionRequest{
275
276
			Prompt:  prompt,
			Format:  req.Format,
Michael Yang's avatar
Michael Yang committed
277
			Images:  images,
278
			Options: opts,
Bruce MacDonald's avatar
Bruce MacDonald committed
279
		}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
280
		if err := runner.llama.Completion(c.Request.Context(), req, fn); err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
281
282
283
284
285
			ch <- gin.H{"error": err.Error()}
		}
	}()

	if req.Stream != nil && !*req.Stream {
286
287
		// Accumulate responses into the final response
		var final api.GenerateResponse
Bruce MacDonald's avatar
Bruce MacDonald committed
288
		var sb strings.Builder
Bruce MacDonald's avatar
Bruce MacDonald committed
289
		for resp := range ch {
290
291
292
293
294
295
296
297
298
299
300
301
302
303
			switch r := resp.(type) {
			case api.GenerateResponse:
				sb.WriteString(r.Response)
				final = r
			case gin.H:
				if errorMsg, ok := r["error"].(string); ok {
					c.JSON(http.StatusInternalServerError, gin.H{"error": errorMsg})
					return
				} else {
					c.JSON(http.StatusInternalServerError, gin.H{"error": "unexpected error format in response"})
					return
				}
			default:
				c.JSON(http.StatusInternalServerError, gin.H{"error": "unexpected error"})
Bruce MacDonald's avatar
Bruce MacDonald committed
304
305
306
				return
			}
		}
307
308
309

		final.Response = sb.String()
		c.JSON(http.StatusOK, final)
Bruce MacDonald's avatar
Bruce MacDonald committed
310
311
312
313
314
315
		return
	}

	streamResponse(c, ch)
}

316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
func getDefaultSessionDuration() time.Duration {
	if t, exists := os.LookupEnv("OLLAMA_KEEP_ALIVE"); exists {
		v, err := strconv.Atoi(t)
		if err != nil {
			d, err := time.ParseDuration(t)
			if err != nil {
				return defaultSessionDuration
			}

			if d < 0 {
				return time.Duration(math.MaxInt64)
			}

			return d
		}

		d := time.Duration(v) * time.Second
		if d < 0 {
			return time.Duration(math.MaxInt64)
		}
		return d
	}

	return defaultSessionDuration
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
342
func (s *Server) EmbeddingsHandler(c *gin.Context) {
Bruce MacDonald's avatar
Bruce MacDonald committed
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
	var req api.EmbeddingRequest
	err := c.ShouldBindJSON(&req)
	switch {
	case errors.Is(err, io.EOF):
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "missing request body"})
		return
	case err != nil:
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	if req.Model == "" {
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "model is required"})
		return
	}

359
	model, err := GetModel(req.Model)
Bruce MacDonald's avatar
Bruce MacDonald committed
360
	if err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
361
		var pErr *fs.PathError
362
		if errors.As(err, &pErr) {
Bruce MacDonald's avatar
Bruce MacDonald committed
363
			c.JSON(http.StatusNotFound, gin.H{"error": fmt.Sprintf("model '%s' not found, try pulling it first", req.Model)})
364
365
366
367
368
369
370
371
372
			return
		}
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	opts, err := modelOptions(model, req.Options)
	if err != nil {
		if errors.Is(err, api.ErrInvalidOpts) {
Bruce MacDonald's avatar
Bruce MacDonald committed
373
			c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
374
			return
Bruce MacDonald's avatar
Bruce MacDonald committed
375
		}
376
377
378
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
379
380
381

	var sessionDuration time.Duration
	if req.KeepAlive == nil {
382
		sessionDuration = getDefaultSessionDuration()
383
384
385
386
	} else {
		sessionDuration = req.KeepAlive.Duration
	}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
387
388
389
390
391
	rCh, eCh := s.sched.GetRunner(c.Request.Context(), model, opts, sessionDuration)
	var runner *runnerRef
	select {
	case runner = <-rCh:
	case err = <-eCh:
392
		handleErrorResponse(c, err)
Bruce MacDonald's avatar
Bruce MacDonald committed
393
394
395
		return
	}

396
397
398
	// an empty request loads the model
	if req.Prompt == "" {
		c.JSON(http.StatusOK, api.EmbeddingResponse{Embedding: []float64{}})
Bruce MacDonald's avatar
Bruce MacDonald committed
399
400
401
		return
	}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
402
	embedding, err := runner.llama.Embedding(c.Request.Context(), req.Prompt)
Bruce MacDonald's avatar
Bruce MacDonald committed
403
	if err != nil {
404
		slog.Info(fmt.Sprintf("embedding generation failed: %v", err))
Bruce MacDonald's avatar
Bruce MacDonald committed
405
406
407
408
409
410
411
412
413
414
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate embedding"})
		return
	}

	resp := api.EmbeddingResponse{
		Embedding: embedding,
	}
	c.JSON(http.StatusOK, resp)
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
415
func (s *Server) PullModelHandler(c *gin.Context) {
Michael Yang's avatar
Michael Yang committed
416
	var req api.PullRequest
Michael Yang's avatar
Michael Yang committed
417
418
419
420
421
422
423
	err := c.ShouldBindJSON(&req)
	switch {
	case errors.Is(err, io.EOF):
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "missing request body"})
		return
	case err != nil:
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
Michael Yang's avatar
Michael Yang committed
424
425
426
		return
	}

Michael Yang's avatar
Michael Yang committed
427
428
429
430
431
432
433
	var model string
	if req.Model != "" {
		model = req.Model
	} else if req.Name != "" {
		model = req.Name
	} else {
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "model is required"})
434
435
436
		return
	}

437
438
439
	ch := make(chan any)
	go func() {
		defer close(ch)
440
441
		fn := func(r api.ProgressResponse) {
			ch <- r
442
		}
443

Michael Yang's avatar
Michael Yang committed
444
		regOpts := &registryOptions{
445
446
447
			Insecure: req.Insecure,
		}

448
449
450
		ctx, cancel := context.WithCancel(c.Request.Context())
		defer cancel()

Michael Yang's avatar
Michael Yang committed
451
		if err := PullModel(ctx, model, regOpts, fn); err != nil {
Michael Yang's avatar
Michael Yang committed
452
			ch <- gin.H{"error": err.Error()}
453
454
455
		}
	}()

456
457
458
459
460
	if req.Stream != nil && !*req.Stream {
		waitForStream(c, ch)
		return
	}

461
462
463
	streamResponse(c, ch)
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
464
func (s *Server) PushModelHandler(c *gin.Context) {
465
	var req api.PushRequest
Michael Yang's avatar
Michael Yang committed
466
467
468
469
470
471
472
	err := c.ShouldBindJSON(&req)
	switch {
	case errors.Is(err, io.EOF):
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "missing request body"})
		return
	case err != nil:
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
Michael Yang's avatar
Michael Yang committed
473
474
		return
	}
Michael Yang's avatar
Michael Yang committed
475

Michael Yang's avatar
Michael Yang committed
476
477
478
479
480
481
482
	var model string
	if req.Model != "" {
		model = req.Model
	} else if req.Name != "" {
		model = req.Name
	} else {
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "model is required"})
483
484
485
		return
	}

486
487
488
	ch := make(chan any)
	go func() {
		defer close(ch)
489
490
		fn := func(r api.ProgressResponse) {
			ch <- r
491
		}
492

Michael Yang's avatar
Michael Yang committed
493
		regOpts := &registryOptions{
494
495
496
			Insecure: req.Insecure,
		}

Michael Yang's avatar
Michael Yang committed
497
498
499
		ctx, cancel := context.WithCancel(c.Request.Context())
		defer cancel()

Michael Yang's avatar
Michael Yang committed
500
		if err := PushModel(ctx, model, regOpts, fn); err != nil {
Michael Yang's avatar
Michael Yang committed
501
			ch <- gin.H{"error": err.Error()}
502
503
504
		}
	}()

505
506
507
508
509
	if req.Stream != nil && !*req.Stream {
		waitForStream(c, ch)
		return
	}

510
511
512
	streamResponse(c, ch)
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
513
func (s *Server) CreateModelHandler(c *gin.Context) {
514
	var req api.CreateRequest
Michael Yang's avatar
Michael Yang committed
515
	if err := c.ShouldBindJSON(&req); errors.Is(err, io.EOF) {
Michael Yang's avatar
Michael Yang committed
516
517
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "missing request body"})
		return
Michael Yang's avatar
Michael Yang committed
518
	} else if err != nil {
Michael Yang's avatar
Michael Yang committed
519
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
Michael Yang's avatar
Michael Yang committed
520
		return
521
522
	}

Michael Yang's avatar
Michael Yang committed
523
524
525
	name := model.ParseName(cmp.Or(req.Model, req.Name))
	if !name.IsValid() {
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "invalid model name"})
526
527
528
		return
	}

Michael Yang's avatar
Michael Yang committed
529
530
	if req.Path == "" && req.Modelfile == "" {
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "path or modelfile are required"})
Michael Yang's avatar
Michael Yang committed
531
532
		return
	}
Michael Yang's avatar
Michael Yang committed
533

Michael Yang's avatar
Michael Yang committed
534
	var r io.Reader = strings.NewReader(req.Modelfile)
Michael Yang's avatar
Michael Yang committed
535
	if req.Path != "" && req.Modelfile == "" {
Michael Yang's avatar
Michael Yang committed
536
		f, err := os.Open(req.Path)
Michael Yang's avatar
Michael Yang committed
537
538
539
540
		if err != nil {
			c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": fmt.Sprintf("error reading modelfile: %s", err)})
			return
		}
Michael Yang's avatar
Michael Yang committed
541
		defer f.Close()
Michael Yang's avatar
Michael Yang committed
542

Michael Yang's avatar
Michael Yang committed
543
		r = f
Michael Yang's avatar
Michael Yang committed
544
	}
Michael Yang's avatar
Michael Yang committed
545

Michael Yang's avatar
Michael Yang committed
546
	modelfile, err := model.ParseFile(r)
Michael Yang's avatar
Michael Yang committed
547
548
549
550
551
	if err != nil {
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

Michael Yang's avatar
Michael Yang committed
552
	ch := make(chan any)
Michael Yang's avatar
Michael Yang committed
553
554
	go func() {
		defer close(ch)
555
556
		fn := func(resp api.ProgressResponse) {
			ch <- resp
557
558
		}

559
560
561
		ctx, cancel := context.WithCancel(c.Request.Context())
		defer cancel()

Michael Yang's avatar
Michael Yang committed
562
		if err := CreateModel(ctx, name.String(), filepath.Dir(req.Path), req.Quantization, modelfile, fn); err != nil {
Michael Yang's avatar
Michael Yang committed
563
			ch <- gin.H{"error": err.Error()}
564
		}
Michael Yang's avatar
Michael Yang committed
565
	}()
Michael Yang's avatar
Michael Yang committed
566

567
568
569
570
571
	if req.Stream != nil && !*req.Stream {
		waitForStream(c, ch)
		return
	}

Michael Yang's avatar
Michael Yang committed
572
	streamResponse(c, ch)
Bruce MacDonald's avatar
Bruce MacDonald committed
573
574
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
575
func (s *Server) DeleteModelHandler(c *gin.Context) {
576
	var req api.DeleteRequest
Michael Yang's avatar
Michael Yang committed
577
578
579
580
581
582
583
	err := c.ShouldBindJSON(&req)
	switch {
	case errors.Is(err, io.EOF):
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "missing request body"})
		return
	case err != nil:
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
584
585
586
		return
	}

Michael Yang's avatar
Michael Yang committed
587
588
589
590
591
592
593
	var model string
	if req.Model != "" {
		model = req.Model
	} else if req.Name != "" {
		model = req.Name
	} else {
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "model is required"})
594
595
596
		return
	}

Michael Yang's avatar
Michael Yang committed
597
	if err := DeleteModel(model); err != nil {
598
		if os.IsNotExist(err) {
Michael Yang's avatar
Michael Yang committed
599
			c.JSON(http.StatusNotFound, gin.H{"error": fmt.Sprintf("model '%s' not found", model)})
600
		} else {
601
602
			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		}
603
604
		return
	}
Michael Yang's avatar
Michael Yang committed
605
606
607
608
609
610
611
612
613
614
615
616

	manifestsPath, err := GetManifestPath()
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	if err := PruneDirectory(manifestsPath); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

617
	c.JSON(http.StatusOK, nil)
618
619
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
620
func (s *Server) ShowModelHandler(c *gin.Context) {
Patrick Devine's avatar
Patrick Devine committed
621
	var req api.ShowRequest
Michael Yang's avatar
Michael Yang committed
622
623
624
625
626
627
628
	err := c.ShouldBindJSON(&req)
	switch {
	case errors.Is(err, io.EOF):
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "missing request body"})
		return
	case err != nil:
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
Patrick Devine's avatar
Patrick Devine committed
629
630
631
		return
	}

Michael Yang's avatar
Michael Yang committed
632
	if req.Model != "" {
Michael Yang's avatar
Michael Yang committed
633
		// noop
Michael Yang's avatar
Michael Yang committed
634
	} else if req.Name != "" {
Michael Yang's avatar
Michael Yang committed
635
		req.Model = req.Name
Michael Yang's avatar
Michael Yang committed
636
	} else {
637
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "model is required"})
638
639
640
		return
	}

641
	resp, err := GetModelInfo(req)
Patrick Devine's avatar
Patrick Devine committed
642
643
	if err != nil {
		if os.IsNotExist(err) {
Michael Yang's avatar
Michael Yang committed
644
			c.JSON(http.StatusNotFound, gin.H{"error": fmt.Sprintf("model '%s' not found", req.Model)})
Patrick Devine's avatar
Patrick Devine committed
645
646
647
648
649
650
651
652
653
		} else {
			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		}
		return
	}

	c.JSON(http.StatusOK, resp)
}

654
655
func GetModelInfo(req api.ShowRequest) (*api.ShowResponse, error) {
	model, err := GetModel(req.Model)
Patrick Devine's avatar
Patrick Devine committed
656
657
658
659
	if err != nil {
		return nil, err
	}

Patrick Devine's avatar
Patrick Devine committed
660
	modelDetails := api.ModelDetails{
661
		ParentModel:       model.ParentModel,
Patrick Devine's avatar
Patrick Devine committed
662
663
664
665
666
667
668
		Format:            model.Config.ModelFormat,
		Family:            model.Config.ModelFamily,
		Families:          model.Config.ModelFamilies,
		ParameterSize:     model.Config.ModelType,
		QuantizationLevel: model.Config.FileType,
	}

669
670
671
672
673
674
675
676
	if req.System != "" {
		model.System = req.System
	}

	if req.Template != "" {
		model.Template = req.Template
	}

677
678
679
680
681
	msgs := make([]api.Message, 0)
	for _, msg := range model.Messages {
		msgs = append(msgs, api.Message{Role: msg.Role, Content: msg.Content})
	}

Patrick Devine's avatar
Patrick Devine committed
682
683
684
685
	resp := &api.ShowResponse{
		License:  strings.Join(model.License, "\n"),
		System:   model.System,
		Template: model.Template,
Patrick Devine's avatar
Patrick Devine committed
686
		Details:  modelDetails,
687
		Messages: msgs,
Patrick Devine's avatar
Patrick Devine committed
688
689
690
691
692
693
694
695
	}

	var params []string
	cs := 30
	for k, v := range model.Options {
		switch val := v.(type) {
		case []interface{}:
			for _, nv := range val {
Patrick Devine's avatar
Patrick Devine committed
696
				params = append(params, fmt.Sprintf("%-*s %#v", cs, k, nv))
Patrick Devine's avatar
Patrick Devine committed
697
			}
Patrick Devine's avatar
Patrick Devine committed
698
699
		default:
			params = append(params, fmt.Sprintf("%-*s %#v", cs, k, v))
Patrick Devine's avatar
Patrick Devine committed
700
701
702
703
		}
	}
	resp.Parameters = strings.Join(params, "\n")

704
705
706
707
708
709
	for k, v := range req.Options {
		if _, ok := req.Options[k]; ok {
			model.Options[k] = v
		}
	}

710
711
712
713
	var sb strings.Builder
	fmt.Fprintln(&sb, "# Modelfile generate by \"ollama show\"")
	fmt.Fprintln(&sb, "# To build a new Modelfile based on this, replace FROM with:")
	fmt.Fprintf(&sb, "# FROM %s\n\n", model.ShortName)
Michael Yang's avatar
Michael Yang committed
714
	fmt.Fprint(&sb, model.String())
715
	resp.Modelfile = sb.String()
716

Patrick Devine's avatar
Patrick Devine committed
717
718
719
	return resp, nil
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
720
func (s *Server) ListModelsHandler(c *gin.Context) {
721
	models := make([]api.ModelResponse, 0)
722
	manifestsPath, err := GetManifestPath()
Patrick Devine's avatar
Patrick Devine committed
723
724
725
726
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
Michael Yang's avatar
Michael Yang committed
727

Patrick Devine's avatar
Patrick Devine committed
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
	modelResponse := func(modelName string) (api.ModelResponse, error) {
		model, err := GetModel(modelName)
		if err != nil {
			return api.ModelResponse{}, err
		}

		modelDetails := api.ModelDetails{
			Format:            model.Config.ModelFormat,
			Family:            model.Config.ModelFamily,
			Families:          model.Config.ModelFamilies,
			ParameterSize:     model.Config.ModelType,
			QuantizationLevel: model.Config.FileType,
		}

		return api.ModelResponse{
Michael Yang's avatar
Michael Yang committed
743
			Model:   model.ShortName,
Patrick Devine's avatar
Patrick Devine committed
744
745
746
747
748
749
750
			Name:    model.ShortName,
			Size:    model.Size,
			Digest:  model.Digest,
			Details: modelDetails,
		}, nil
	}

Michael Yang's avatar
Michael Yang committed
751
	walkFunc := func(path string, info os.FileInfo, _ error) error {
Patrick Devine's avatar
Patrick Devine committed
752
		if !info.IsDir() {
753
754
755
756
			path, tag := filepath.Split(path)
			model := strings.Trim(strings.TrimPrefix(path, manifestsPath), string(os.PathSeparator))
			modelPath := strings.Join([]string{model, tag}, ":")
			canonicalModelPath := strings.ReplaceAll(modelPath, string(os.PathSeparator), "/")
757

758
			resp, err := modelResponse(canonicalModelPath)
Patrick Devine's avatar
Patrick Devine committed
759
			if err != nil {
760
				slog.Info(fmt.Sprintf("skipping file: %s", canonicalModelPath))
Michael Yang's avatar
Michael Yang committed
761
				// nolint: nilerr
762
				return nil
Patrick Devine's avatar
Patrick Devine committed
763
			}
Michael Yang's avatar
Michael Yang committed
764

Patrick Devine's avatar
Patrick Devine committed
765
766
			resp.ModifiedAt = info.ModTime()
			models = append(models, resp)
Patrick Devine's avatar
Patrick Devine committed
767
		}
Michael Yang's avatar
Michael Yang committed
768

Patrick Devine's avatar
Patrick Devine committed
769
		return nil
Michael Yang's avatar
Michael Yang committed
770
771
	}

772
	if err := filepath.Walk(manifestsPath, walkFunc); err != nil {
Patrick Devine's avatar
Patrick Devine committed
773
774
775
776
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

Michael Yang's avatar
Michael Yang committed
777
	c.JSON(http.StatusOK, api.ListResponse{Models: models})
Patrick Devine's avatar
Patrick Devine committed
778
779
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
780
func (s *Server) CopyModelHandler(c *gin.Context) {
Michael Yang's avatar
Michael Yang committed
781
782
	var r api.CopyRequest
	if err := c.ShouldBindJSON(&r); errors.Is(err, io.EOF) {
Michael Yang's avatar
Michael Yang committed
783
784
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "missing request body"})
		return
Michael Yang's avatar
Michael Yang committed
785
	} else if err != nil {
Michael Yang's avatar
Michael Yang committed
786
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
Patrick Devine's avatar
Patrick Devine committed
787
788
789
		return
	}

Michael Yang's avatar
Michael Yang committed
790
791
	src := model.ParseName(r.Source)
	if !src.IsValid() {
Michael Yang's avatar
Michael Yang committed
792
793
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": fmt.Sprintf("source %q is invalid", r.Source)})
		return
794
795
	}

Michael Yang's avatar
Michael Yang committed
796
797
	dst := model.ParseName(r.Destination)
	if !dst.IsValid() {
Michael Yang's avatar
Michael Yang committed
798
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": fmt.Sprintf("destination %q is invalid", r.Source)})
Patrick Devine's avatar
Patrick Devine committed
799
800
		return
	}
Michael Yang's avatar
Michael Yang committed
801
802
803
804
805
806

	if err := CopyModel(src, dst); errors.Is(err, os.ErrNotExist) {
		c.JSON(http.StatusNotFound, gin.H{"error": fmt.Sprintf("model %q not found", r.Source)})
	} else if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
	}
Patrick Devine's avatar
Patrick Devine committed
807
808
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
809
func (s *Server) HeadBlobHandler(c *gin.Context) {
Michael Yang's avatar
Michael Yang committed
810
811
812
813
814
815
816
817
818
819
820
	path, err := GetBlobsPath(c.Param("digest"))
	if err != nil {
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	if _, err := os.Stat(path); err != nil {
		c.AbortWithStatusJSON(http.StatusNotFound, gin.H{"error": fmt.Sprintf("blob %q not found", c.Param("digest"))})
		return
	}

Michael Yang's avatar
Michael Yang committed
821
	c.Status(http.StatusOK)
Michael Yang's avatar
Michael Yang committed
822
823
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
824
func (s *Server) CreateBlobHandler(c *gin.Context) {
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
	path, err := GetBlobsPath(c.Param("digest"))
	if err != nil {
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	_, err = os.Stat(path)
	switch {
	case errors.Is(err, os.ErrNotExist):
		// noop
	case err != nil:
		c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	default:
		c.Status(http.StatusOK)
		return
	}

843
	layer, err := NewLayer(c.Request.Body, "")
Michael Yang's avatar
Michael Yang committed
844
845
846
847
848
	if err != nil {
		c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

849
850
	if layer.Digest != c.Param("digest") {
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": fmt.Sprintf("digest mismatch, expected %q, got %q", c.Param("digest"), layer.Digest)})
Michael Yang's avatar
Michael Yang committed
851
852
853
		return
	}

854
	if _, err := layer.Commit(); err != nil {
Michael Yang's avatar
Michael Yang committed
855
856
857
858
		c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

Michael Yang's avatar
Michael Yang committed
859
	c.Status(http.StatusCreated)
Michael Yang's avatar
Michael Yang committed
860
861
}

Michael Yang's avatar
Michael Yang committed
862
863
864
865
866
867
var defaultAllowOrigins = []string{
	"localhost",
	"127.0.0.1",
	"0.0.0.0",
}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
func isLocalIP(ip netip.Addr) bool {
	if interfaces, err := net.Interfaces(); err == nil {
		for _, iface := range interfaces {
			addrs, err := iface.Addrs()
			if err != nil {
				continue
			}

			for _, a := range addrs {
				if parsed, _, err := net.ParseCIDR(a.String()); err == nil {
					if parsed.String() == ip.String() {
						return true
					}
				}
			}
		}
	}

	return false
}

889
func allowedHost(host string) bool {
Jeffrey Morgan's avatar
Jeffrey Morgan committed
890
	if host == "" || host == "localhost" {
891
892
893
894
895
896
897
898
		return true
	}

	if hostname, err := os.Hostname(); err == nil && host == hostname {
		return true
	}

	var tlds = []string{
Jeffrey Morgan's avatar
Jeffrey Morgan committed
899
900
901
		"localhost",
		"local",
		"internal",
902
	}
903

Jeffrey Morgan's avatar
Jeffrey Morgan committed
904
	// check if the host is a local TLD
905
906
907
908
909
910
	for _, tld := range tlds {
		if strings.HasSuffix(host, "."+tld) {
			return true
		}
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
911
	return false
Jeffrey Morgan's avatar
Jeffrey Morgan committed
912
}
913

Jeffrey Morgan's avatar
Jeffrey Morgan committed
914
915
916
func allowedHostsMiddleware(addr net.Addr) gin.HandlerFunc {
	return func(c *gin.Context) {
		if addr == nil {
917
918
919
920
			c.Next()
			return
		}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
921
		if addr, err := netip.ParseAddrPort(addr.String()); err == nil && !addr.Addr().IsLoopback() {
922
923
924
925
926
927
928
929
930
			c.Next()
			return
		}

		host, _, err := net.SplitHostPort(c.Request.Host)
		if err != nil {
			host = c.Request.Host
		}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
931
		if addr, err := netip.ParseAddr(host); err == nil {
Jeffrey Morgan's avatar
Jeffrey Morgan committed
932
			if addr.IsLoopback() || addr.IsPrivate() || addr.IsUnspecified() || isLocalIP(addr) {
Jeffrey Morgan's avatar
Jeffrey Morgan committed
933
934
935
936
937
				c.Next()
				return
			}
		}

938
939
940
941
942
943
944
		if allowedHost(host) {
			c.Next()
			return
		}

		c.AbortWithStatus(http.StatusForbidden)
	}
945
}
946

947
func (s *Server) GenerateRoutes() http.Handler {
Michael Yang's avatar
Michael Yang committed
948
949
	config := cors.DefaultConfig()
	config.AllowWildcard = true
950
	config.AllowBrowserExtensions = true
Michael Yang's avatar
Michael Yang committed
951

952
953
954
955
	if allowedOrigins := strings.Trim(os.Getenv("OLLAMA_ORIGINS"), "\"'"); allowedOrigins != "" {
		config.AllowOrigins = strings.Split(allowedOrigins, ",")
	}

Michael Yang's avatar
Michael Yang committed
956
957
958
959
960
961
962
963
	for _, allowOrigin := range defaultAllowOrigins {
		config.AllowOrigins = append(config.AllowOrigins,
			fmt.Sprintf("http://%s", allowOrigin),
			fmt.Sprintf("https://%s", allowOrigin),
			fmt.Sprintf("http://%s:*", allowOrigin),
			fmt.Sprintf("https://%s:*", allowOrigin),
		)
	}
Michael Yang's avatar
Michael Yang committed
964

Bruce MacDonald's avatar
Bruce MacDonald committed
965
	r := gin.Default()
966
967
	r.Use(
		cors.New(config),
968
		allowedHostsMiddleware(s.addr),
969
	)
Bruce MacDonald's avatar
Bruce MacDonald committed
970

Daniel Hiltgen's avatar
Daniel Hiltgen committed
971
972
973
974
975
976
977
978
979
980
981
	r.POST("/api/pull", s.PullModelHandler)
	r.POST("/api/generate", s.GenerateHandler)
	r.POST("/api/chat", s.ChatHandler)
	r.POST("/api/embeddings", s.EmbeddingsHandler)
	r.POST("/api/create", s.CreateModelHandler)
	r.POST("/api/push", s.PushModelHandler)
	r.POST("/api/copy", s.CopyModelHandler)
	r.DELETE("/api/delete", s.DeleteModelHandler)
	r.POST("/api/show", s.ShowModelHandler)
	r.POST("/api/blobs/:digest", s.CreateBlobHandler)
	r.HEAD("/api/blobs/:digest", s.HeadBlobHandler)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
982

983
	// Compatibility endpoints
Daniel Hiltgen's avatar
Daniel Hiltgen committed
984
	r.POST("/v1/chat/completions", openai.Middleware(), s.ChatHandler)
985

Michael Yang's avatar
Michael Yang committed
986
987
988
989
990
	for _, method := range []string{http.MethodGet, http.MethodHead} {
		r.Handle(method, "/", func(c *gin.Context) {
			c.String(http.StatusOK, "Ollama is running")
		})

Daniel Hiltgen's avatar
Daniel Hiltgen committed
991
		r.Handle(method, "/api/tags", s.ListModelsHandler)
Michael Yang's avatar
Michael Yang committed
992
993
994
		r.Handle(method, "/api/version", func(c *gin.Context) {
			c.JSON(http.StatusOK, gin.H{"version": version.Version})
		})
Michael Yang's avatar
Michael Yang committed
995
996
	}

997
998
999
1000
	return r
}

func Serve(ln net.Listener) error {
Michael Yang's avatar
Michael Yang committed
1001
	level := slog.LevelInfo
1002
	if debug := os.Getenv("OLLAMA_DEBUG"); debug != "" {
Michael Yang's avatar
Michael Yang committed
1003
		level = slog.LevelDebug
1004
	}
Michael Yang's avatar
Michael Yang committed
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020

	handler := slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{
		Level:     level,
		AddSource: true,
		ReplaceAttr: func(_ []string, attr slog.Attr) slog.Attr {
			if attr.Key == slog.SourceKey {
				source := attr.Value.Any().(*slog.Source)
				source.File = filepath.Base(source.File)
			}

			return attr
		},
	})

	slog.SetDefault(slog.New(handler))

1021
1022
1023
1024
1025
1026
1027
1028
	blobsDir, err := GetBlobsPath("")
	if err != nil {
		return err
	}
	if err := fixBlobs(blobsDir); err != nil {
		return err
	}

1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
	if noprune := os.Getenv("OLLAMA_NOPRUNE"); noprune == "" {
		// clean up unused layers and manifests
		if err := PruneLayers(); err != nil {
			return err
		}

		manifestsPath, err := GetManifestPath()
		if err != nil {
			return err
		}

		if err := PruneDirectory(manifestsPath); err != nil {
			return err
		}
	}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1045
1046
1047
	ctx, done := context.WithCancel(context.Background())
	sched := InitScheduler(ctx)
	s := &Server{addr: ln.Addr(), sched: sched}
1048
1049
	r := s.GenerateRoutes()

1050
	slog.Info(fmt.Sprintf("Listening on %s (version %s)", ln.Addr(), version.Version))
1051
	srvr := &http.Server{
Jeffrey Morgan's avatar
Jeffrey Morgan committed
1052
1053
1054
		Handler: r,
	}

1055
1056
	// listen for a ctrl+c and stop any loaded llm
	signals := make(chan os.Signal, 1)
1057
	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
1058
1059
	go func() {
		<-signals
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1060
1061
		done()
		sched.unloadAllRunners()
1062
		gpu.Cleanup()
1063
1064
1065
		os.Exit(0)
	}()

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1066
	if err := llm.Init(); err != nil {
1067
1068
		return fmt.Errorf("unable to initialize llm library %w", err)
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1069
1070
1071
1072
1073
1074

	s.sched.Run(ctx)

	// At startup we retrieve GPU information so we can get log messages before loading a model
	// This will log warnings to the log in case we have problems with detected GPUs
	_ = gpu.GetGPUInfo()
1075

1076
	return srvr.Serve(ln)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
1077
}
Michael Yang's avatar
Michael Yang committed
1078

1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
func waitForStream(c *gin.Context, ch chan interface{}) {
	c.Header("Content-Type", "application/json")
	for resp := range ch {
		switch r := resp.(type) {
		case api.ProgressResponse:
			if r.Status == "success" {
				c.JSON(http.StatusOK, r)
				return
			}
		case gin.H:
			if errorMsg, ok := r["error"].(string); ok {
				c.JSON(http.StatusInternalServerError, gin.H{"error": errorMsg})
				return
			} else {
				c.JSON(http.StatusInternalServerError, gin.H{"error": "unexpected error format in progress response"})
				return
			}
		default:
			c.JSON(http.StatusInternalServerError, gin.H{"error": "unexpected progress response"})
			return
		}
	}
	c.JSON(http.StatusInternalServerError, gin.H{"error": "unexpected end of progress response"})
}

Michael Yang's avatar
Michael Yang committed
1104
func streamResponse(c *gin.Context, ch chan any) {
1105
	c.Header("Content-Type", "application/x-ndjson")
Michael Yang's avatar
Michael Yang committed
1106
1107
1108
1109
1110
1111
1112
1113
	c.Stream(func(w io.Writer) bool {
		val, ok := <-ch
		if !ok {
			return false
		}

		bts, err := json.Marshal(val)
		if err != nil {
1114
			slog.Info(fmt.Sprintf("streamResponse: json.Marshal failed with %s", err))
Michael Yang's avatar
Michael Yang committed
1115
1116
1117
			return false
		}

1118
		// Delineate chunks with new-line delimiter
Michael Yang's avatar
Michael Yang committed
1119
1120
		bts = append(bts, '\n')
		if _, err := w.Write(bts); err != nil {
1121
			slog.Info(fmt.Sprintf("streamResponse: w.Write failed with %s", err))
Michael Yang's avatar
Michael Yang committed
1122
1123
1124
1125
1126
1127
			return false
		}

		return true
	})
}
Bruce MacDonald's avatar
Bruce MacDonald committed
1128

1129
// ChatPrompt builds up a prompt from a series of messages for the currently `loaded` model
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1130
func chatPrompt(ctx context.Context, runner *runnerRef, template string, messages []api.Message, numCtx int) (string, error) {
1131
	encode := func(s string) ([]int, error) {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1132
		return runner.llama.Tokenize(ctx, s)
1133
1134
	}

1135
	prompt, err := ChatPrompt(template, messages, numCtx, encode)
1136
1137
1138
1139
1140
1141
1142
	if err != nil {
		return "", err
	}

	return prompt, nil
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1143
func (s *Server) ChatHandler(c *gin.Context) {
Bruce MacDonald's avatar
Bruce MacDonald committed
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
	checkpointStart := time.Now()

	var req api.ChatRequest
	err := c.ShouldBindJSON(&req)
	switch {
	case errors.Is(err, io.EOF):
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "missing request body"})
		return
	case err != nil:
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	// validate the request
	switch {
	case req.Model == "":
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "model is required"})
		return
	case len(req.Format) > 0 && req.Format != "json":
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "format must be json"})
		return
	}

1167
	model, err := GetModel(req.Model)
Bruce MacDonald's avatar
Bruce MacDonald committed
1168
1169
	if err != nil {
		var pErr *fs.PathError
1170
		if errors.As(err, &pErr) {
Bruce MacDonald's avatar
Bruce MacDonald committed
1171
			c.JSON(http.StatusNotFound, gin.H{"error": fmt.Sprintf("model '%s' not found, try pulling it first", req.Model)})
1172
1173
1174
1175
1176
1177
			return
		}
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

1178
	if model.IsEmbedding() {
1179
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "embedding models do not support chat"})
1180
1181
1182
		return
	}

1183
1184
1185
	opts, err := modelOptions(model, req.Options)
	if err != nil {
		if errors.Is(err, api.ErrInvalidOpts) {
Bruce MacDonald's avatar
Bruce MacDonald committed
1186
			c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
1187
			return
Bruce MacDonald's avatar
Bruce MacDonald committed
1188
		}
1189
1190
1191
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
1192
1193
1194

	var sessionDuration time.Duration
	if req.KeepAlive == nil {
1195
		sessionDuration = getDefaultSessionDuration()
1196
1197
1198
1199
	} else {
		sessionDuration = req.KeepAlive.Duration
	}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1200
1201
1202
1203
1204
	rCh, eCh := s.sched.GetRunner(c.Request.Context(), model, opts, sessionDuration)
	var runner *runnerRef
	select {
	case runner = <-rCh:
	case err = <-eCh:
1205
		handleErrorResponse(c, err)
Bruce MacDonald's avatar
Bruce MacDonald committed
1206
1207
1208
1209
1210
		return
	}

	checkpointLoaded := time.Now()

1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
	// if the first message is not a system message, then add the model's default system message
	if len(req.Messages) > 0 && req.Messages[0].Role != "system" {
		req.Messages = append([]api.Message{
			{
				Role:    "system",
				Content: model.System,
			},
		}, req.Messages...)
	}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1221
	prompt, err := chatPrompt(c.Request.Context(), runner, model.Template, req.Messages, opts.NumCtx)
Bruce MacDonald's avatar
Bruce MacDonald committed
1222
1223
1224
1225
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}
Michael Yang's avatar
Michael Yang committed
1226

1227
	// an empty request loads the model
1228
	if len(req.Messages) == 0 || prompt == "" {
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
		resp := api.ChatResponse{
			CreatedAt: time.Now().UTC(),
			Model:     req.Model,
			Done:      true,
			Message:   api.Message{Role: "assistant"},
		}
		c.JSON(http.StatusOK, resp)
		return
	}

1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
	// only send images that are in the prompt
	var i int
	var images []llm.ImageData
	for _, m := range req.Messages {
		for _, img := range m.Images {
			if !isSupportedImageType(img) {
				c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "unsupported image format"})
				return
			}

			if strings.Contains(prompt, fmt.Sprintf("[img-%d]", i)) {
				images = append(images, llm.ImageData{Data: img, ID: i})
			}
			i += 1
		}
	}

	slog.Debug("chat handler", "prompt", prompt, "images", len(images))
1257

Bruce MacDonald's avatar
Bruce MacDonald committed
1258
1259
1260
1261
1262
	ch := make(chan any)

	go func() {
		defer close(ch)

1263
		fn := func(r llm.CompletionResponse) {
Bruce MacDonald's avatar
Bruce MacDonald committed
1264
1265

			resp := api.ChatResponse{
1266
				Model:     req.Model,
1267
				CreatedAt: time.Now().UTC(),
1268
				Message:   api.Message{Role: "assistant", Content: r.Content},
Bruce MacDonald's avatar
Bruce MacDonald committed
1269
1270
1271
1272
1273
1274
1275
1276
1277
				Done:      r.Done,
				Metrics: api.Metrics{
					PromptEvalCount:    r.PromptEvalCount,
					PromptEvalDuration: r.PromptEvalDuration,
					EvalCount:          r.EvalCount,
					EvalDuration:       r.EvalDuration,
				},
			}

1278
1279
1280
			if r.Done {
				resp.TotalDuration = time.Since(checkpointStart)
				resp.LoadDuration = checkpointLoaded.Sub(checkpointStart)
Bruce MacDonald's avatar
Bruce MacDonald committed
1281
1282
1283
1284
1285
			}

			ch <- resp
		}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1286
		if err := runner.llama.Completion(c.Request.Context(), llm.CompletionRequest{
1287
1288
			Prompt:  prompt,
			Format:  req.Format,
Michael Yang's avatar
Michael Yang committed
1289
			Images:  images,
1290
			Options: opts,
1291
		}, fn); err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
1292
1293
1294
1295
1296
			ch <- gin.H{"error": err.Error()}
		}
	}()

	if req.Stream != nil && !*req.Stream {
1297
1298
		// Accumulate responses into the final response
		var final api.ChatResponse
Bruce MacDonald's avatar
Bruce MacDonald committed
1299
1300
		var sb strings.Builder
		for resp := range ch {
1301
1302
			switch r := resp.(type) {
			case api.ChatResponse:
1303
				sb.WriteString(r.Message.Content)
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
				final = r
			case gin.H:
				if errorMsg, ok := r["error"].(string); ok {
					c.JSON(http.StatusInternalServerError, gin.H{"error": errorMsg})
					return
				} else {
					c.JSON(http.StatusInternalServerError, gin.H{"error": "unexpected error format in response"})
					return
				}
			default:
				c.JSON(http.StatusInternalServerError, gin.H{"error": "unexpected error"})
				return
Bruce MacDonald's avatar
Bruce MacDonald committed
1316
1317
			}
		}
1318

1319
		final.Message = api.Message{Role: "assistant", Content: sb.String()}
1320
		c.JSON(http.StatusOK, final)
Bruce MacDonald's avatar
Bruce MacDonald committed
1321
1322
1323
1324
1325
		return
	}

	streamResponse(c, ch)
}
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337

func handleErrorResponse(c *gin.Context, err error) {
	if errors.Is(err, context.Canceled) {
		c.JSON(499, gin.H{"error": "request canceled"})
		return
	}
	if errors.Is(err, ErrMaxQueue) {
		c.JSON(http.StatusServiceUnavailable, gin.H{"error": err.Error()})
		return
	}
	c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
}