routes.go 34.6 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
	"slices"
20
	"strconv"
Michael Yang's avatar
Michael Yang committed
21
	"strings"
22
	"syscall"
23
	"time"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
24

Michael Yang's avatar
Michael Yang committed
25
	"github.com/gin-contrib/cors"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
26
27
	"github.com/gin-gonic/gin"

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

Michael Yang's avatar
Michael Yang committed
39
40
var mode string = gin.DebugMode

41
type Server struct {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
42
43
	addr  net.Addr
	sched *Scheduler
44
45
}

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

	gin.SetMode(mode)
}

58
59
var defaultSessionDuration = 5 * time.Minute

60
61
62
63
64
65
66
67
68
69
70
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
71
72
}

73
74
75
76
77
78
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
79
func (s *Server) GenerateHandler(c *gin.Context) {
Bruce MacDonald's avatar
Bruce MacDonald committed
80
81
	checkpointStart := time.Now()
	var req api.GenerateRequest
Michael Yang's avatar
Michael Yang committed
82
	err := c.ShouldBindJSON(&req)
Patrick Devine's avatar
Patrick Devine committed
83

Michael Yang's avatar
Michael Yang committed
84
85
86
87
88
89
	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
90
91
92
		return
	}

93
94
95
	// validate the request
	switch {
	case req.Model == "":
96
97
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "model is required"})
		return
98
99
100
	case len(req.Format) > 0 && req.Format != "json":
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "format must be json"})
		return
101
102
103
	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
104
105
	}

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

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

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

129
130
131
132
133
134
	opts, err := modelOptions(model, req.Options)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

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

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

Bruce MacDonald's avatar
Bruce MacDonald committed
151
	// an empty request loads the model
152
153
	// 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
154
	if req.Prompt == "" && req.Template == "" && req.System == "" {
155
		c.JSON(http.StatusOK, api.GenerateResponse{
156
157
158
159
			CreatedAt:  time.Now().UTC(),
			Model:      req.Model,
			Done:       true,
			DoneReason: "load",
Michael Yang's avatar
Michael Yang committed
160
		})
Bruce MacDonald's avatar
Bruce MacDonald committed
161
162
163
164
165
		return
	}

	checkpointLoaded := time.Now()

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

175
176
177
178
179
180
181
182
183
		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
184
185
186
187
188
189
190
191
192
193
194
195
196
		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
197
		if req.Context != nil {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
198
			prev, err := runner.llama.Detokenize(c.Request.Context(), req.Context)
Bruce MacDonald's avatar
Bruce MacDonald committed
199
200
201
202
203
			if err != nil {
				c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
				return
			}

204
			sb.WriteString(prev)
205
206
		}

207
208
209
		sb.WriteString(p)

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

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

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

219
		fn := func(r llm.CompletionResponse) {
Bruce MacDonald's avatar
Bruce MacDonald committed
220
221
222
223
			// 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
224
225
			}

