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
26
	"github.com/gin-gonic/gin"

Jeffrey Morgan's avatar
Jeffrey Morgan committed
27
	"github.com/jmorganca/ollama/api"
28
	"github.com/jmorganca/ollama/llm"
Michael Yang's avatar
Michael Yang committed
29
	"github.com/jmorganca/ollama/version"
Jeffrey Morgan's avatar
Jeffrey Morgan committed
30
31
)

Michael Yang's avatar
Michael Yang committed
32
33
34
35
36
37
38
39
40
41
42
43
44
45
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
46
var loaded struct {
Michael Yang's avatar
Michael Yang committed
47
48
	mu sync.Mutex

49
	runner llm.LLM
Michael Yang's avatar
Michael Yang committed
50
51
52

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

54
55
	*Model
	*api.Options
Michael Yang's avatar
Michael Yang committed
56
57
}

58
59
var defaultSessionDuration = 5 * time.Minute

Bruce MacDonald's avatar
Bruce MacDonald committed
60
// 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
61
func load(ctx context.Context, workDir string, model *Model, reqOpts map[string]interface{}, sessionDuration time.Duration) error {
62
63
64
	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
65
		return err
66
67
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
68
69
	if err := opts.FromMap(reqOpts); err != nil {
		return err
70
71
	}

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

84
85
86
87
88
89
90
	needLoad := loaded.runner == nil || // is there a model loaded?
		loaded.ModelPath != model.ModelPath || // has the base model changed?
		!reflect.DeepEqual(loaded.AdapterPaths, model.AdapterPaths) || // have the adapters changed?
		!reflect.DeepEqual(loaded.Options.Runner, opts.Runner) // have the runner options changed?

	if needLoad {
		if loaded.runner != nil {
91
			log.Println("changing loaded model")
92
93
94
95
			loaded.runner.Close()
			loaded.runner = nil
			loaded.Model = nil
			loaded.Options = nil
Michael Yang's avatar
Michael Yang committed
96
		}
Michael Yang's avatar
Michael Yang committed
97

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

103
104
105
		loaded.Model = model
		loaded.runner = llmRunner
		loaded.Options = &opts
Michael Yang's avatar
Michael Yang committed
106
	}
107

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

Jeffrey Morgan's avatar
Jeffrey Morgan committed
110
111
112
113
	if loaded.expireTimer == nil {
		loaded.expireTimer = time.AfterFunc(sessionDuration, func() {
			loaded.mu.Lock()
			defer loaded.mu.Unlock()
Michael Yang's avatar
Michael Yang committed
114

Jeffrey Morgan's avatar
Jeffrey Morgan committed
115
			if time.Now().Before(loaded.expireAt) {
Michael Yang's avatar
Michael Yang committed
116
117
118
				return
			}

119
120
			if loaded.runner != nil {
				loaded.runner.Close()
Michael Yang's avatar
Michael Yang committed
121
122
			}

123
124
125
			loaded.runner = nil
			loaded.Model = nil
			loaded.Options = nil
Michael Yang's avatar
Michael Yang committed
126
		})
Michael Yang's avatar
Michael Yang committed
127
	}
128

Jeffrey Morgan's avatar
Jeffrey Morgan committed
129
	loaded.expireTimer.Reset(sessionDuration)
Bruce MacDonald's avatar
Bruce MacDonald committed
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
	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
	}

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

