routes.go 15 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
145
146
	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 {
147
148
149
150
151
		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
152
153
154
155
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

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

Michael Yang's avatar
Michael Yang committed
169
170
	checkpointLoaded := time.Now()

171
	prompt, err := model.Prompt(req)
Michael Yang's avatar
Michael Yang committed
172
173
174
175
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
Jeffrey Morgan's avatar
Jeffrey Morgan committed
176

Michael Yang's avatar
Michael Yang committed
177
178
179
	ch := make(chan any)
	go func() {
		defer close(ch)
Michael Yang's avatar
Michael Yang committed
180
		fn := func(r api.GenerateResponse) {
Jeffrey Morgan's avatar
Jeffrey Morgan committed
181
182
			loaded.expireAt = time.Now().Add(sessionDuration)
			loaded.expireTimer.Reset(sessionDuration)
Michael Yang's avatar
Michael Yang committed
183

Michael Yang's avatar
Michael Yang committed
184
185
186
			r.Model = req.Model
			r.CreatedAt = time.Now().UTC()
			if r.Done {
Michael Yang's avatar
Michael Yang committed
187
188
				r.TotalDuration = time.Since(checkpointStart)
				r.LoadDuration = checkpointLoaded.Sub(checkpointStart)
Michael Yang's avatar
Michael Yang committed
189
190
191
			}

			ch <- r
Michael Yang's avatar
Michael Yang committed
192
193
		}

194
195
		// an empty request loads the model
		if req.Prompt == "" && req.Template == "" && req.System == "" {
Patrick Devine's avatar
Patrick Devine committed
196
197
			ch <- api.GenerateResponse{Model: req.Model, Done: true}
		} else {
198
			if err := loaded.runner.Predict(c.Request.Context(), req.Context, prompt, fn); err != nil {
Patrick Devine's avatar
Patrick Devine committed
199
200
				ch <- gin.H{"error": err.Error()}
			}
Michael Yang's avatar
Michael Yang committed
201
		}
Michael Yang's avatar
Michael Yang committed
202
	}()
Michael Yang's avatar
Michael Yang committed
203

204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
	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
221
	streamResponse(c, ch)
Michael Yang's avatar
Michael Yang committed
222
}
Michael Yang's avatar
Michael Yang committed
223

Bruce MacDonald's avatar
Bruce MacDonald committed
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
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
	}
239
240
241

	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
242
243
244
245
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

246
	if !loaded.Options.EmbeddingOnly {
Bruce MacDonald's avatar
Bruce MacDonald committed
247
248
249
250
		c.JSON(http.StatusBadRequest, gin.H{"error": "embedding option must be set to true"})
		return
	}

251
	embedding, err := loaded.runner.Embedding(c.Request.Context(), req.Prompt)
Bruce MacDonald's avatar
Bruce MacDonald committed
252
253
254
255
256
257
258
259
260
261
262
263
	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)
}

264
func PullModelHandler(c *gin.Context) {
Michael Yang's avatar
Michael Yang committed
265
266
267
268
269
270
	var req api.PullRequest
	if err := c.ShouldBindJSON(&req); err != nil {
		c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
		return
	}

271
272
273
	ch := make(chan any)
	go func() {
		defer close(ch)
274
275
		fn := func(r api.ProgressResponse) {
			ch <- r
276
		}
277

278
279
280
281
		regOpts := &RegistryOptions{
			Insecure: req.Insecure,
		}

282
283
284
285
		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
286
			ch <- gin.H{"error": err.Error()}
287
288
289
		}
	}()

290
291
292
293
294
	if req.Stream != nil && !*req.Stream {
		waitForStream(c, ch)
		return
	}

295
296
297
	streamResponse(c, ch)
}

298
func PushModelHandler(c *gin.Context) {
299
300
301
	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
302
303
		return
	}
Michael Yang's avatar
Michael Yang committed
304

305
306
307
	ch := make(chan any)
	go func() {
		defer close(ch)
308
309
		fn := func(r api.ProgressResponse) {
			ch <- r
310
		}
311

312
313
314
315
		regOpts := &RegistryOptions{
			Insecure: req.Insecure,
		}

316
317
		ctx := context.Background()
		if err := PushModel(ctx, req.Name, regOpts, fn); err != nil {
Michael Yang's avatar
Michael Yang committed
318
			ch <- gin.H{"error": err.Error()}
319
320
321
		}
	}()

322
323
324
325
326
	if req.Stream != nil && !*req.Stream {
		waitForStream(c, ch)
		return
	}

327
328
329
	streamResponse(c, ch)
}

330
func CreateModelHandler(c *gin.Context) {
331
332
333
	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
334
		return
335
336
	}

337
338
	workDir := c.GetString("workDir")

Michael Yang's avatar
Michael Yang committed
339
	ch := make(chan any)
Michael Yang's avatar
Michael Yang committed
340
341
	go func() {
		defer close(ch)
342
343
		fn := func(resp api.ProgressResponse) {
			ch <- resp
344
345
		}

346
347
348
		ctx, cancel := context.WithCancel(c.Request.Context())
		defer cancel()

349
		if err := CreateModel(ctx, workDir, req.Name, req.Path, fn); err != nil {
Michael Yang's avatar
Michael Yang committed
350
			ch <- gin.H{"error": err.Error()}
351
		}
Michael Yang's avatar
Michael Yang committed
352
	}()
Michael Yang's avatar
Michael Yang committed
353

354
355
356
357
358
	if req.Stream != nil && !*req.Stream {
		waitForStream(c, ch)
		return
	}

