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

import (
4
	"context"
Michael Yang's avatar
Michael Yang committed
5
	"encoding/json"
6
	"errors"
7
	"fmt"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
8
	"io"
9
	"io/fs"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
10
11
12
	"log"
	"net"
	"net/http"
13
	"os"
14
	"os/signal"
Michael Yang's avatar
Michael Yang committed
15
	"path/filepath"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
16
	"reflect"
17
	"runtime"
Patrick Devine's avatar
Patrick Devine committed
18
	"strconv"
Michael Yang's avatar
Michael Yang committed
19
	"strings"
Michael Yang's avatar
Michael Yang committed
20
	"sync"
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"
Bruce MacDonald's avatar
Bruce MacDonald committed
26
	"gonum.org/v1/gonum/mat"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
27

Jeffrey Morgan's avatar
Jeffrey Morgan committed
28
	"github.com/jmorganca/ollama/api"
29
	"github.com/jmorganca/ollama/llm"
30
	"github.com/jmorganca/ollama/vector"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
31
32
)

Michael Yang's avatar
Michael Yang committed
33
34
35
36
37
38
39
40
41
42
43
44
45
46
var mode string = gin.DebugMode

func init() {
	switch mode {
	case gin.DebugMode:
	case gin.ReleaseMode:
	case gin.TestMode:
	default:
		mode = gin.DebugMode
	}

	gin.SetMode(mode)
}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
47
var loaded struct {
Michael Yang's avatar
Michael Yang committed
48
49
	mu sync.Mutex

50
	llm        llm.LLM
51
	Embeddings []vector.Embedding
Michael Yang's avatar
Michael Yang committed
52
53
54

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

56
57
	digest  string
	options api.Options
Michael Yang's avatar
Michael Yang committed
58
59
}

60
61
var defaultSessionDuration = 5 * time.Minute

Bruce MacDonald's avatar
Bruce MacDonald committed
62
// load a model into memory if it is not already loaded, it is up to the caller to lock loaded.mu before calling this function
63
func load(ctx context.Context, workDir string, model *Model, reqOpts map[string]interface{}, sessionDuration time.Duration) error {
64
65
66
	opts := api.DefaultOptions()
	if err := opts.FromMap(model.Options); err != nil {
		log.Printf("could not load model options: %v", err)
Bruce MacDonald's avatar
Bruce MacDonald committed
67
		return err
68
69
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
70
	if err := opts.FromMap(reqOpts); err != nil {
71
		log.Printf("could not merge model options: %v", err)
Bruce MacDonald's avatar
Bruce MacDonald committed
72
		return err
73
74
	}

75
76
77
78
79
80
81
82
83
84
85
	// check if the loaded model is still running in a subprocess, in case something unexpected happened
	if loaded.llm != nil {
		if err := loaded.llm.Ping(ctx); err != nil {
			log.Print("loaded llm process not responding, closing now")
			// the subprocess is no longer running, so close it
			loaded.llm.Close()
			loaded.llm = nil
			loaded.digest = ""
		}
	}

86
	if model.Digest != loaded.digest || !reflect.DeepEqual(loaded.options, opts) {
Jeffrey Morgan's avatar
Jeffrey Morgan committed
87
		if loaded.llm != nil {
88
			log.Println("changing loaded model")
Jeffrey Morgan's avatar
Jeffrey Morgan committed
89
90
			loaded.llm.Close()
			loaded.llm = nil
91
			loaded.digest = ""
Michael Yang's avatar
Michael Yang committed
92
		}
Michael Yang's avatar
Michael Yang committed
93

94
95
96
97
98
		if model.Embeddings != nil && len(model.Embeddings) > 0 {
			opts.EmbeddingOnly = true // this is requried to generate embeddings, completions will still work
			loaded.Embeddings = model.Embeddings
		}

99
		llmModel, err := llm.New(workDir, model.ModelPath, model.AdapterPaths, opts)
Michael Yang's avatar
Michael Yang committed
100
		if err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
101
			return err
Michael Yang's avatar
Michael Yang committed
102
103
		}

104
105
106
107
108
		// set cache values before modifying opts
		loaded.llm = llmModel
		loaded.digest = model.Digest
		loaded.options = opts

109
		if opts.NumKeep < 0 {
Bruce MacDonald's avatar
Bruce MacDonald committed
110
			promptWithSystem, err := model.Prompt(api.GenerateRequest{}, "")
111
			if err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
112
				return err
113
114
			}

Bruce MacDonald's avatar
Bruce MacDonald committed
115
			promptNoSystem, err := model.Prompt(api.GenerateRequest{Context: []int{0}}, "")
116
			if err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
117
				return err
118
119
			}

120
121
122
123
			tokensWithSystem, err := llmModel.Encode(ctx, promptWithSystem)
			if err != nil {
				return err
			}
Michael Yang's avatar
Michael Yang committed
124

125
126
127
128
			tokensNoSystem, err := llmModel.Encode(ctx, promptNoSystem)
			if err != nil {
				return err
			}
129

Michael Yang's avatar
Michael Yang committed
130
			opts.NumKeep = len(tokensWithSystem) - len(tokensNoSystem)
131

132
133
			llmModel.SetOptions(opts)
		}
