"official/vision/tasks/image_classification.py" did not exist on "51113cc5f7f46e0c37407fd22e1a3dfaacb92383"
routes.go 32.9 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"
32
	"github.com/ollama/ollama/server/envconfig"
33
	"github.com/ollama/ollama/types/errtypes"
Michael Yang's avatar
Michael Yang committed
34
	"github.com/ollama/ollama/types/model"
35
	"github.com/ollama/ollama/version"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
36
37
)

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

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

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

	gin.SetMode(mode)
}

57
58
var defaultSessionDuration = 5 * time.Minute

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

72
73
74
75
76
77
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
78
func (s *Server) GenerateHandler(c *gin.Context) {
Bruce MacDonald's avatar
Bruce MacDonald committed
79
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
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
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
	}

Michael Yang's avatar
Michael Yang committed
423
424
425
426
427
428
429
	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"})
430
431
432
		return
	}

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

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

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

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

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

457
458
459
	streamResponse(c, ch)
}

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

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

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

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

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

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

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

506
507
508
	streamResponse(c, ch)
}

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

Michael Yang's avatar
Michael Yang committed
519
520
	name := model.ParseName(cmp.Or(req.Model, req.Name))
	if !name.IsValid() {
521
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": errtypes.InvalidModelNameErrMsg})
522
523
524
		return
	}

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

Michael Yang's avatar
Michael Yang committed
530
	var r io.Reader = strings.NewReader(req.Modelfile)
Michael Yang's avatar
Michael Yang committed
531
	if req.Path != "" && req.Modelfile == "" {
Michael Yang's avatar
Michael Yang committed
532
		f, err := os.Open(req.Path)
Michael Yang's avatar
Michael Yang committed
533
534
535
536
		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
537
		defer f.Close()
Michael Yang's avatar
Michael Yang committed
538

Michael Yang's avatar
Michael Yang committed
539
		r = f
Michael Yang's avatar
Michael Yang committed
540
	}
Michael Yang's avatar
Michael Yang committed
541

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

Michael Yang's avatar
Michael Yang committed
548
	ch := make(chan any)
Michael Yang's avatar
Michael Yang committed
549
550
	go func() {
		defer close(ch)
551
552
		fn := func(resp api.ProgressResponse) {
			ch <- resp
553
554
		}

555
556
557
		ctx, cancel := context.WithCancel(c.Request.Context())
		defer cancel()

558
559
560
561
562
563
		quantization := req.Quantization
		if req.Quantize != "" {
			quantization = req.Quantize
		}

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

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

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

Daniel Hiltgen's avatar
Daniel Hiltgen committed
576
func (s *Server) DeleteModelHandler(c *gin.Context) {
577
	var req api.DeleteRequest
Michael Yang's avatar
Michael Yang committed
578
579
580
581
582
583
584
	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()})
585
586
587
		return
	}

Michael Yang's avatar
Michael Yang committed
588
589
590
591
592
593
594
	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"})
595
596
597
		return
	}

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

	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
	}

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

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

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

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

	c.JSON(http.StatusOK, resp)
}

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

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

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

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

678
679
680
681
682
	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
683
684
685
686
	resp := &api.ShowResponse{
		License:  strings.Join(model.License, "\n"),
		System:   model.System,
		Template: model.Template,
Patrick Devine's avatar
Patrick Devine committed
687
		Details:  modelDetails,
688
		Messages: msgs,
Patrick Devine's avatar
Patrick Devine committed
689
690
691
692
693
694
695
696
	}

	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
697
				params = append(params, fmt.Sprintf("%-*s %#v", cs, k, nv))
Patrick Devine's avatar
Patrick Devine committed
698
			}
Patrick Devine's avatar
Patrick Devine committed
699
700
		default:
			params = append(params, fmt.Sprintf("%-*s %#v", cs, k, v))
Patrick Devine's avatar
Patrick Devine committed
701
702
703
704
		}
	}
	resp.Parameters = strings.Join(params, "\n")

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

711
712
713
714
	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
715
	fmt.Fprint(&sb, model.String())
716
	resp.Modelfile = sb.String()
717

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