Bruce MacDonald's avatar
Bruce MacDonald committed
226
			resp := api.GenerateResponse{
227
228
229
230
231
				Model:      req.Model,
				CreatedAt:  time.Now().UTC(),
				Done:       r.Done,
				Response:   r.Content,
				DoneReason: r.DoneReason,
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
func getDefaultSessionDuration() time.Duration {
317
318
	if envconfig.KeepAlive != "" {
		v, err := strconv.Atoi(envconfig.KeepAlive)
319
		if err != nil {
320
			d, err := time.ParseDuration(envconfig.KeepAlive)
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
			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
373
374
			return
		}
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

	opts, err := modelOptions(model, req.Options)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
375
376
377

	var sessionDuration time.Duration
	if req.KeepAlive == nil {
378
		sessionDuration = getDefaultSessionDuration()
379
380
381
382
	} else {
		sessionDuration = req.KeepAlive.Duration
	}

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

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

Daniel Hiltgen's avatar
Daniel Hiltgen committed
398
	embedding, err := runner.llama.Embedding(c.Request.Context(), req.Prompt)
Bruce MacDonald's avatar
Bruce MacDonald committed
399
	if err != nil {
400
		slog.Info(fmt.Sprintf("embedding generation failed: %v", err))
Bruce MacDonald's avatar
Bruce MacDonald committed
401
402
403
404
405
406
407
408
409
410
		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
411
func (s *Server) PullModelHandler(c *gin.Context) {
Michael Yang's avatar
Michael Yang committed
412
	var req api.PullRequest
Michael Yang's avatar
Michael Yang committed
413
414
415
416
417
418
419
	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
420
421
422
		return
	}

423
424
425
426
427
428
429
430
	name := model.ParseName(cmp.Or(req.Model, req.Name))
	if !name.IsValid() {
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "invalid model name"})
		return
	}

	if err := checkNameExists(name); err != nil {
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
431
432
433
		return
	}

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

Michael Yang's avatar
Michael Yang committed
441
		regOpts := &registryOptions{
442
443
444
			Insecure: req.Insecure,
		}

445
446
447
		ctx, cancel := context.WithCancel(c.Request.Context())
		defer cancel()

448
		if err := PullModel(ctx, name.DisplayShortest(), regOpts, fn); err != nil {
Michael Yang's avatar
Michael Yang committed
449
			ch <- gin.H{"error": err.Error()}
450
451
452
		}
	}()

453
454
455
456
457
	if req.Stream != nil && !*req.Stream {
		waitForStream(c, ch)
		return
	}

458
459
460
	streamResponse(c, ch)
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
461
func (s *Server) PushModelHandler(c *gin.Context) {
462
	var req api.PushRequest
Michael Yang's avatar
Michael Yang committed
463
464
465
466
467
468
469
	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
470
471
		return
	}
Michael Yang's avatar
Michael Yang committed
472

Michael Yang's avatar
Michael Yang committed
473
474
475
476
477
478
479
	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"})
480
481
482
		return
	}

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

Michael Yang's avatar
Michael Yang committed
490
		regOpts := &registryOptions{
491
492
493
			Insecure: req.Insecure,
		}

Michael Yang's avatar
Michael Yang committed
494
495
496
		ctx, cancel := context.WithCancel(c.Request.Context())
		defer cancel()

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

502
503
504
505
506
	if req.Stream != nil && !*req.Stream {
		waitForStream(c, ch)
		return
	}

507
508
509
	streamResponse(c, ch)
}

510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
func checkNameExists(name model.Name) error {
	names, err := Manifests()
	if err != nil {
		return err
	}

	for n := range names {
		if strings.EqualFold(n.Filepath(), name.Filepath()) && n != name {
			return fmt.Errorf("a model with that name already exists")
		}
	}

	return nil
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
525
func (s *Server) CreateModelHandler(c *gin.Context) {
526
527
	var r api.CreateRequest
	if err := c.ShouldBindJSON(&r); errors.Is(err, io.EOF) {
Michael Yang's avatar
Michael Yang committed
528
529
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "missing request body"})
		return
Michael Yang's avatar
Michael Yang committed
530
	} else if err != nil {
Michael Yang's avatar
Michael Yang committed
531
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
Michael Yang's avatar
Michael Yang committed
532
		return
533
534
	}

535
	name := model.ParseName(cmp.Or(r.Model, r.Name))
Michael Yang's avatar
Michael Yang committed
536
	if !name.IsValid() {
537
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": errtypes.InvalidModelNameErrMsg})
538
539
540
		return
	}

541
542
543
544
545
	if err := checkNameExists(name); err != nil {
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

546
	if r.Path == "" && r.Modelfile == "" {
Michael Yang's avatar
Michael Yang committed
547
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "path or modelfile are required"})
Michael Yang's avatar
Michael Yang committed
548
549
		return
	}
Michael Yang's avatar
Michael Yang committed
550

551
552
553
	var sr io.Reader = strings.NewReader(r.Modelfile)
	if r.Path != "" && r.Modelfile == "" {
		f, err := os.Open(r.Path)
Michael Yang's avatar
Michael Yang committed
554
555
556
557
		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
558
		defer f.Close()
Michael Yang's avatar
Michael Yang committed
559

560
		sr = f
Michael Yang's avatar
Michael Yang committed
561
	}
Michael Yang's avatar
Michael Yang committed
562

563
	f, err := parser.ParseFile(sr)