Michael Yang's avatar
Michael Yang committed
134
	}
135

Jeffrey Morgan's avatar
Jeffrey Morgan committed
136
	loaded.expireAt = time.Now().Add(sessionDuration)
Bruce MacDonald's avatar
Bruce MacDonald committed
137

Jeffrey Morgan's avatar
Jeffrey Morgan committed
138
139
140
141
	if loaded.expireTimer == nil {
		loaded.expireTimer = time.AfterFunc(sessionDuration, func() {
			loaded.mu.Lock()
			defer loaded.mu.Unlock()
Michael Yang's avatar
Michael Yang committed
142

Jeffrey Morgan's avatar
Jeffrey Morgan committed
143
			if time.Now().Before(loaded.expireAt) {
Michael Yang's avatar
Michael Yang committed
144
145
146
				return
			}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
147
			if loaded.llm == nil {
Michael Yang's avatar
Michael Yang committed
148
149
150
				return
			}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
151
152
			loaded.llm.Close()
			loaded.llm = nil
153
			loaded.digest = ""
Michael Yang's avatar
Michael Yang committed
154
		})
Michael Yang's avatar
Michael Yang committed
155
	}
156

Jeffrey Morgan's avatar
Jeffrey Morgan committed
157
	loaded.expireTimer.Reset(sessionDuration)
Bruce MacDonald's avatar
Bruce MacDonald committed
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
	return nil
}