Daniel Hiltgen's avatar
Daniel Hiltgen committed
721
func (s *Server) ListModelsHandler(c *gin.Context) {
722
	manifests, 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

728
729
	var models []api.ModelResponse
	if err := filepath.Walk(manifests, func(path string, info os.FileInfo, _ error) error {
Patrick Devine's avatar
Patrick Devine committed
730
		if !info.IsDir() {
731
732
733
734
			rel, err := filepath.Rel(manifests, path)
			if err != nil {
				return err
			}
735

736
737
738
739
740
741
			if hidden, err := filepath.Match(".*", filepath.Base(rel)); err != nil {
				return err
			} else if hidden {
				return nil
			}

742
			n := model.ParseNameFromFilepath(rel)
Michael Yang's avatar
Michael Yang committed
743
			if !n.IsValid() {
Michael Yang's avatar
Michael Yang committed
744
				slog.Warn("bad manifest filepath", "path", rel)
Michael Yang's avatar
Michael Yang committed
745
746
747
				return nil
			}

748
			m, err := ParseNamedManifest(n)
Patrick Devine's avatar
Patrick Devine committed
749
			if err != nil {
Michael Yang's avatar
Michael Yang committed
750
751
				slog.Warn("bad manifest", "name", n, "error", err)
				return nil
Patrick Devine's avatar
Patrick Devine committed
752
			}
Michael Yang's avatar
Michael Yang committed
753

754
755
			f, err := m.Config.Open()
			if err != nil {
Michael Yang's avatar
Michael Yang committed
756
757
				slog.Warn("bad manifest config filepath", "name", n, "error", err)
				return nil
758
759
760
761
762
			}
			defer f.Close()

			var c ConfigV2
			if err := json.NewDecoder(f).Decode(&c); err != nil {
Michael Yang's avatar
Michael Yang committed
763
764
				slog.Warn("bad manifest config", "name", n, "error", err)
				return nil
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
			}

			// tag should never be masked
			models = append(models, api.ModelResponse{
				Model:      n.DisplayShortest(),
				Name:       n.DisplayShortest(),
				Size:       m.Size(),
				Digest:     m.Digest,
				ModifiedAt: info.ModTime(),
				Details: api.ModelDetails{
					Format:            c.ModelFormat,
					Family:            c.ModelFamily,
					Families:          c.ModelFamilies,
					ParameterSize:     c.ModelType,
					QuantizationLevel: c.FileType,
				},
			})
Patrick Devine's avatar
Patrick Devine committed
782
		}
Michael Yang's avatar
Michael Yang committed
783

Patrick Devine's avatar
Patrick Devine committed
784
		return nil
785
	}); err != nil {
Patrick Devine's avatar
Patrick Devine committed
786
787
788
789
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

790
791
792
793
794
	slices.SortStableFunc(models, func(i, j api.ModelResponse) int {
		// most recently modified first
		return cmp.Compare(j.ModifiedAt.Unix(), i.ModifiedAt.Unix())
	})

Michael Yang's avatar
Michael Yang committed
795
	c.JSON(http.StatusOK, api.ListResponse{Models: models})
Patrick Devine's avatar
Patrick Devine committed
796
797
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
798
func (s *Server) CopyModelHandler(c *gin.Context) {
Michael Yang's avatar
Michael Yang committed
799
800
	var r api.CopyRequest
	if err := c.ShouldBindJSON(&r); errors.Is(err, io.EOF) {
Michael Yang's avatar
Michael Yang committed
801
802
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "missing request body"})
		return
Michael Yang's avatar
Michael Yang committed
803
	} else if err != nil {
Michael Yang's avatar
Michael Yang committed
804
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": err.Error()})
Patrick Devine's avatar
Patrick Devine committed
805
806
807
		return
	}

Michael Yang's avatar
Michael Yang committed
808
809
	src := model.ParseName(r.Source)
	if !src.IsValid() {
Michael Yang's avatar
Michael Yang committed
810
811
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": fmt.Sprintf("source %q is invalid", r.Source)})
		return
812
813
	}

Michael Yang's avatar
Michael Yang committed
814
815
	dst := model.ParseName(r.Destination)
	if !dst.IsValid() {
816
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": fmt.Sprintf("destination %q is invalid", r.Destination)})
Patrick Devine's avatar
Patrick Devine committed
817
818
		return
	}
Michael Yang's avatar
Michael Yang committed
819
820
821
822
823
824

	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
825
826
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
827
func (s *Server) HeadBlobHandler(c *gin.Context) {
Michael Yang's avatar
Michael Yang committed
828
829
830
831
832
833
834
835
836
837
838
	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
839
	c.Status(http.StatusOK)
Michael Yang's avatar
Michael Yang committed
840
841
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
842
func (s *Server) CreateBlobHandler(c *gin.Context) {
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
	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
	}

861
	layer, err := NewLayer(c.Request.Body, "")
Michael Yang's avatar
Michael Yang committed
862
863
864
865
866
	if err != nil {
		c.AbortWithStatusJSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

867
868
	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
869
870
871
		return
	}

Michael Yang's avatar
Michael Yang committed
872
	c.Status(http.StatusCreated)
Michael Yang's avatar
Michael Yang committed
873
874
}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
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
}