Michael Yang's avatar
Michael Yang committed
564
565
566
567
568
	if err != nil {
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

Michael Yang's avatar
Michael Yang committed
569
	ch := make(chan any)
Michael Yang's avatar
Michael Yang committed
570
571
	go func() {
		defer close(ch)
572
573
		fn := func(resp api.ProgressResponse) {
			ch <- resp
574
575
		}

576
577
578
		ctx, cancel := context.WithCancel(c.Request.Context())
		defer cancel()

579
580
		quantization := cmp.Or(r.Quantize, r.Quantization)
		if err := CreateModel(ctx, name, filepath.Dir(r.Path), strings.ToUpper(quantization), f, fn); err != nil {
Michael Yang's avatar
Michael Yang committed
581
			ch <- gin.H{"error": err.Error()}
582
		}
Michael Yang's avatar
Michael Yang committed
583
	}()
Michael Yang's avatar
Michael Yang committed
584

585
	if r.Stream != nil && !*r.Stream {
586
587
588
589
		waitForStream(c, ch)
		return
	}

Michael Yang's avatar
Michael Yang committed
590
	streamResponse(c, ch)
Bruce MacDonald's avatar
Bruce MacDonald committed
591
592
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
593
func (s *Server) DeleteModelHandler(c *gin.Context) {
594
595
	var r api.DeleteRequest
	if err := c.ShouldBindJSON(&r); errors.Is(err, io.EOF) {
Michael Yang's avatar
Michael Yang committed
596
597
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "missing request body"})
		return
598
	} else if err != nil {
Michael Yang's avatar
Michael Yang committed
599
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
600
601
602
		return
	}

603
604
605
	n := model.ParseName(cmp.Or(r.Model, r.Name))
	if !n.IsValid() {
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": fmt.Sprintf("name %q is invalid", cmp.Or(r.Model, r.Name))})
606
607
		return
	}
Michael Yang's avatar
Michael Yang committed
608

609
	m, err := ParseNamedManifest(n)
Michael Yang's avatar
Michael Yang committed
610
611
612
613
614
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

615
	if err := m.Remove(); err != nil {
Michael Yang's avatar
Michael Yang committed
616
617
618
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
619
620
621
622
623

	if err := m.RemoveLayers(); err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
624
625
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
626
func (s *Server) ShowModelHandler(c *gin.Context) {
Patrick Devine's avatar
Patrick Devine committed
627
	var req api.ShowRequest
Michael Yang's avatar
Michael Yang committed
628
629
630
631
632
633
634
	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
635
636
637
		return
	}

Michael Yang's avatar
Michael Yang committed
638
	if req.Model != "" {
Michael Yang's avatar
Michael Yang committed
639
		// noop
Michael Yang's avatar
Michael Yang committed
640
	} else if req.Name != "" {
Michael Yang's avatar
Michael Yang committed
641
		req.Model = req.Name
Michael Yang's avatar
Michael Yang committed
642
	} else {
643
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "model is required"})
644
645
646
		return
	}

647
	resp, err := GetModelInfo(req)
Patrick Devine's avatar
Patrick Devine committed
648
	if err != nil {
649
650
		switch {
		case os.IsNotExist(err):
Michael Yang's avatar
Michael Yang committed
651
			c.JSON(http.StatusNotFound, gin.H{"error": fmt.Sprintf("model '%s' not found", req.Model)})
652
653
654
		case err.Error() == "invalid model name":
			c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		default:
Patrick Devine's avatar
Patrick Devine committed
655
656
657
658
659
660
661
662
			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		}
		return
	}

	c.JSON(http.StatusOK, resp)
}