func GenerateHandler(c *gin.Context) {
	loaded.mu.Lock()
	defer loaded.mu.Unlock()

	checkpointStart := time.Now()

	var req api.GenerateRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	model, err := GetModel(req.Model)
	if err != nil {
175
176
177
178
179
		var pErr *fs.PathError
		if errors.As(err, &pErr) {
			c.JSON(http.StatusNotFound, gin.H{"error": fmt.Sprintf("model '%s' not found, try pulling it first", req.Model)})
			return
		}
Bruce MacDonald's avatar
Bruce MacDonald committed
180
181
182
183
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

184
185
186
187
188
	workDir := c.GetString("workDir")

	// TODO: set this duration from the request if specified
	sessionDuration := defaultSessionDuration
	if err := load(c.Request.Context(), workDir, model, req.Options, sessionDuration); err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
189
190
191
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
Michael Yang's avatar
Michael Yang committed
192

Michael Yang's avatar
Michael Yang committed
193
194
	checkpointLoaded := time.Now()

Bruce MacDonald's avatar
Bruce MacDonald committed
195
196
	embedding := ""
	if model.Embeddings != nil && len(model.Embeddings) > 0 {
197
		promptEmbed, err := loaded.llm.Embedding(c.Request.Context(), req.Prompt)
Bruce MacDonald's avatar
Bruce MacDonald committed
198
199
200
201
202
203
204
205
206
207
208
209
210
		if err != nil {
			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
			return
		}
		// TODO: set embed_top from specified parameters in modelfile
		embed_top := 3
		topK := vector.TopK(embed_top, mat.NewVecDense(len(promptEmbed), promptEmbed), loaded.Embeddings)
		for _, e := range topK {
			embedding = fmt.Sprintf("%s %s", embedding, e.Embedding.Data)
		}
	}

	prompt, err := model.Prompt(req, embedding)
Michael Yang's avatar
Michael Yang committed
211
212
213
214
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
215

Michael Yang's avatar
Michael Yang committed
216
217
218
	ch := make(chan any)
	go func() {
		defer close(ch)
Michael Yang's avatar
Michael Yang committed
219
		fn := func(r api.GenerateResponse) {
Jeffrey Morgan's avatar
Jeffrey Morgan committed
220
221
			loaded.expireAt = time.Now().Add(sessionDuration)
			loaded.expireTimer.Reset(sessionDuration)
Michael Yang's avatar
Michael Yang committed
222

Michael Yang's avatar
Michael Yang committed
223
224
225
			r.Model = req.Model
			r.CreatedAt = time.Now().UTC()
			if r.Done {
Michael Yang's avatar
Michael Yang committed
226
227
				r.TotalDuration = time.Since(checkpointStart)
				r.LoadDuration = checkpointLoaded.Sub(checkpointStart)
Michael Yang's avatar
Michael Yang committed
228
229
230
			}

			ch <- r
Michael Yang's avatar
Michael Yang committed
231
232
		}

233
234
		// an empty request loads the model
		if req.Prompt == "" && req.Template == "" && req.System == "" {
Patrick Devine's avatar
Patrick Devine committed
235
236
237
238
239
			ch <- api.GenerateResponse{Model: req.Model, Done: true}
		} else {
			if err := loaded.llm.Predict(c.Request.Context(), req.Context, prompt, fn); err != nil {
				ch <- gin.H{"error": err.Error()}
			}
Michael Yang's avatar
Michael Yang committed
240
		}
Michael Yang's avatar
Michael Yang committed
241
	}()
Michael Yang's avatar
Michael Yang committed
242

243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
	if req.Stream != nil && !*req.Stream {
		var response api.GenerateResponse
		generated := ""
		for resp := range ch {
			if r, ok := resp.(api.GenerateResponse); ok {
				generated += r.Response
				response = r
			} else {
				c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
				return
			}
		}
		response.Response = generated
		c.JSON(http.StatusOK, response)
		return
	}

Michael Yang's avatar
Michael Yang committed
260
	streamResponse(c, ch)
Michael Yang's avatar
Michael Yang committed
261
}
Michael Yang's avatar
Michael Yang committed
262

Bruce MacDonald's avatar
Bruce MacDonald committed
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
func EmbeddingHandler(c *gin.Context) {
	loaded.mu.Lock()
	defer loaded.mu.Unlock()

	var req api.EmbeddingRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	model, err := GetModel(req.Model)
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}
278
279
280

	workDir := c.GetString("workDir")
	if err := load(c.Request.Context(), workDir, model, req.Options, 5*time.Minute); err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
281
282
283
284
285
286
287
288
289
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

	if !loaded.options.EmbeddingOnly {
		c.JSON(http.StatusBadRequest, gin.H{"error": "embedding option must be set to true"})
		return
	}

290
	embedding, err := loaded.llm.Embedding(c.Request.Context(), req.Prompt)
Bruce MacDonald's avatar
Bruce MacDonald committed
291
292
293
294
295
296
297
298
299
300
301
302
	if err != nil {
		log.Printf("embedding generation failed: %v", err)
		c.JSON(http.StatusInternalServerError, gin.H{"error": "failed to generate embedding"})
		return
	}

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

303
func PullModelHandler(c *gin.Context) {
Michael Yang's avatar
Michael Yang committed
304
305
306
307
308
309
	var req api.PullRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

310
311
312
	ch := make(chan any)
	go func() {
		defer close(ch)
313
314
		fn := func(r api.ProgressResponse) {
			ch <- r
315
		}
316

317
318
319
320
		regOpts := &RegistryOptions{
			Insecure: req.Insecure,
		}

321
322
323
324
		ctx, cancel := context.WithCancel(c.Request.Context())
		defer cancel()

		if err := PullModel(ctx, req.Name, regOpts, fn); err != nil {
Michael Yang's avatar
Michael Yang committed
325
			ch <- gin.H{"error": err.Error()}
326
327
328
		}
	}()

329
330
331
332
333
	if req.Stream != nil && !*req.Stream {
		waitForStream(c, ch)
		return
	}

334
335
336
	streamResponse(c, ch)
}

337
func PushModelHandler(c *gin.Context) {
338
339
340
	var req api.PushRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
Michael Yang's avatar
Michael Yang committed
341
342
		return
	}
Michael Yang's avatar
Michael Yang committed
343

344
345
346
	ch := make(chan any)
	go func() {
		defer close(ch)
347
348
		fn := func(r api.ProgressResponse) {
			ch <- r
349
		}
350

351
352
353
354
		regOpts := &RegistryOptions{
			Insecure: req.Insecure,
		}

355
356
		ctx := context.Background()
		if err := PushModel(ctx, req.Name, regOpts, fn); err != nil {
Michael Yang's avatar
Michael Yang committed
357
			ch <- gin.H{"error": err.Error()}
358
359
360
		}
	}()

361
362
363
364
365
	if req.Stream != nil && !*req.Stream {
		waitForStream(c, ch)
		return
	}

366
367
368
	streamResponse(c, ch)
}