896
func allowedHost(host string) bool {
Jeffrey Morgan's avatar
Jeffrey Morgan committed
897
	if host == "" || host == "localhost" {
898
899
900
901
902
903
904
905
		return true
	}

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

	var tlds = []string{
Jeffrey Morgan's avatar
Jeffrey Morgan committed
906
907
908
		"localhost",
		"local",
		"internal",
909
	}
910

Jeffrey Morgan's avatar
Jeffrey Morgan committed
911
	// check if the host is a local TLD
912
913
914
915
916
917
	for _, tld := range tlds {
		if strings.HasSuffix(host, "."+tld) {
			return true
		}
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
918
	return false
Jeffrey Morgan's avatar
Jeffrey Morgan committed
919
}
920

Jeffrey Morgan's avatar
Jeffrey Morgan committed
921
922
923
func allowedHostsMiddleware(addr net.Addr) gin.HandlerFunc {
	return func(c *gin.Context) {
		if addr == nil {
924
925
926
927
			c.Next()
			return
		}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
928
		if addr, err := netip.ParseAddrPort(addr.String()); err == nil && !addr.Addr().IsLoopback() {
929
930
931
932
933
934
935
936
937
			c.Next()
			return
		}

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

Jeffrey Morgan's avatar
Jeffrey Morgan committed
938
		if addr, err := netip.ParseAddr(host); err == nil {
Jeffrey Morgan's avatar
Jeffrey Morgan committed
939
			if addr.IsLoopback() || addr.IsPrivate() || addr.IsUnspecified() || isLocalIP(addr) {
Jeffrey Morgan's avatar
Jeffrey Morgan committed
940
941
942
943
944
				c.Next()
				return
			}
		}

945
		if allowedHost(host) {
946
947
948
949
950
			if c.Request.Method == "OPTIONS" {
				c.AbortWithStatus(http.StatusNoContent)
				return
			}

951
952
953
954
955
956
			c.Next()
			return
		}

		c.AbortWithStatus(http.StatusForbidden)
	}
957
}
958

959
func (s *Server) GenerateRoutes() http.Handler {
Michael Yang's avatar
Michael Yang committed
960
961
	config := cors.DefaultConfig()
	config.AllowWildcard = true
962
	config.AllowBrowserExtensions = true
963
	config.AllowHeaders = []string{"Authorization", "Content-Type", "User-Agent", "Accept", "X-Requested-With"}
964
	config.AllowOrigins = envconfig.AllowOrigins
Michael Yang's avatar
Michael Yang committed
965

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

Daniel Hiltgen's avatar
Daniel Hiltgen committed
972
973
974
975
976
977
978
979
980
981
982
	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)
983
	r.GET("/api/ps", s.ProcessHandler)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
984

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

Michael Yang's avatar
Michael Yang committed
988
989
990
991
992
	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
993
		r.Handle(method, "/api/tags", s.ListModelsHandler)
Michael Yang's avatar
Michael Yang committed
994
995
996
		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
997
998
	}