663
func GetModelInfo(req api.ShowRequest) (*api.ShowResponse, error) {
664
	m, err := GetModel(req.Model)
Patrick Devine's avatar
Patrick Devine committed
665
666
667
668
	if err != nil {
		return nil, err
	}

Patrick Devine's avatar
Patrick Devine committed
669
	modelDetails := api.ModelDetails{
670
671
672
673
674
675
		ParentModel:       m.ParentModel,
		Format:            m.Config.ModelFormat,
		Family:            m.Config.ModelFamily,
		Families:          m.Config.ModelFamilies,
		ParameterSize:     m.Config.ModelType,
		QuantizationLevel: m.Config.FileType,
Patrick Devine's avatar
Patrick Devine committed
676
677
	}

678
	if req.System != "" {
679
		m.System = req.System
680
681
682
	}

	if req.Template != "" {
683
		m.Template = req.Template
684
685
	}

686
	msgs := make([]api.Message, 0)
687
	for _, msg := range m.Messages {
688
689
690
		msgs = append(msgs, api.Message{Role: msg.Role, Content: msg.Content})
	}

691
692
693
694
695
696
697
698
699
700
	n := model.ParseName(req.Model)
	if !n.IsValid() {
		return nil, fmt.Errorf("invalid model name")
	}

	manifest, err := ParseNamedManifest(n)
	if err != nil {
		return nil, err
	}

Patrick Devine's avatar
Patrick Devine committed
701
	resp := &api.ShowResponse{
702
703
704
705
706
707
		License:    strings.Join(m.License, "\n"),
		System:     m.System,
		Template:   m.Template,
		Details:    modelDetails,
		Messages:   msgs,
		ModifiedAt: manifest.fi.ModTime(),
Patrick Devine's avatar
Patrick Devine committed
708
709
710
711
	}

	var params []string
	cs := 30
712
	for k, v := range m.Options {
Patrick Devine's avatar
Patrick Devine committed
713
714
715
		switch val := v.(type) {
		case []interface{}:
			for _, nv := range val {
Patrick Devine's avatar
Patrick Devine committed
716
				params = append(params, fmt.Sprintf("%-*s %#v", cs, k, nv))
Patrick Devine's avatar
Patrick Devine committed
717
			}
Patrick Devine's avatar
Patrick Devine committed
718
719
		default:
			params = append(params, fmt.Sprintf("%-*s %#v", cs, k, v))
Patrick Devine's avatar
Patrick Devine committed
720
721
722
723
		}
	}
	resp.Parameters = strings.Join(params, "\n")

724
725
	for k, v := range req.Options {
		if _, ok := req.Options[k]; ok {
726
			m.Options[k] = v
727
728
729
		}
	}

730
	var sb strings.Builder
731
	fmt.Fprintln(&sb, "# Modelfile generated by \"ollama show\"")
732
	fmt.Fprintln(&sb, "# To build a new Modelfile based on this, replace FROM with:")
733
734
	fmt.Fprintf(&sb, "# FROM %s\n\n", m.ShortName)
	fmt.Fprint(&sb, m.String())
735
	resp.Modelfile = sb.String()
736

737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
	kvData, err := getKVData(m.ModelPath, req.Verbose)
	if err != nil {
		return nil, err
	}
	delete(kvData, "general.name")
	delete(kvData, "tokenizer.chat_template")
	resp.ModelInfo = kvData

	if len(m.ProjectorPaths) > 0 {
		projectorData, err := getKVData(m.ProjectorPaths[0], req.Verbose)
		if err != nil {
			return nil, err
		}
		resp.ProjectorInfo = projectorData
	}

Patrick Devine's avatar
Patrick Devine committed
753
754
755
	return resp, nil
}