Bruce MacDonald's avatar
Bruce MacDonald committed
150
151
	model, err := GetModel(req.Model)
	if err != nil {
152
153
154
155
156
		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
157
158
159
160
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

161
162
163
164
165
	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 {
166
167
168
169
		if errors.Is(err, api.ErrInvalidOpts) {
			c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
			return
		}
Bruce MacDonald's avatar
Bruce MacDonald committed
170
171
172
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
Michael Yang's avatar
Michael Yang committed
173

Michael Yang's avatar
Michael Yang committed
174
175
	checkpointLoaded := time.Now()

176
	prompt, err := model.Prompt(req)
Michael Yang's avatar
Michael Yang committed
177
178
179
180
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
181

Michael Yang's avatar
Michael Yang committed
182
183
184
	ch := make(chan any)
	go func() {
		defer close(ch)
Michael Yang's avatar
Michael Yang committed
185
186
187
188
189
190
		// an empty request loads the model
		if req.Prompt == "" && req.Template == "" && req.System == "" {
			ch <- api.GenerateResponse{CreatedAt: time.Now().UTC(), Model: req.Model, Done: true}
			return
		}

Michael Yang's avatar
Michael Yang committed
191
		fn := func(r api.GenerateResponse) {
Jeffrey Morgan's avatar
Jeffrey Morgan committed
192
193
			loaded.expireAt = time.Now().Add(sessionDuration)
			loaded.expireTimer.Reset(sessionDuration)
Michael Yang's avatar
Michael Yang committed
194

Michael Yang's avatar
Michael Yang committed
195
196
197
			r.Model = req.Model
			r.CreatedAt = time.Now().UTC()
			if r.Done {
Michael Yang's avatar
Michael Yang committed
198
199
				r.TotalDuration = time.Since(checkpointStart)
				r.LoadDuration = checkpointLoaded.Sub(checkpointStart)
Michael Yang's avatar
Michael Yang committed
200
201
202
			}

			ch <- r
Michael Yang's avatar
Michael Yang committed
203
204
		}

Michael Yang's avatar
Michael Yang committed
205
206
		if err := loaded.runner.Predict(c.Request.Context(), req.Context, prompt, fn); err != nil {
			ch <- gin.H{"error": err.Error()}
Michael Yang's avatar
Michael Yang committed
207
		}
Michael Yang's avatar
Michael Yang committed
208
	}()
Michael Yang's avatar
Michael Yang committed
209

210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
	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
227
	streamResponse(c, ch)
Michael Yang's avatar
Michael Yang committed
228
}
Michael Yang's avatar
Michael Yang committed
229

Bruce MacDonald's avatar
Bruce MacDonald committed
230
231
232
233
234
235
236
237
238
239
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
	}

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

Bruce MacDonald's avatar
Bruce MacDonald committed
245
246
247
248
249
	model, err := GetModel(req.Model)
	if err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}
250
251
252

	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
253
254
255
256
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

257
	if !loaded.Options.EmbeddingOnly {
Bruce MacDonald's avatar
Bruce MacDonald committed
258
259
260
261
		c.JSON(http.StatusBadRequest, gin.H{"error": "embedding option must be set to true"})
		return
	}

262
	embedding, err := loaded.runner.Embedding(c.Request.Context(), req.Prompt)
Bruce MacDonald's avatar
Bruce MacDonald committed
263
264
265
266
267
268
269
270
271
272
273
274
	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)
}

275
func PullModelHandler(c *gin.Context) {
Michael Yang's avatar
Michael Yang committed
276
277
278
279
280
281
	var req api.PullRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

282
283
284
285
286
	if req.Name == "" {
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "name is required"})
		return
	}

287
288
289
	ch := make(chan any)
	go func() {
		defer close(ch)
290
291
		fn := func(r api.ProgressResponse) {
			ch <- r
292
		}
293

294
295
296
297
		regOpts := &RegistryOptions{
			Insecure: req.Insecure,
		}

298
299
300
301
		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
302
			ch <- gin.H{"error": err.Error()}
303
304
305
		}
	}()

306
307
308
309
310
	if req.Stream != nil && !*req.Stream {
		waitForStream(c, ch)
		return
	}

311
312
313
	streamResponse(c, ch)
}

314
func PushModelHandler(c *gin.Context) {
315
316
317
	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
318
319
		return
	}
Michael Yang's avatar
Michael Yang committed
320

321
322
323
324
325
	if req.Name == "" {
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "name is required"})
		return
	}

326
327
328
	ch := make(chan any)
	go func() {
		defer close(ch)
329
330
		fn := func(r api.ProgressResponse) {
			ch <- r
331
		}
332

333
334
335
336
		regOpts := &RegistryOptions{
			Insecure: req.Insecure,
		}

337
338
		ctx := context.Background()
		if err := PushModel(ctx, req.Name, regOpts, fn); err != nil {
Michael Yang's avatar
Michael Yang committed
339
			ch <- gin.H{"error": err.Error()}
340
341
342
		}
	}()

343
344
345
346
347
	if req.Stream != nil && !*req.Stream {
		waitForStream(c, ch)
		return
	}

348
349
350
	streamResponse(c, ch)
}