999
1000
1001
1002
	return r
}

func Serve(ln net.Listener) error {
Michael Yang's avatar
Michael Yang committed
1003
	level := slog.LevelInfo
1004
	if envconfig.Debug {
Michael Yang's avatar
Michael Yang committed
1005
		level = slog.LevelDebug
1006
	}
Michael Yang's avatar
Michael Yang committed
1007

1008
	slog.Info("server config", "env", envconfig.AsMap())
Michael Yang's avatar
Michael Yang committed
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
	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))

1024
1025
1026
1027
1028
1029
1030
1031
	blobsDir, err := GetBlobsPath("")
	if err != nil {
		return err
	}
	if err := fixBlobs(blobsDir); err != nil {
		return err
	}

1032
	if !envconfig.NoPrune {
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
		// 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
1048
	ctx, done := context.WithCancel(context.Background())
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1049
1050
	schedCtx, schedDone := context.WithCancel(ctx)
	sched := InitScheduler(schedCtx)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1051
	s := &Server{addr: ln.Addr(), sched: sched}
1052
1053
	r := s.GenerateRoutes()

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

1059
1060
	// listen for a ctrl+c and stop any loaded llm
	signals := make(chan os.Signal, 1)
1061
	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
1062
1063
	go func() {
		<-signals
1064
		srvr.Close()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1065
		schedDone()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1066
		sched.unloadAllRunners()
1067
		gpu.Cleanup()
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1068
		done()
1069
1070
	}()

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

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1075
	s.sched.Run(schedCtx)
Daniel Hiltgen's avatar
Daniel Hiltgen committed
1076
1077
1078

	// 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
1079
1080
	gpus := gpu.GetGPUInfo()
	gpus.LogDetails()
1081

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1082
1083
1084
1085
1086
1087
1088
1089
	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()
	return err
Jeffrey Morgan's avatar
Jeffrey Morgan committed
1090
}
Michael Yang's avatar
Michael Yang committed
1091

1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
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
1117
func streamResponse(c *gin.Context, ch chan any) {
1118
	c.Header("Content-Type", "application/x-ndjson")
Michael Yang's avatar
Michael Yang committed
1119
1120
1121
1122
1123
1124
1125
1126
	c.Stream(func(w io.Writer) bool {
		val, ok := <-ch
		if !ok {
			return false
		}

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

1131
		// Delineate chunks with new-line delimiter
Michael Yang's avatar
Michael Yang committed
1132
1133
		bts = append(bts, '\n')
		if _, err := w.Write(bts); err != nil {
1134
			slog.Info(fmt.Sprintf("streamResponse: w.Write failed with %s", err))
Michael Yang's avatar
Michael Yang committed
1135
1136
1137
1138
1139
1140
			return false
		}

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

1142
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
1169
func (s *Server) ProcessHandler(c *gin.Context) {
	models := []api.ModelResponse{}

	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,
		}

		mr := api.ModelResponse{
			Model:     model.ShortName,
			Name:      model.ShortName,
			Size:      int64(v.estimatedTotal),
			SizeVRAM:  int64(v.estimatedVRAM),
			Digest:    model.Digest,
			Details:   modelDetails,
			ExpiresAt: v.expiresAt,
		}
		models = append(models, mr)
	}

	c.JSON(http.StatusOK, api.ListResponse{Models: models})
}

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

1176
	prompt, err := ChatPrompt(template, messages, numCtx, encode)
1177
1178
1179
1180
1181
1182
1183
	if err != nil {
		return "", err
	}

	return prompt, nil
}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1184
func (s *Server) ChatHandler(c *gin.Context) {
Bruce MacDonald's avatar
Bruce MacDonald committed
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
	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
	}

1208
	model, err := GetModel(req.Model)
Bruce MacDonald's avatar
Bruce MacDonald committed
1209
1210
	if err != nil {
		var pErr *fs.PathError
1211
		if errors.As(err, &pErr) {
Bruce MacDonald's avatar
Bruce MacDonald committed
1212
			c.JSON(http.StatusNotFound, gin.H{"error": fmt.Sprintf("model '%s' not found, try pulling it first", req.Model)})
1213
1214
1215
1216
1217
1218
			return
		}
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

1219
	if model.IsEmbedding() {
1220
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "embedding models do not support chat"})
1221
1222
1223
		return
	}