756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
func getKVData(digest string, verbose bool) (llm.KV, error) {
	kvData, err := llm.LoadModel(digest)
	if err != nil {
		return nil, err
	}

	kv := kvData.KV()

	if !verbose {
		for k := range kv {
			if t, ok := kv[k].([]any); len(t) > 5 && ok {
				kv[k] = []any{}
			}
		}
	}

	return kv, nil
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
775
func (s *Server) ListModelsHandler(c *gin.Context) {
776
	ms, err := Manifests()
Patrick Devine's avatar
Patrick Devine committed
777
778
779
780
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
Michael Yang's avatar
Michael Yang committed
781

782
	models := []api.ListModelResponse{}
783
784
785
786
787
788
789
	for n, m := range ms {
		f, err := m.Config.Open()
		if err != nil {
			slog.Warn("bad manifest filepath", "name", n, "error", err)
			continue
		}
		defer f.Close()
790

791
792
793
794
		var cf ConfigV2
		if err := json.NewDecoder(f).Decode(&cf); err != nil {
			slog.Warn("bad manifest config", "name", n, "error", err)
			continue
Patrick Devine's avatar
Patrick Devine committed
795
		}
Michael Yang's avatar
Michael Yang committed
796

797
		// tag should never be masked
798
		models = append(models, api.ListModelResponse{
799
800
801
802
803
804
805
806
807
808
809
810
811
			Model:      n.DisplayShortest(),
			Name:       n.DisplayShortest(),
			Size:       m.Size(),
			Digest:     m.digest,
			ModifiedAt: m.fi.ModTime(),
			Details: api.ModelDetails{
				Format:            cf.ModelFormat,
				Family:            cf.ModelFamily,
				Families:          cf.ModelFamilies,
				ParameterSize:     cf.ModelType,
				QuantizationLevel: cf.FileType,
			},
		})
Patrick Devine's avatar
Patrick Devine committed
812
813
	}

814
	slices.SortStableFunc(models, func(i, j api.ListModelResponse) int {
815
816
817
818
		// most recently modified first
		return cmp.Compare(j.ModifiedAt.Unix(), i.ModifiedAt.Unix())
	})

Michael Yang's avatar
Michael Yang committed
819
	c.JSON(http.StatusOK, api.ListResponse{Models: models})
Patrick Devine's avatar
Patrick Devine committed
820
821
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
822
func (s *Server) CopyModelHandler(c *gin.Context) {
Michael Yang's avatar
Michael Yang committed
823
824
	var r api.CopyRequest
	if err := c.ShouldBindJSON(&r); errors.Is(err, io.EOF) {
Michael Yang's avatar
Michael Yang committed
825
826
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "missing request body"})
		return
Michael Yang's avatar
Michael Yang committed
827
	} else if err != nil {
Michael Yang's avatar
Michael Yang committed
828
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
Patrick Devine's avatar
Patrick Devine committed
829
830
831
		return
	}

Michael Yang's avatar
Michael Yang committed
832
833
	src := model.ParseName(r.Source)
	if !src.IsValid() {
Michael Yang's avatar
Michael Yang committed
834
835
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": fmt.Sprintf("source %q is invalid", r.Source)})
		return
836
837
	}

Michael Yang's avatar
Michael Yang committed
838
839
	dst := model.ParseName(r.Destination)
	if !dst.IsValid() {
840
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": fmt.Sprintf("destination %q is invalid", r.Destination)})
Patrick Devine's avatar
Patrick Devine committed
841
842
		return
	}
Michael Yang's avatar
Michael Yang committed
843

844
845
846
847
848
	if err := checkNameExists(dst); err != nil {
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

Michael Yang's avatar
Michael Yang committed
849
850
851
852
853
	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
854
855
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
856
func (s *Server) HeadBlobHandler(c *gin.Context) {
Michael Yang's avatar
Michael Yang committed
857
858
859
860
861
862
863
864
865
866
867
	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
868
	c.Status(http.StatusOK)
Michael Yang's avatar
Michael Yang committed
869
870
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
871
func (s *Server) CreateBlobHandler(c *gin.Context) {
Michael Yang's avatar
Michael Yang committed
872
873
	if ib, ok := intermediateBlobs[c.Param("digest")]; ok {
		p, err := GetBlobsPath(ib)
874
875
876
877
878
879
		if err != nil {
			c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
			return
		}

		if _, err := os.Stat(p); errors.Is(err, os.ErrNotExist) {
Michael Yang's avatar
Michael Yang committed
880
881
			slog.Info("evicting intermediate blob which no longer exists", "digest", ib)
			delete(intermediateBlobs, c.Param("digest"))
882
883
884
885
886
887
888
889
890
		} else if err != nil {
			c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
			return
		} else {
			c.Status(http.StatusOK)
			return
		}
	}

891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
	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
	}

909
	layer, err := NewLayer(c.Request.Body, "")
Michael Yang's avatar
Michael Yang committed
910
911
912
913
914
	if err != nil {
		c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

915
916
	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
917
918
919
		return
	}

Michael Yang's avatar
Michael Yang committed
920
	c.Status(http.StatusCreated)
Michael Yang's avatar
Michael Yang committed
921
922
}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
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
}

944
func allowedHost(host string) bool {
Jeffrey Morgan's avatar
Jeffrey Morgan committed
945
	if host == "" || host == "localhost" {
946
947
948
949
950
951
952
953
		return true
	}

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

	var tlds = []string{
Jeffrey Morgan's avatar
Jeffrey Morgan committed
954
955
956
		"localhost",
		"local",
		"internal",
957
	}
958

Jeffrey Morgan's avatar
Jeffrey Morgan committed
959
	// check if the host is a local TLD
960
961
962
963
964
965
	for _, tld := range tlds {
		if strings.HasSuffix(host, "."+tld) {
			return true
		}
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
966
	return false
Jeffrey Morgan's avatar
Jeffrey Morgan committed
967
}
968

Jeffrey Morgan's avatar
Jeffrey Morgan committed
969
970
971
func allowedHostsMiddleware(addr net.Addr) gin.HandlerFunc {
	return func(c *gin.Context) {
		if addr == nil {
972
973
974
975
			c.Next()
			return
		}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
976
		if addr, err := netip.ParseAddrPort(addr.String()); err == nil && !addr.Addr().IsLoopback() {
977
978
979
980
981
982
983
984
985
			c.Next()
			return
		}

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

Jeffrey Morgan's avatar
Jeffrey Morgan committed
986
		if addr, err := netip.ParseAddr(host); err == nil {
Jeffrey Morgan's avatar
Jeffrey Morgan committed
987
			if addr.IsLoopback() || addr.IsPrivate() || addr.IsUnspecified() || isLocalIP(addr) {
Jeffrey Morgan's avatar
Jeffrey Morgan committed
988
989
990
991
992
				c.Next()
				return
			}
		}

993
		if allowedHost(host) {
Michael Yang's avatar
lint  
Michael Yang committed
994
			if c.Request.Method == http.MethodOptions {
995
996
997
998
				c.AbortWithStatus(http.StatusNoContent)
				return
			}

999
1000
1001
1002
1003
1004
			c.Next()
			return
		}

		c.AbortWithStatus(http.StatusForbidden)
	}
1005
}
1006

1007
func (s *Server) GenerateRoutes() http.Handler {
Michael Yang's avatar
Michael Yang committed
1008
1009
	config := cors.DefaultConfig()
	config.AllowWildcard = true
1010
	config.AllowBrowserExtensions = true
1011
	config.AllowHeaders = []string{"Authorization", "Content-Type", "User-Agent", "Accept", "X-Requested-With"}
royjhan's avatar
royjhan committed
1012
1013
1014
1015
	openAIProperties := []string{"lang", "package-version", "os", "arch", "runtime", "runtime-version", "async"}
	for _, prop := range openAIProperties {
		config.AllowHeaders = append(config.AllowHeaders, "x-stainless-"+prop)
	}
1016
	config.AllowOrigins = envconfig.AllowOrigins
Michael Yang's avatar
Michael Yang committed
1017

Bruce MacDonald's avatar
Bruce MacDonald committed
1018
	r := gin.Default()
1019
1020
	r.Use(
		cors.New(config),
1021
		allowedHostsMiddleware(s.addr),
1022
	)
Bruce MacDonald's avatar
Bruce MacDonald committed
1023

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
	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)
1035
	r.GET("/api/ps", s.ProcessHandler)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
1036

1037
	// Compatibility endpoints
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1038
	r.POST("/v1/chat/completions", openai.Middleware(), s.ChatHandler)
1039

Michael Yang's avatar
Michael Yang committed
1040
1041
1042
1043
1044
	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
1045
		r.Handle(method, "/api/tags", s.ListModelsHandler)
Michael Yang's avatar
Michael Yang committed
1046
1047
1048
		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
1049
1050
	}

1051
1052
1053
1054
	return r
}