351
func CreateModelHandler(c *gin.Context) {
352
353
	var req api.CreateRequest
	if err := c.ShouldBindJSON(&req); err != nil {
Michael Yang's avatar
Michael Yang committed
354
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
Michael Yang's avatar
Michael Yang committed
355
		return
356
357
	}

358
359
360
361
362
	if req.Name == "" || req.Path == "" {
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "name and path are required"})
		return
	}

Michael Yang's avatar
Michael Yang committed
363
	ch := make(chan any)
Michael Yang's avatar
Michael Yang committed
364
365
	go func() {
		defer close(ch)
366
367
		fn := func(resp api.ProgressResponse) {
			ch <- resp
368
369
		}

370
371
372
		ctx, cancel := context.WithCancel(c.Request.Context())
		defer cancel()

373
		if err := CreateModel(ctx, req.Name, req.Path, fn); err != nil {
Michael Yang's avatar
Michael Yang committed
374
			ch <- gin.H{"error": err.Error()}
375
		}
Michael Yang's avatar
Michael Yang committed
376
	}()
Michael Yang's avatar
Michael Yang committed
377

378
379
380
381
382
	if req.Stream != nil && !*req.Stream {
		waitForStream(c, ch)
		return
	}

Michael Yang's avatar
Michael Yang committed
383
	streamResponse(c, ch)
Bruce MacDonald's avatar
Bruce MacDonald committed
384
385
}

386
387
388
389
390
391
392
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
	}

393
394
395
396
397
	if req.Name == "" {
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "name is required"})
		return
	}

398
399
400
401
	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 {
402
403
			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		}
404
405
		return
	}
Michael Yang's avatar
Michael Yang committed
406
407
408
409
410
411
412
413
414
415
416
417

	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
	}

418
	c.JSON(http.StatusOK, nil)
419
420
}

Patrick Devine's avatar
Patrick Devine committed
421
422
423
424
425
426
427
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
	}

428
429
430
431
432
	if req.Name == "" {
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "name is required"})
		return
	}

Patrick Devine's avatar
Patrick Devine committed
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
	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
}

497
func ListModelsHandler(c *gin.Context) {
498
	models := make([]api.ModelResponse, 0)
Patrick Devine's avatar
Patrick Devine committed
499
500
501
502
503
	fp, err := GetManifestPath()
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
Michael Yang's avatar
Michael Yang committed
504
505

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

511
			mp := ParseModelPath(tag)
Patrick Devine's avatar
Patrick Devine committed
512
			manifest, digest, err := GetManifest(mp)
Patrick Devine's avatar
Patrick Devine committed
513
			if err != nil {
514
515
				log.Printf("skipping file: %s", fp)
				return nil
Patrick Devine's avatar
Patrick Devine committed
516
			}
Michael Yang's avatar
Michael Yang committed
517
518

			models = append(models, api.ModelResponse{
Patrick Devine's avatar
Patrick Devine committed
519
520
				Name:       mp.GetShortTagname(),
				Size:       manifest.GetTotalSize(),
Patrick Devine's avatar
Patrick Devine committed
521
				Digest:     digest,
Michael Yang's avatar
Michael Yang committed
522
523
				ModifiedAt: info.ModTime(),
			})
Patrick Devine's avatar
Patrick Devine committed
524
		}
Michael Yang's avatar
Michael Yang committed
525

Patrick Devine's avatar
Patrick Devine committed
526
		return nil
Michael Yang's avatar
Michael Yang committed
527
528
529
	}

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

Michael Yang's avatar
Michael Yang committed
534
	c.JSON(http.StatusOK, api.ListResponse{Models: models})
Patrick Devine's avatar
Patrick Devine committed
535
536
}

Patrick Devine's avatar
Patrick Devine committed
537
538
539
540
541
542
543
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
	}

544
545
546
547
548
	if req.Source == "" || req.Destination == "" {
		c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "source add destination are required"})
		return
	}

Patrick Devine's avatar
Patrick Devine committed
549
550
551
552
553
554
555
556
557
558
	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)
	}

Michael Yang's avatar
Michael Yang committed
611
	log.Printf("Listening on %s (version %s)", ln.Addr(), version.Version)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
612
613
614
615
	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
		if loaded.runner != nil {
			loaded.runner.Close()
623
		}
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
	})
}