369
func CreateModelHandler(c *gin.Context) {
370
371
372
	var req api.CreateRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"message": err.Error()})
Michael Yang's avatar
Michael Yang committed
373
		return
374
375
	}

376
377
	workDir := c.GetString("workDir")

Michael Yang's avatar
Michael Yang committed
378
	ch := make(chan any)
Michael Yang's avatar
Michael Yang committed
379
380
	go func() {
		defer close(ch)
381
382
		fn := func(resp api.ProgressResponse) {
			ch <- resp
383
384
		}

385
386
387
		ctx, cancel := context.WithCancel(c.Request.Context())
		defer cancel()

388
		if err := CreateModel(ctx, workDir, req.Name, req.Path, fn); err != nil {
Michael Yang's avatar
Michael Yang committed
389
			ch <- gin.H{"error": err.Error()}
390
		}
Michael Yang's avatar
Michael Yang committed
391
	}()
Michael Yang's avatar
Michael Yang committed
392

393
394
395
396
397
	if req.Stream != nil && !*req.Stream {
		waitForStream(c, ch)
		return
	}

Michael Yang's avatar
Michael Yang committed
398
	streamResponse(c, ch)
Bruce MacDonald's avatar
Bruce MacDonald committed
399
400
}

401
402
403
404
405
406
407
func DeleteModelHandler(c *gin.Context) {
	var req api.DeleteRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

408
409
410
411
	if err := DeleteModel(req.Name); err != nil {
		if os.IsNotExist(err) {
			c.JSON(http.StatusNotFound, gin.H{"error": fmt.Sprintf("model '%s' not found", req.Name)})
		} else {
412
413
			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		}
414
415
		return
	}
Michael Yang's avatar
Michael Yang committed
416
417
418
419
420
421
422
423
424
425
426
427

	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
	}

428
	c.JSON(http.StatusOK, nil)
429
430
}

Patrick Devine's avatar
Patrick Devine committed
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
func ShowModelHandler(c *gin.Context) {
	var req api.ShowRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

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

	c.JSON(http.StatusOK, resp)
}