1224
1225
1226
1227
1228
	opts, err := modelOptions(model, req.Options)
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
1229
1230
1231

	var sessionDuration time.Duration
	if req.KeepAlive == nil {
1232
		sessionDuration = getDefaultSessionDuration()
1233
1234
1235
1236
	} else {
		sessionDuration = req.KeepAlive.Duration
	}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1237
1238
1239
1240
1241
	rCh, eCh := s.sched.GetRunner(c.Request.Context(), model, opts, sessionDuration)
	var runner *runnerRef
	select {
	case runner = <-rCh:
	case err = <-eCh:
1242
		handleErrorResponse(c, err)
Bruce MacDonald's avatar
Bruce MacDonald committed
1243
1244
1245
1246
1247
		return
	}

	checkpointLoaded := time.Now()

1248
1249
1250
1251
1252
1253
1254
1255
1256
1257
	// 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
1258
	prompt, err := chatPrompt(c.Request.Context(), runner, model.Template, req.Messages, opts.NumCtx)
Bruce MacDonald's avatar
Bruce MacDonald committed
1259
1260
1261
1262
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}
Michael Yang's avatar
Michael Yang committed
1263

1264
	// an empty request loads the model
1265
	if len(req.Messages) == 0 || prompt == "" {
1266
		resp := api.ChatResponse{
1267
1268
1269
1270
1271
			CreatedAt:  time.Now().UTC(),
			Model:      req.Model,
			Done:       true,
			DoneReason: "load",
			Message:    api.Message{Role: "assistant"},
1272
1273
1274
1275
1276
		}
		c.JSON(http.StatusOK, resp)
		return
	}

1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
	// 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))
1295

Bruce MacDonald's avatar
Bruce MacDonald committed
1296
1297
1298
1299
1300
	ch := make(chan any)

	go func() {
		defer close(ch)

1301
		fn := func(r llm.CompletionResponse) {
Bruce MacDonald's avatar
Bruce MacDonald committed
1302
1303

			resp := api.ChatResponse{
1304
1305
1306
1307
1308
				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
1309
1310
1311
1312
1313
1314
1315
1316
				Metrics: api.Metrics{
					PromptEvalCount:    r.PromptEvalCount,
					PromptEvalDuration: r.PromptEvalDuration,
					EvalCount:          r.EvalCount,
					EvalDuration:       r.EvalDuration,
				},
			}

1317
1318
1319
			if r.Done {
				resp.TotalDuration = time.Since(checkpointStart)
				resp.LoadDuration = checkpointLoaded.Sub(checkpointStart)
Bruce MacDonald's avatar
Bruce MacDonald committed
1320
1321
1322
1323
1324
			}

			ch <- resp
		}

Daniel Hiltgen's avatar
Daniel Hiltgen committed
1325
		if err := runner.llama.Completion(c.Request.Context(), llm.CompletionRequest{
1326
1327
			Prompt:  prompt,
			Format:  req.Format,
Michael Yang's avatar
Michael Yang committed
1328
			Images:  images,
1329
			Options: opts,
1330
		}, fn); err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
1331
1332
1333
1334
1335
			ch <- gin.H{"error": err.Error()}
		}
	}()

	if req.Stream != nil && !*req.Stream {
1336
1337
		// Accumulate responses into the final response
		var final api.ChatResponse
Bruce MacDonald's avatar
Bruce MacDonald committed
1338
1339
		var sb strings.Builder
		for resp := range ch {
1340
1341
			switch r := resp.(type) {
			case api.ChatResponse:
1342
				sb.WriteString(r.Message.Content)
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
				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
1355
1356
			}
		}
1357

1358
		final.Message = api.Message{Role: "assistant", Content: sb.String()}
1359
		c.JSON(http.StatusOK, final)
Bruce MacDonald's avatar
Bruce MacDonald committed
1360
1361
1362
1363
1364
		return
	}

	streamResponse(c, ch)
}
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376

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