func Serve(ln net.Listener) error {
Michael Yang's avatar
Michael Yang committed
1055
	level := slog.LevelInfo
1056
	if envconfig.Debug {
Michael Yang's avatar
Michael Yang committed
1057
		level = slog.LevelDebug
1058
	}
Michael Yang's avatar
Michael Yang committed
1059

1060
	slog.Info("server config", "env", envconfig.Values())
Michael Yang's avatar
Michael Yang committed
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
	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))

1076
1077
1078
1079
1080
1081
1082
1083
	blobsDir, err := GetBlobsPath("")
	if err != nil {
		return err
	}
	if err := fixBlobs(blobsDir); err != nil {
		return err
	}

1084
	if !envconfig.NoPrune {
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
		// 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
1100
	ctx, done := context.WithCancel(context.Background())
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1101
1102
	schedCtx, schedDone := context.WithCancel(ctx)
	sched := InitScheduler(schedCtx)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1103
	s := &Server{addr: ln.Addr(), sched: sched}
1104
1105
	r := s.GenerateRoutes()

1106
	slog.Info(fmt.Sprintf("Listening on %s (version %s)", ln.Addr(), version.Version))
1107
	srvr := &http.Server{
Jeffrey Morgan's avatar
Jeffrey Morgan committed
1108
1109
1110
		Handler: r,
	}

1111
1112
	// listen for a ctrl+c and stop any loaded llm
	signals := make(chan os.Signal, 1)
1113
	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
1114
1115
	go func() {
		<-signals
1116
		srvr.Close()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1117
		schedDone()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1118
		sched.unloadAllRunners()
1119
		gpu.Cleanup()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1120
		done()
1121
1122
	}()

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1123
	if err := llm.Init(); err != nil {
1124
1125
		return fmt.Errorf("unable to initialize llm library %w", err)
	}
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1126

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1127
	s.sched.Run(schedCtx)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1128
1129
1130

	// 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
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1131
1132
	gpus := gpu.GetGPUInfo()
	gpus.LogDetails()
1133

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1134
1135
1136
1137
1138
1139
1140
	err = srvr.Serve(ln)
	// If server is closed from the signal handler, wait for the ctx to be done
	// otherwise error out quickly
	if !errors.Is(err, http.ErrServerClosed) {
		return err
	}
	<-ctx.Done()
1141
	return nil
Jeffrey Morgan's avatar
Jeffrey Morgan committed
1142
}
Michael Yang's avatar
Michael Yang committed
1143