func GetModelInfo(name string) (*api.ShowResponse, error) {
	model, err := GetModel(name)
	if err != nil {
		return nil, err
	}

	resp := &api.ShowResponse{
		License:  strings.Join(model.License, "\n"),
		System:   model.System,
		Template: model.Template,
	}

	mf, err := ShowModelfile(model)
	if err != nil {
		return nil, err
	}

	resp.Modelfile = mf

	var params []string
	cs := 30
	for k, v := range model.Options {
		switch val := v.(type) {
		case string:
			params = append(params, fmt.Sprintf("%-*s %s", cs, k, val))
		case int:
			params = append(params, fmt.Sprintf("%-*s %s", cs, k, strconv.Itoa(val)))
		case float64:
			params = append(params, fmt.Sprintf("%-*s %s", cs, k, strconv.FormatFloat(val, 'f', 0, 64)))
		case bool:
			params = append(params, fmt.Sprintf("%-*s %s", cs, k, strconv.FormatBool(val)))
		case []interface{}:
			for _, nv := range val {
				switch nval := nv.(type) {
				case string:
					params = append(params, fmt.Sprintf("%-*s %s", cs, k, nval))
				case int:
					params = append(params, fmt.Sprintf("%-*s %s", cs, k, strconv.Itoa(nval)))
				case float64:
					params = append(params, fmt.Sprintf("%-*s %s", cs, k, strconv.FormatFloat(nval, 'f', 0, 64)))
				case bool:
					params = append(params, fmt.Sprintf("%-*s %s", cs, k, strconv.FormatBool(nval)))
				}
			}
		}
	}
	resp.Parameters = strings.Join(params, "\n")

	return resp, nil
}

502
func ListModelsHandler(c *gin.Context) {
503
	var models []api.ModelResponse
Patrick Devine's avatar
Patrick Devine committed
504
505
506
507
508
	fp, err := GetManifestPath()
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
Michael Yang's avatar
Michael Yang committed
509
510

	walkFunc := func(path string, info os.FileInfo, _ error) error {
Patrick Devine's avatar
Patrick Devine committed
511
		if !info.IsDir() {
Michael Yang's avatar
Michael Yang committed
512
513
514
			dir, file := filepath.Split(path)
			dir = strings.Trim(strings.TrimPrefix(dir, fp), string(os.PathSeparator))
			tag := strings.Join([]string{dir, file}, ":")
515

516
			mp := ParseModelPath(tag)
Patrick Devine's avatar
Patrick Devine committed
517
			manifest, digest, err := GetManifest(mp)
Patrick Devine's avatar
Patrick Devine committed
518
			if err != nil {
519
520
				log.Printf("skipping file: %s", fp)
				return nil
Patrick Devine's avatar
Patrick Devine committed
521
			}
Michael Yang's avatar
Michael Yang committed
522
523

			models = append(models, api.ModelResponse{
Patrick Devine's avatar
Patrick Devine committed
524
525
				Name:       mp.GetShortTagname(),
				Size:       manifest.GetTotalSize(),
Patrick Devine's avatar
Patrick Devine committed
526
				Digest:     digest,
Michael Yang's avatar
Michael Yang committed
527
528
				ModifiedAt: info.ModTime(),
			})
Patrick Devine's avatar
Patrick Devine committed
529
		}
Michael Yang's avatar
Michael Yang committed
530

Patrick Devine's avatar
Patrick Devine committed
531
		return nil
Michael Yang's avatar
Michael Yang committed
532
533
534
	}

	if err := filepath.Walk(fp, walkFunc); err != nil {
Patrick Devine's avatar
Patrick Devine committed
535
536
537
538
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}

Michael Yang's avatar
Michael Yang committed
539
	c.JSON(http.StatusOK, api.ListResponse{Models: models})
Patrick Devine's avatar
Patrick Devine committed
540
541
}