Michael Yang's avatar
Michael Yang committed
359
	streamResponse(c, ch)
Bruce MacDonald's avatar
Bruce MacDonald committed
360
361
}

362
363
364
365
366
367
368
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
	}

369
370
371
372
	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 {
373
374
			c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		}
375
376
		return
	}
Michael Yang's avatar
Michael Yang committed
377
378
379
380
381
382
383
384
385
386
387
388

	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
	}

389
	c.JSON(http.StatusOK, nil)
390
391
}

Patrick Devine's avatar
Patrick Devine committed
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
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
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
}

463
func ListModelsHandler(c *gin.Context) {
464
	models := make([]api.ModelResponse, 0)
Patrick Devine's avatar
Patrick Devine committed
465
466
467
468
469
	fp, err := GetManifestPath()
	if err != nil {
		c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
		return
	}
Michael Yang's avatar
Michael Yang committed
470
471

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

477
			mp := ParseModelPath(tag)
Patrick Devine's avatar
Patrick Devine committed
478
			manifest, digest, err := GetManifest(mp)
Patrick Devine's avatar
Patrick Devine committed
479
			if err != nil {
480
481
				log.Printf("skipping file: %s", fp)
				return nil
Patrick Devine's avatar
Patrick Devine committed
482
			}
Michael Yang's avatar
Michael Yang committed
483
484

			models = append(models, api.ModelResponse{
Patrick Devine's avatar
Patrick Devine committed
485
486
				Name:       mp.GetShortTagname(),
				Size:       manifest.GetTotalSize(),
Patrick Devine's avatar
Patrick Devine committed
487
				Digest:     digest,
Michael Yang's avatar
Michael Yang committed
488
489
				ModifiedAt: info.ModTime(),
			})
Patrick Devine's avatar
Patrick Devine committed
490
		}
Michael Yang's avatar
Michael Yang committed
491

Patrick Devine's avatar
Patrick Devine committed
492
		return nil
Michael Yang's avatar
Michael Yang committed
493
494
495
	}

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

Michael Yang's avatar
Michael Yang committed
500
	c.JSON(http.StatusOK, api.ListResponse{Models: models})
Patrick Devine's avatar
Patrick Devine committed
501
502
}

Patrick Devine's avatar
Patrick Devine committed
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
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
520
521
522
523
524
525
526
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
527
528
	config := cors.DefaultConfig()
	config.AllowWildcard = true
Michael Yang's avatar
Michael Yang committed
529
530
531
532
533
534
535
536
537
538

	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
539

540
541
542
543
544
545
	workDir, err := os.MkdirTemp("", "ollama")
	if err != nil {
		return err
	}
	defer os.RemoveAll(workDir)

Bruce MacDonald's avatar
Bruce MacDonald committed
546
	r := gin.Default()
547
548
549
550
551
552
553
	r.Use(
		cors.New(config),
		func(c *gin.Context) {
			c.Set("workDir", workDir)
			c.Next()
		},
	)
Bruce MacDonald's avatar
Bruce MacDonald committed
554

555
556
	r.POST("/api/pull", PullModelHandler)
	r.POST("/api/generate", GenerateHandler)
Bruce MacDonald's avatar
Bruce MacDonald committed
557
	r.POST("/api/embeddings", EmbeddingHandler)
558
559
	r.POST("/api/create", CreateModelHandler)
	r.POST("/api/push", PushModelHandler)
Patrick Devine's avatar
Patrick Devine committed
560
	r.POST("/api/copy", CopyModelHandler)
561
	r.DELETE("/api/delete", DeleteModelHandler)
Patrick Devine's avatar
Patrick Devine committed
562
	r.POST("/api/show", ShowModelHandler)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
563

Michael Yang's avatar
Michael Yang committed
564
565
566
567
568
569
570
571
	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
572
	log.Printf("Listening on %s (version %s)", ln.Addr(), version.Version)
Jeffrey Morgan's avatar
Jeffrey Morgan committed
573
574
575
576
	s := &http.Server{
		Handler: r,
	}

577
578
	// listen for a ctrl+c and stop any loaded llm
	signals := make(chan os.Signal, 1)
579
	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
580
581
	go func() {
		<-signals
582
583
		if loaded.runner != nil {
			loaded.runner.Close()
584
		}
585
		os.RemoveAll(workDir)
586
587
588
		os.Exit(0)
	}()

589
590
591
	if runtime.GOOS == "linux" {
		// check compatibility to log warnings
		if _, err := llm.CheckVRAM(); err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
592
			log.Printf("Warning: GPU support may not enabled, check you have installed install GPU drivers: %v", err)
593
594
595
		}
	}

Jeffrey Morgan's avatar
Jeffrey Morgan committed
596
597
	return s.Serve(ln)
}
Michael Yang's avatar
Michael Yang committed
598

599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
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
624
func streamResponse(c *gin.Context, ch chan any) {
625
	c.Header("Content-Type", "application/x-ndjson")
Michael Yang's avatar
Michael Yang committed
626
627
628
629
630
631
632
633
	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
634
			log.Printf("streamResponse: json.Marshal failed with %s", err)
Michael Yang's avatar
Michael Yang committed
635
636
637
			return false
		}

638
		// Delineate chunks with new-line delimiter
Michael Yang's avatar
Michael Yang committed
639
640
		bts = append(bts, '\n')
		if _, err := w.Write(bts); err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
641
			log.Printf("streamResponse: w.Write failed with %s", err)
Michael Yang's avatar
Michael Yang committed
642
643
644
645
646
647
			return false
		}

		return true
	})
}