1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
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
1169
func streamResponse(c *gin.Context, ch chan any) {
1170
	c.Header("Content-Type", "application/x-ndjson")
Michael Yang's avatar
Michael Yang committed
1171
1172
1173
1174
1175
1176
1177
1178
	c.Stream(func(w io.Writer) bool {
		val, ok := <-ch
		if !ok {
			return false
		}

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

1183
		// Delineate chunks with new-line delimiter
Michael Yang's avatar
Michael Yang committed
1184
1185
		bts = append(bts, '\n')
		if _, err := w.Write(bts); err != nil {
1186
			slog.Info(fmt.Sprintf("streamResponse: w.Write failed with %s", err))
Michael Yang's avatar
Michael Yang committed
1187
1188
1189
1190
1191
1192
			return false
		}

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

1194
func (s *Server) ProcessHandler(c *gin.Context) {
1195
	models := []api.ProcessModelResponse{}
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206

	for _, v := range s.sched.loaded {
		model := v.model
		modelDetails := api.ModelDetails{
			Format:            model.Config.ModelFormat,
			Family:            model.Config.ModelFamily,
			Families:          model.Config.ModelFamilies,
			ParameterSize:     model.Config.ModelType,
			QuantizationLevel: model.Config.FileType,
		}

1207
		mr := api.ProcessModelResponse{
1208
1209
1210
1211
1212
1213
1214
1215
			Model:     model.ShortName,
			Name:      model.ShortName,
			Size:      int64(v.estimatedTotal),
			SizeVRAM:  int64(v.estimatedVRAM),
			Digest:    model.Digest,
			Details:   modelDetails,
			ExpiresAt: v.expiresAt,
		}
1216
1217
1218
1219
1220
1221
1222
1223
		// The scheduler waits to set expiresAt, so if a model is loading it's
		// possible that it will be set to the unix epoch. For those cases, just
		// calculate the time w/ the sessionDuration instead.
		var epoch time.Time
		if v.expiresAt == epoch {
			mr.ExpiresAt = time.Now().Add(v.sessionDuration)
		}

1224
1225
1226
		models = append(models, mr)
	}

1227
	c.JSON(http.StatusOK, api.ProcessResponse{Models: models})
1228
1229
}

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

1236
	prompt, err := ChatPrompt(template, messages, numCtx, encode)
1237
1238
1239
1240
1241
1242
1243
	if err != nil {
		return "", err
	}

	return prompt, nil
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1244
func (s *Server) ChatHandler(c *gin.Context) {
Bruce MacDonald's avatar
Bruce MacDonald committed
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
	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
	}

1268
	model, err := GetModel(req.Model)
Bruce MacDonald's avatar
Bruce MacDonald committed
1269
1270
	if err != nil {
		var pErr *fs.PathError
1271
		if errors.As(err, &pErr) {
Bruce MacDonald's avatar
Bruce MacDonald committed
1272
			c.JSON(http.StatusNotFound, gin.H{"error": fmt.Sprintf("model '%s' not found, try pulling it first", req.Model)})
1273
1274
1275
1276
1277
1278
			return
		}
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

1279
	if model.IsEmbedding() {
1280
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "embedding models do not support chat"})
1281
1282
1283
		return
	}

1284
1285
1286
1287
1288
	opts, err := modelOptions(model, req.Options)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
1289
1290
1291

	var sessionDuration time.Duration
	if req.KeepAlive == nil {
1292
		sessionDuration = getDefaultSessionDuration()
1293
1294
1295
1296
	} else {
		sessionDuration = req.KeepAlive.Duration
	}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1297
1298
1299
1300
1301
	rCh, eCh := s.sched.GetRunner(c.Request.Context(), model, opts, sessionDuration)
	var runner *runnerRef
	select {
	case runner = <-rCh:
	case err = <-eCh:
1302
		handleErrorResponse(c, err)
Bruce MacDonald's avatar
Bruce MacDonald committed
1303
1304
1305
1306
1307
		return
	}

	checkpointLoaded := time.Now()

1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
	// 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
1318
	prompt, err := chatPrompt(c.Request.Context(), runner, model.Template, req.Messages, opts.NumCtx)
Bruce MacDonald's avatar
Bruce MacDonald committed
1319
1320
1321
1322
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}
Michael Yang's avatar
Michael Yang committed
1323

1324
	// an empty request loads the model
1325
	if len(req.Messages) == 0 || prompt == "" {
1326
		resp := api.ChatResponse{
1327
1328
1329
1330
1331
			CreatedAt:  time.Now().UTC(),
			Model:      req.Model,
			Done:       true,
			DoneReason: "load",
			Message:    api.Message{Role: "assistant"},
1332
1333
1334
1335
1336
		}
		c.JSON(http.StatusOK, resp)
		return
	}

1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
	// 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))
1355

Bruce MacDonald's avatar
Bruce MacDonald committed
1356
1357
1358
1359
1360
	ch := make(chan any)

	go func() {
		defer close(ch)

1361
		fn := func(r llm.CompletionResponse) {
Bruce MacDonald's avatar
Bruce MacDonald committed
1362
			resp := api.ChatResponse{
1363
1364
1365
1366
1367
				Model:      req.Model,
				CreatedAt:  time.Now().UTC(),
				Message:    api.Message{Role: "assistant", Content: r.Content},
				Done:       r.Done,
				DoneReason: r.DoneReason,
Bruce MacDonald's avatar
Bruce MacDonald committed
1368
1369
1370
1371
1372
1373
1374
1375
				Metrics: api.Metrics{
					PromptEvalCount:    r.PromptEvalCount,
					PromptEvalDuration: r.PromptEvalDuration,
					EvalCount:          r.EvalCount,
					EvalDuration:       r.EvalDuration,
				},
			}

1376
1377
1378
			if r.Done {
				resp.TotalDuration = time.Since(checkpointStart)
				resp.LoadDuration = checkpointLoaded.Sub(checkpointStart)
Bruce MacDonald's avatar
Bruce MacDonald committed
1379
1380
1381
1382
1383
			}

			ch <- resp
		}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1384
		if err := runner.llama.Completion(c.Request.Context(), llm.CompletionRequest{
1385
1386
			Prompt:  prompt,
			Format:  req.Format,
Michael Yang's avatar
Michael Yang committed
1387
			Images:  images,
1388
			Options: opts,
1389
		}, fn); err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
1390
1391
1392
1393
1394
			ch <- gin.H{"error": err.Error()}
		}
	}()

	if req.Stream != nil && !*req.Stream {
1395
1396
		// Accumulate responses into the final response
		var final api.ChatResponse
Bruce MacDonald's avatar
Bruce MacDonald committed
1397
1398
		var sb strings.Builder
		for resp := range ch {
1399
1400
			switch r := resp.(type) {
			case api.ChatResponse:
1401
				sb.WriteString(r.Message.Content)
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
				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
1414
1415
			}
		}
1416

1417
		final.Message = api.Message{Role: "assistant", Content: sb.String()}
1418
		c.JSON(http.StatusOK, final)
Bruce MacDonald's avatar
Bruce MacDonald committed
1419
1420
1421
1422
1423
		return
	}

	streamResponse(c, ch)
}
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435

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()})
}