Patrick Devine's avatar
Patrick Devine committed
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
func CopyModelHandler(c *gin.Context) {
	var req api.CopyRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

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

Michael Yang's avatar
Michael Yang committed
559
560
561
562
563
564
565
var defaultAllowOrigins = []string{
	"localhost",
	"127.0.0.1",
	"0.0.0.0",
}

func Serve(ln net.Listener, allowOrigins []string) error {
Michael Yang's avatar
Michael Yang committed
566
567
	config := cors.DefaultConfig()
	config.AllowWildcard = true
Michael Yang's avatar
Michael Yang committed
568
569
570
571
572
573
574
575
576
577

	config.AllowOrigins = allowOrigins
	for _, allowOrigin := range defaultAllowOrigins {
		config.AllowOrigins = append(config.AllowOrigins,
			fmt.Sprintf("http://%s", allowOrigin),
			fmt.Sprintf("https://%s", allowOrigin),
			fmt.Sprintf("http://%s:*", allowOrigin),
			fmt.Sprintf("https://%s:*", allowOrigin),
		)
	}
Michael Yang's avatar
Michael Yang committed
578

579
580
581
582
583
584
	workDir, err := os.MkdirTemp("", "ollama")
	if err != nil {
		return err
	}
	defer os.RemoveAll(workDir)

Bruce MacDonald's avatar
Bruce MacDonald committed
585
	r := gin.Default()
586
587
588
589
590
591
592
	r.Use(
		cors.New(config),
		func(c *gin.Context) {
			c.Set("workDir", workDir)
			c.Next()
		},
	)
Bruce MacDonald's avatar
Bruce MacDonald committed
593

594
595
	r.POST("/api/pull", PullModelHandler)
	r.POST("/api/generate", GenerateHandler)
Bruce MacDonald's avatar
Bruce MacDonald committed
596
	r.POST("/api/embeddings", EmbeddingHandler)
597
598
	r.POST("/api/create", CreateModelHandler)
	r.POST("/api/push", PushModelHandler)
Patrick Devine's avatar
Patrick Devine committed
599
	r.POST("/api/copy", CopyModelHandler)
600
	r.DELETE("/api/delete", DeleteModelHandler)
Patrick Devine's avatar
Patrick Devine committed
601
	r.POST("/api/show", ShowModelHandler)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
602

Michael Yang's avatar
Michael Yang committed
603
604
605
606
607
608
609
610
	for _, method := range []string{http.MethodGet, http.MethodHead} {
		r.Handle(method, "/", func(c *gin.Context) {
			c.String(http.StatusOK, "Ollama is running")
		})

		r.Handle(method, "/api/tags", ListModelsHandler)
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
611
612
613
614
615
	log.Printf("Listening on %s", ln.Addr())
	s := &http.Server{
		Handler: r,
	}

616
617
	// listen for a ctrl+c and stop any loaded llm
	signals := make(chan os.Signal, 1)
618
	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
619
620
	go func() {
		<-signals
621
622
623
		if loaded.llm != nil {
			loaded.llm.Close()
		}
624
		os.RemoveAll(workDir)
625
626
627
		os.Exit(0)
	}()

628
629
630
	if runtime.GOOS == "linux" {
		// check compatibility to log warnings
		if _, err := llm.CheckVRAM(); err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
631
			log.Printf("Warning: GPU support may not enabled, check you have installed install GPU drivers: %v", err)
632
633
634
		}
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
635
636
	return s.Serve(ln)
}
Michael Yang's avatar
Michael Yang committed
637

638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
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
663
func streamResponse(c *gin.Context, ch chan any) {
664
	c.Header("Content-Type", "application/x-ndjson")
Michael Yang's avatar
Michael Yang committed
665
666
667
668
669
670
671
672
	c.Stream(func(w io.Writer) bool {
		val, ok := <-ch
		if !ok {
			return false
		}

		bts, err := json.Marshal(val)
		if err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
673
			log.Printf("streamResponse: json.Marshal failed with %s", err)
Michael Yang's avatar
Michael Yang committed
674
675
676
			return false
		}

677
		// Delineate chunks with new-line delimiter
Michael Yang's avatar
Michael Yang committed
678
679
		bts = append(bts, '\n')
		if _, err := w.Write(bts); err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
680
			log.Printf("streamResponse: w.Write failed with %s", err)
Michael Yang's avatar
Michael Yang committed
681
682
683
684
685
686
			return false
		}

		return true
	})
}