images.go 38 KB
Newer Older
1
2
3
package server

import (
4
	"bufio"
5
	"bytes"
6
	"context"
7
	"crypto/sha256"
Patrick Devine's avatar
Patrick Devine committed
8
	"encoding/hex"
9
10
11
12
13
14
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
Michael Yang's avatar
Michael Yang committed
15
	"net/url"
16
17
	"os"
	"path/filepath"
Michael Yang's avatar
Michael Yang committed
18
	"reflect"
Michael Yang's avatar
Michael Yang committed
19
	"runtime"
20
21
	"strconv"
	"strings"
Quinn Slack's avatar
Quinn Slack committed
22
	"text/template"
23

Michael Yang's avatar
Michael Yang committed
24
25
	"golang.org/x/exp/slices"

26
	"github.com/jmorganca/ollama/api"
27
	"github.com/jmorganca/ollama/llm"
28
	"github.com/jmorganca/ollama/parser"
29
	"github.com/jmorganca/ollama/vector"
Michael Yang's avatar
Michael Yang committed
30
	"github.com/jmorganca/ollama/version"
31
32
)

33
34
35
36
type RegistryOptions struct {
	Insecure bool
	Username string
	Password string
Patrick Devine's avatar
Patrick Devine committed
37
	Token    string
38
39
}

40
type Model struct {
Patrick Devine's avatar
Patrick Devine committed
41
42
43
44
45
46
47
48
49
50
51
52
	Name          string `json:"name"`
	ShortName     string
	ModelPath     string
	OriginalModel string
	AdapterPaths  []string
	Template      string
	System        string
	License       []string
	Digest        string
	ConfigDigest  string
	Options       map[string]interface{}
	Embeddings    []vector.Embedding
53
54
}

Bruce MacDonald's avatar
Bruce MacDonald committed
55
func (m *Model) Prompt(request api.GenerateRequest, embedding string) (string, error) {
56
57
58
59
60
61
	t := m.Template
	if request.Template != "" {
		t = request.Template
	}

	tmpl, err := template.New("").Parse(t)
62
63
64
65
66
	if err != nil {
		return "", err
	}

	var vars struct {
Michael Yang's avatar
Michael Yang committed
67
		First  bool
68
69
		System string
		Prompt string
70
		Embed  string
71
72
73

		// deprecated: versions <= 0.0.7 used this to omit the system prompt
		Context []int
74
75
	}

Michael Yang's avatar
Michael Yang committed
76
	vars.First = len(request.Context) == 0
77
78
	vars.System = m.System
	vars.Prompt = request.Prompt
Michael Yang's avatar
Michael Yang committed
79
	vars.Context = request.Context
Bruce MacDonald's avatar
Bruce MacDonald committed
80
	vars.Embed = embedding
81

82
83
84
85
	if request.System != "" {
		vars.System = request.System
	}

86
87
88
89
90
91
92
93
	var sb strings.Builder
	if err := tmpl.Execute(&sb, vars); err != nil {
		return "", err
	}

	return sb.String(), nil
}

94
95
96
97
98
99
100
101
102
103
type ManifestV2 struct {
	SchemaVersion int      `json:"schemaVersion"`
	MediaType     string   `json:"mediaType"`
	Config        Layer    `json:"config"`
	Layers        []*Layer `json:"layers"`
}

type Layer struct {
	MediaType string `json:"mediaType"`
	Digest    string `json:"digest"`
Michael Yang's avatar
Michael Yang committed
104
	Size      int64  `json:"size"`
Michael Yang's avatar
Michael Yang committed
105
	From      string `json:"from,omitempty"`
106
107
}

Michael Yang's avatar
Michael Yang committed
108
type LayerReader struct {
109
	Layer
Michael Yang's avatar
Michael Yang committed
110
	io.Reader
111
112
113
}

type ConfigV2 struct {
Michael Yang's avatar
Michael Yang committed
114
115
116
117
118
	ModelFormat string `json:"model_format"`
	ModelFamily string `json:"model_family"`
	ModelType   string `json:"model_type"`
	FileType    string `json:"file_type"`
	RootFS      RootFS `json:"rootfs"`
119
120

	// required by spec
121
122
123
124
125
126
127
128
129
	Architecture string `json:"architecture"`
	OS           string `json:"os"`
}

type RootFS struct {
	Type    string   `json:"type"`
	DiffIDs []string `json:"diff_ids"`
}

Michael Yang's avatar
Michael Yang committed
130
func (m *ManifestV2) GetTotalSize() (total int64) {
Patrick Devine's avatar
Patrick Devine committed
131
132
133
	for _, layer := range m.Layers {
		total += layer.Size
	}
Michael Yang's avatar
Michael Yang committed
134

Patrick Devine's avatar
Patrick Devine committed
135
136
137
138
	total += m.Config.Size
	return total
}

Patrick Devine's avatar
Patrick Devine committed
139
func GetManifest(mp ModelPath) (*ManifestV2, string, error) {
Patrick Devine's avatar
Patrick Devine committed
140
	fp, err := mp.GetManifestPath(false)
141
	if err != nil {
Patrick Devine's avatar
Patrick Devine committed
142
		return nil, "", err
143
	}
144

145
	if _, err = os.Stat(fp); err != nil {
Patrick Devine's avatar
Patrick Devine committed
146
		return nil, "", err
147
148
149
150
	}

	var manifest *ManifestV2

151
	bts, err := os.ReadFile(fp)
152
	if err != nil {
Patrick Devine's avatar
Patrick Devine committed
153
		return nil, "", fmt.Errorf("couldn't open file '%s'", fp)
154
155
	}

Patrick Devine's avatar
Patrick Devine committed
156
157
158
	shaSum := sha256.Sum256(bts)
	shaStr := hex.EncodeToString(shaSum[:])

159
	if err := json.Unmarshal(bts, &manifest); err != nil {
Patrick Devine's avatar
Patrick Devine committed
160
		return nil, "", err
161
162
	}

Patrick Devine's avatar
Patrick Devine committed
163
	return manifest, shaStr, nil
164
165
166
}

func GetModel(name string) (*Model, error) {
167
	mp := ParseModelPath(name)
Patrick Devine's avatar
Patrick Devine committed
168
	manifest, digest, err := GetManifest(mp)
169
170
171
172
173
	if err != nil {
		return nil, err
	}

	model := &Model{
Patrick Devine's avatar
Patrick Devine committed
174
		Name:         mp.GetFullTagname(),
Patrick Devine's avatar
Patrick Devine committed
175
		ShortName:    mp.GetShortTagname(),
Patrick Devine's avatar
Patrick Devine committed
176
177
178
		Digest:       digest,
		ConfigDigest: manifest.Config.Digest,
		Template:     "{{ .Prompt }}",
Patrick Devine's avatar
Patrick Devine committed
179
		License:      []string{},
180
181
182
	}

	for _, layer := range manifest.Layers {
Patrick Devine's avatar
Patrick Devine committed
183
		filename, err := GetBlobsPath(layer.Digest)
184
185
186
187
		if err != nil {
			return nil, err
		}

188
189
190
		switch layer.MediaType {
		case "application/vnd.ollama.image.model":
			model.ModelPath = filename
Patrick Devine's avatar
Patrick Devine committed
191
			model.OriginalModel = layer.From
192
193
194
195
196
197
198
199
200
201
		case "application/vnd.ollama.image.embed":
			file, err := os.Open(filename)
			if err != nil {
				return nil, fmt.Errorf("failed to open file: %s", filename)
			}
			defer file.Close()

			if err = json.NewDecoder(file).Decode(&model.Embeddings); err != nil {
				return nil, err
			}
202
203
		case "application/vnd.ollama.image.adapter":
			model.AdapterPaths = append(model.AdapterPaths, filename)
204
205
206
207
208
209
210
211
212
		case "application/vnd.ollama.image.template":
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}

			model.Template = string(bts)
		case "application/vnd.ollama.image.system":
			bts, err := os.ReadFile(filename)
213
214
215
			if err != nil {
				return nil, err
			}
216
217

			model.System = string(bts)
218
219
220
221
222
223
224
		case "application/vnd.ollama.image.prompt":
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}

			model.Template = string(bts)
225
		case "application/vnd.ollama.image.params":
Michael Yang's avatar
Michael Yang committed
226
227
228
229
230
			params, err := os.Open(filename)
			if err != nil {
				return nil, err
			}
			defer params.Close()
231

232
			// parse model options parameters into a map so that we can see which fields have been specified explicitly
233
			if err = json.NewDecoder(params).Decode(&model.Options); err != nil {
234
235
				return nil, err
			}
Patrick Devine's avatar
Patrick Devine committed
236
237
238
239
240
241
		case "application/vnd.ollama.image.license":
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}
			model.License = append(model.License, string(bts))
242
243
244
245
246
247
		}
	}

	return model, nil
}

248
249
func filenameWithPath(path, f string) (string, error) {
	// if filePath starts with ~/, replace it with the user's home directory.
Michael Yang's avatar
Michael Yang committed
250
251
	if strings.HasPrefix(f, fmt.Sprintf("~%s", string(os.PathSeparator))) {
		parts := strings.Split(f, string(os.PathSeparator))
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
		home, err := os.UserHomeDir()
		if err != nil {
			return "", fmt.Errorf("failed to open file: %v", err)
		}

		f = filepath.Join(home, filepath.Join(parts[1:]...))
	}

	// if filePath is not an absolute path, make it relative to the modelfile path
	if !filepath.IsAbs(f) {
		f = filepath.Join(filepath.Dir(path), f)
	}

	return f, nil
}

268
func CreateModel(ctx context.Context, workDir, name string, path string, fn func(resp api.ProgressResponse)) error {
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
	mp := ParseModelPath(name)

	var manifest *ManifestV2
	var err error
	var noprune string

	// build deleteMap to prune unused layers
	deleteMap := make(map[string]bool)

	if noprune = os.Getenv("OLLAMA_NOPRUNE"); noprune == "" {
		manifest, _, err = GetManifest(mp)
		if err != nil && !errors.Is(err, os.ErrNotExist) {
			return err
		}

		if manifest != nil {
			for _, l := range manifest.Layers {
				deleteMap[l.Digest] = true
			}
			deleteMap[manifest.Config.Digest] = true
		}
	}

292
293
	mf, err := os.Open(path)
	if err != nil {
294
		fn(api.ProgressResponse{Status: fmt.Sprintf("couldn't open modelfile '%s'", path)})
295
		return fmt.Errorf("failed to open file: %w", err)
296
	}
297
	defer mf.Close()
298

299
	fn(api.ProgressResponse{Status: "parsing modelfile"})
300
301
302
303
304
	commands, err := parser.Parse(mf)
	if err != nil {
		return err
	}

305
306
307
308
309
	config := ConfigV2{
		Architecture: "amd64",
		OS:           "linux",
	}

Michael Yang's avatar
Michael Yang committed
310
	var layers []*LayerReader
311
	params := make(map[string][]string)
Michael Yang's avatar
Michael Yang committed
312
	var sourceParams map[string]any
313
	embed := EmbeddingParams{fn: fn}
314
	for _, c := range commands {
315
		log.Printf("[%s] - %s\n", c.Name, c.Args)
316
317
		switch c.Name {
		case "model":
318
			fn(api.ProgressResponse{Status: "looking for model"})
319
			embed.model = c.Args
320

321
			mp := ParseModelPath(c.Args)
Patrick Devine's avatar
Patrick Devine committed
322
			mf, _, err := GetManifest(mp)
323
			if err != nil {
324
325
326
				modelFile, err := filenameWithPath(path, c.Args)
				if err != nil {
					return err
327
				}
328
				if _, err := os.Stat(modelFile); err != nil {
329
330
331
					// the model file does not exist, try pulling it
					if errors.Is(err, os.ErrNotExist) {
						fn(api.ProgressResponse{Status: "pulling model file"})
332
						if err := PullModel(ctx, c.Args, &RegistryOptions{}, fn); err != nil {
333
334
							return err
						}
Patrick Devine's avatar
Patrick Devine committed
335
						mf, _, err = GetManifest(mp)
336
337
338
339
340
341
342
						if err != nil {
							return fmt.Errorf("failed to open file after pull: %v", err)
						}
					} else {
						return err
					}
				} else {
343
					embed.model = modelFile
344
345
					// create a model from this specified file
					fn(api.ProgressResponse{Status: "creating model layer"})
346
					file, err := os.Open(modelFile)
347
348
349
350
351
					if err != nil {
						return fmt.Errorf("failed to open file: %v", err)
					}
					defer file.Close()

Bruce MacDonald's avatar
Bruce MacDonald committed
352
					ggml, err := llm.DecodeGGML(file)
353
354
355
356
					if err != nil {
						return err
					}

357
					config.ModelFormat = ggml.Name()
Michael Yang's avatar
Michael Yang committed
358
359
360
					config.ModelFamily = ggml.ModelFamily()
					config.ModelType = ggml.ModelType()
					config.FileType = ggml.FileType()
361
362
363
364

					// reset the file
					file.Seek(0, io.SeekStart)

365
366
367
368
369
370
					l, err := CreateLayer(file)
					if err != nil {
						return fmt.Errorf("failed to create layer: %v", err)
					}
					l.MediaType = "application/vnd.ollama.image.model"
					layers = append(layers, l)
371
				}
372
			}
373

374
			if mf != nil {
Michael Yang's avatar
Michael Yang committed
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
				sourceBlobPath, err := GetBlobsPath(mf.Config.Digest)
				if err != nil {
					return err
				}

				sourceBlob, err := os.Open(sourceBlobPath)
				if err != nil {
					return err
				}
				defer sourceBlob.Close()

				var source ConfigV2
				if err := json.NewDecoder(sourceBlob).Decode(&source); err != nil {
					return err
				}

391
				// copy the model metadata
Michael Yang's avatar
Michael Yang committed
392
393
				config.ModelFamily = source.ModelFamily
				config.ModelType = source.ModelType
394
				config.ModelFormat = source.ModelFormat
Michael Yang's avatar
Michael Yang committed
395
396
				config.FileType = source.FileType

397
				for _, l := range mf.Layers {
Michael Yang's avatar
Michael Yang committed
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
					if l.MediaType == "application/vnd.ollama.image.params" {
						sourceParamsBlobPath, err := GetBlobsPath(l.Digest)
						if err != nil {
							return err
						}

						sourceParamsBlob, err := os.Open(sourceParamsBlobPath)
						if err != nil {
							return err
						}
						defer sourceParamsBlob.Close()

						if err := json.NewDecoder(sourceParamsBlob).Decode(&sourceParams); err != nil {
							return err
						}
					}

415
416
417
418
					newLayer, err := GetLayerWithBufferFromLayer(l)
					if err != nil {
						return err
					}
Michael Yang's avatar
Michael Yang committed
419
					newLayer.From = mp.GetNamespaceRepository()
420
421
422
					layers = append(layers, newLayer)
				}
			}
423
424
		case "embed":
			embedFilePath, err := filenameWithPath(path, c.Args)
Michael Yang's avatar
Michael Yang committed
425
426
427
			if err != nil {
				return err
			}
428
			embed.files = append(embed.files, embedFilePath)
429
430
431
		case "adapter":
			fn(api.ProgressResponse{Status: fmt.Sprintf("creating model %s layer", c.Name)})

Michael Yang's avatar
Michael Yang committed
432
433
434
			fp, err := filenameWithPath(path, c.Args)
			if err != nil {
				return err
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
			}

			// create a model from this specified file
			fn(api.ProgressResponse{Status: "creating model layer"})

			file, err := os.Open(fp)
			if err != nil {
				return fmt.Errorf("failed to open file: %v", err)
			}
			defer file.Close()

			l, err := CreateLayer(file)
			if err != nil {
				return fmt.Errorf("failed to create layer: %v", err)
			}
			l.MediaType = "application/vnd.ollama.image.adapter"
			layers = append(layers, l)
Bruce MacDonald's avatar
Bruce MacDonald committed
452
453
454
455
456
457
458
459
460
		case "license":
			fn(api.ProgressResponse{Status: fmt.Sprintf("creating model %s layer", c.Name)})
			mediaType := fmt.Sprintf("application/vnd.ollama.image.%s", c.Name)

			layer, err := CreateLayer(strings.NewReader(c.Args))
			if err != nil {
				return err
			}

461
462
463
464
			if layer.Size > 0 {
				layer.MediaType = mediaType
				layers = append(layers, layer)
			}
Bruce MacDonald's avatar
Bruce MacDonald committed
465
		case "template", "system", "prompt":
466
			fn(api.ProgressResponse{Status: fmt.Sprintf("creating model %s layer", c.Name)})
467
			// remove the layer if one exists
468
469
			mediaType := fmt.Sprintf("application/vnd.ollama.image.%s", c.Name)
			layers = removeLayerFromLayers(layers, mediaType)
470

471
			layer, err := CreateLayer(strings.NewReader(c.Args))
472
			if err != nil {
473
				return err
474
			}
475

476
477
478
479
			if layer.Size > 0 {
				layer.MediaType = mediaType
				layers = append(layers, layer)
			}
480
		default:
481
			// runtime parameters, build a list of args for each parameter to allow multiple values to be specified (ex: multiple stop sequences)
482
			params[c.Name] = append(params[c.Name], c.Args)
483
484
485
486
		}
	}

	// Create a single layer for the parameters
Michael Yang's avatar
Michael Yang committed
487
	if len(params) > 0 {
488
		fn(api.ProgressResponse{Status: "creating parameter layer"})
Michael Yang's avatar
Michael Yang committed
489

490
		layers = removeLayerFromLayers(layers, "application/vnd.ollama.image.params")
491
		formattedParams, err := formatParams(params)
492
493
494
		if err != nil {
			return fmt.Errorf("couldn't create params json: %v", err)
		}
495

Michael Yang's avatar
Michael Yang committed
496
497
498
499
500
501
		for k, v := range sourceParams {
			if _, ok := formattedParams[k]; !ok {
				formattedParams[k] = v
			}
		}

Michael Yang's avatar
Michael Yang committed
502
503
504
505
506
507
		if config.ModelType == "65B" {
			if numGQA, ok := formattedParams["num_gqa"].(int); ok && numGQA == 8 {
				config.ModelType = "70B"
			}
		}

508
509
510
511
512
513
		bts, err := json.Marshal(formattedParams)
		if err != nil {
			return err
		}

		l, err := CreateLayer(bytes.NewReader(bts))
514
515
516
517
518
		if err != nil {
			return fmt.Errorf("failed to create layer: %v", err)
		}
		l.MediaType = "application/vnd.ollama.image.params"
		layers = append(layers, l)
519
520

		// apply these parameters to the embedding options, in case embeddings need to be generated using this model
521
		embed.opts = formattedParams
522
523
	}

524
	// generate the embedding layers
525
	embeddingLayers, err := embeddingLayers(workDir, embed)
526
527
528
529
530
	if err != nil {
		return err
	}
	layers = append(layers, embeddingLayers...)

531
532
533
534
535
536
537
538
	digests, err := getLayerDigests(layers)
	if err != nil {
		return err
	}

	var manifestLayers []*Layer
	for _, l := range layers {
		manifestLayers = append(manifestLayers, &l.Layer)
539
		delete(deleteMap, l.Layer.Digest)
540
541
542
	}

	// Create a layer for the config object
543
	fn(api.ProgressResponse{Status: "creating config layer"})
544
	cfg, err := createConfigLayer(config, digests)
545
546
547
548
	if err != nil {
		return err
	}
	layers = append(layers, cfg)
549
	delete(deleteMap, cfg.Layer.Digest)
550

Michael Yang's avatar
Michael Yang committed
551
	if err := SaveLayers(layers, fn, false); err != nil {
552
553
554
555
		return err
	}

	// Create the manifest
556
	fn(api.ProgressResponse{Status: "writing manifest"})
557
558
559
560
561
	err = CreateManifest(name, cfg, manifestLayers)
	if err != nil {
		return err
	}

562
563
564
565
566
567
568
569
	if noprune == "" {
		fn(api.ProgressResponse{Status: "removing any unused layers"})
		err = deleteUnusedLayers(nil, deleteMap, false)
		if err != nil {
			return err
		}
	}

570
	fn(api.ProgressResponse{Status: "success"})
571
572
573
	return nil
}

574
575
type EmbeddingParams struct {
	model string
576
	opts  map[string]interface{}
577
578
579
580
581
	files []string // paths to files to embed
	fn    func(resp api.ProgressResponse)
}

// embeddingLayers loads the associated LLM and generates the embeddings to be stored from an input file
582
func embeddingLayers(workDir string, e EmbeddingParams) ([]*LayerReader, error) {
583
584
	layers := []*LayerReader{}
	if len(e.files) > 0 {
585
586
587
588
589
		// check if the model is a file path or a model name
		model, err := GetModel(e.model)
		if err != nil {
			if !strings.Contains(err.Error(), "couldn't open file") {
				return nil, fmt.Errorf("unexpected error opening model to generate embeddings: %v", err)
590
			}
591
592
			// the model may be a file path, create a model from this file
			model = &Model{ModelPath: e.model}
593
594
		}

595
		if err := load(context.Background(), workDir, model, e.opts, defaultSessionDuration); err != nil {
596
597
598
			return nil, fmt.Errorf("load model to generate embeddings: %v", err)
		}

599
		// this will be used to check if we already have embeddings for a file
600
		modelInfo, err := os.Stat(model.ModelPath)
601
		if err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
602
			return nil, fmt.Errorf("failed to get model file info: %v", err)
603
604
		}

Bruce MacDonald's avatar
Bruce MacDonald committed
605
606
607
		addedFiles := make(map[string]bool) // keep track of files that have already been added
		for _, filePattern := range e.files {
			matchingFiles, err := filepath.Glob(filePattern)
608
			if err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
609
				return nil, fmt.Errorf("could not find files with pattern %s: %w", filePattern, err)
610
611
			}

Bruce MacDonald's avatar
Bruce MacDonald committed
612
613
			for _, filePath := range matchingFiles {
				if addedFiles[filePath] {
614
615
					continue
				}
Bruce MacDonald's avatar
Bruce MacDonald committed
616
				addedFiles[filePath] = true
Bruce MacDonald's avatar
Bruce MacDonald committed
617
				// check if we already have embeddings for this file path
618
				layerIdentifier := fmt.Sprintf("%s:%s:%s:%d", filePath, e.model, modelInfo.ModTime().Format("2006-01-02 15:04:05"), modelInfo.Size())
Bruce MacDonald's avatar
Bruce MacDonald committed
619
620
621
622
623
624
				digest, _ := GetSHA256Digest(strings.NewReader(layerIdentifier))
				existing, err := existingFileEmbeddings(digest)
				if err != nil {
					return nil, fmt.Errorf("failed to check existing embeddings for file %s: %v", filePath, err)
				}

Bruce MacDonald's avatar
Bruce MacDonald committed
625
626
				// TODO: check file type
				f, err := os.Open(filePath)
627
				if err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
628
					return nil, fmt.Errorf("could not open embed file: %w", err)
629
				}
Bruce MacDonald's avatar
Bruce MacDonald committed
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
				scanner := bufio.NewScanner(f)
				scanner.Split(bufio.ScanLines)

				data := []string{}
				for scanner.Scan() {
					data = append(data, scanner.Text())
				}
				f.Close()

				// the digest of the file is set here so that the client knows a new operation is in progress
				fileDigest, _ := GetSHA256Digest(bytes.NewReader([]byte(filePath)))

				embeddings := []vector.Embedding{}
				for i, d := range data {
					if strings.TrimSpace(d) == "" {
						continue
					}
					e.fn(api.ProgressResponse{
						Status:    fmt.Sprintf("creating embeddings for file %s", filePath),
						Digest:    fileDigest,
Michael Yang's avatar
Michael Yang committed
650
651
						Total:     int64(len(data) - 1),
						Completed: int64(i),
Bruce MacDonald's avatar
Bruce MacDonald committed
652
					})
Bruce MacDonald's avatar
Bruce MacDonald committed
653
654
655
656
657
					if len(existing[d]) > 0 {
						// already have an embedding for this line
						embeddings = append(embeddings, vector.Embedding{Data: d, Vector: existing[d]})
						continue
					}
658
					embed, err := loaded.llm.Embedding(context.Background(), d)
Bruce MacDonald's avatar
Bruce MacDonald committed
659
					if err != nil {
660
661
						log.Printf("failed to generate embedding for '%s' line %d: %v", filePath, i+1, err)
						continue
Bruce MacDonald's avatar
Bruce MacDonald committed
662
663
					}
					embeddings = append(embeddings, vector.Embedding{Data: d, Vector: embed})
664
665
				}

Bruce MacDonald's avatar
Bruce MacDonald committed
666
667
668
669
670
				b, err := json.Marshal(embeddings)
				if err != nil {
					return nil, fmt.Errorf("failed to encode embeddings: %w", err)
				}
				r := bytes.NewReader(b)
671

Bruce MacDonald's avatar
Bruce MacDonald committed
672
673
674
675
				layer := &LayerReader{
					Layer: Layer{
						MediaType: "application/vnd.ollama.image.embed",
						Digest:    digest,
Michael Yang's avatar
Michael Yang committed
676
						Size:      r.Size(),
Bruce MacDonald's avatar
Bruce MacDonald committed
677
678
679
					},
					Reader: r,
				}
680

Bruce MacDonald's avatar
Bruce MacDonald committed
681
682
				layers = append(layers, layer)
			}
683
684
685
686
687
		}
	}
	return layers, nil
}

Bruce MacDonald's avatar
Bruce MacDonald committed
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
// existingFileEmbeddings checks if we already have embeddings for a file and loads them into a look-up map
func existingFileEmbeddings(digest string) (map[string][]float64, error) {
	path, err := GetBlobsPath(digest)
	if err != nil {
		return nil, fmt.Errorf("embeddings blobs path: %w", err)
	}
	existingFileEmbeddings := make(map[string][]float64)
	if _, err := os.Stat(path); err == nil {
		// already have some embeddings for this file, load embeddings previously generated
		file, err := os.Open(path)
		if err != nil {
			return nil, fmt.Errorf("failed to open existing embedding file: %s", err)
		}
		defer file.Close()

		existing := []vector.Embedding{}
		if err = json.NewDecoder(file).Decode(&existing); err != nil {
			return nil, err
		}
		for _, e := range existing {
			existingFileEmbeddings[e.Data] = e.Vector
		}
	}
	return existingFileEmbeddings, nil
}

Michael Yang's avatar
Michael Yang committed
714
func removeLayerFromLayers(layers []*LayerReader, mediaType string) []*LayerReader {
Michael Yang's avatar
Michael Yang committed
715
716
717
	return slices.DeleteFunc(layers, func(layer *LayerReader) bool {
		return layer.MediaType == mediaType
	})
718
719
}

720
func SaveLayers(layers []*LayerReader, fn func(resp api.ProgressResponse), force bool) error {
721
722
	// Write each of the layers to disk
	for _, layer := range layers {
Patrick Devine's avatar
Patrick Devine committed
723
		fp, err := GetBlobsPath(layer.Digest)
724
725
726
		if err != nil {
			return err
		}
727
728

		_, err = os.Stat(fp)
Bruce MacDonald's avatar
Bruce MacDonald committed
729
730
		// note: embed layers are always written since their digest doesnt indicate anything about the contents
		if os.IsNotExist(err) || force || layer.MediaType == "application/vnd.ollama.image.embed" {
731
732
			fn(api.ProgressResponse{Status: fmt.Sprintf("writing layer %s", layer.Digest)})

733
734
735
736
737
738
739
			out, err := os.Create(fp)
			if err != nil {
				log.Printf("couldn't create %s", fp)
				return err
			}
			defer out.Close()

Michael Yang's avatar
Michael Yang committed
740
			if _, err = io.Copy(out, layer.Reader); err != nil {
741
742
				return err
			}
Michael Yang's avatar
Michael Yang committed
743

744
		} else {
745
			fn(api.ProgressResponse{Status: fmt.Sprintf("using already created layer %s", layer.Digest)})
746
747
748
749
750
751
		}
	}

	return nil
}

Michael Yang's avatar
Michael Yang committed
752
func CreateManifest(name string, cfg *LayerReader, layers []*Layer) error {
753
	mp := ParseModelPath(name)
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
	manifest := ManifestV2{
		SchemaVersion: 2,
		MediaType:     "application/vnd.docker.distribution.manifest.v2+json",
		Config: Layer{
			MediaType: cfg.MediaType,
			Size:      cfg.Size,
			Digest:    cfg.Digest,
		},
		Layers: layers,
	}

	manifestJSON, err := json.Marshal(manifest)
	if err != nil {
		return err
	}

Patrick Devine's avatar
Patrick Devine committed
770
	fp, err := mp.GetManifestPath(true)
771
772
773
	if err != nil {
		return err
	}
774
	return os.WriteFile(fp, manifestJSON, 0o644)
775
776
}

Michael Yang's avatar
Michael Yang committed
777
func GetLayerWithBufferFromLayer(layer *Layer) (*LayerReader, error) {
Patrick Devine's avatar
Patrick Devine committed
778
	fp, err := GetBlobsPath(layer.Digest)
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
	if err != nil {
		return nil, err
	}

	file, err := os.Open(fp)
	if err != nil {
		return nil, fmt.Errorf("could not open blob: %w", err)
	}
	defer file.Close()

	newLayer, err := CreateLayer(file)
	if err != nil {
		return nil, err
	}
	newLayer.MediaType = layer.MediaType
	return newLayer, nil
}

797
798
// formatParams converts specified parameter options to their correct types
func formatParams(params map[string][]string) (map[string]interface{}, error) {
799
800
801
	opts := api.Options{}
	valueOpts := reflect.ValueOf(&opts).Elem() // names of the fields in the options struct
	typeOpts := reflect.TypeOf(opts)           // types of the fields in the options struct
Michael Yang's avatar
Michael Yang committed
802

803
	// build map of json struct tags to their types
Michael Yang's avatar
Michael Yang committed
804
805
806
807
808
809
810
811
	jsonOpts := make(map[string]reflect.StructField)
	for _, field := range reflect.VisibleFields(typeOpts) {
		jsonTag := strings.Split(field.Tag.Get("json"), ",")[0]
		if jsonTag != "" {
			jsonOpts[jsonTag] = field
		}
	}

812
	out := make(map[string]interface{})
Michael Yang's avatar
Michael Yang committed
813
	// iterate params and set values based on json struct tags
814
	for key, vals := range params {
Michael Yang's avatar
Michael Yang committed
815
816
817
818
819
		if opt, ok := jsonOpts[key]; ok {
			field := valueOpts.FieldByName(opt.Name)
			if field.IsValid() && field.CanSet() {
				switch field.Kind() {
				case reflect.Float32:
820
					floatVal, err := strconv.ParseFloat(vals[0], 32)
Michael Yang's avatar
Michael Yang committed
821
					if err != nil {
822
						return nil, fmt.Errorf("invalid float value %s", vals)
Michael Yang's avatar
Michael Yang committed
823
824
					}

Michael Yang's avatar
Michael Yang committed
825
					out[key] = float32(floatVal)
Michael Yang's avatar
Michael Yang committed
826
				case reflect.Int:
Michael Yang's avatar
Michael Yang committed
827
					intVal, err := strconv.ParseInt(vals[0], 10, 64)
Michael Yang's avatar
Michael Yang committed
828
					if err != nil {
829
						return nil, fmt.Errorf("invalid int value %s", vals)
Michael Yang's avatar
Michael Yang committed
830
831
					}

832
					out[key] = intVal
Michael Yang's avatar
Michael Yang committed
833
				case reflect.Bool:
834
					boolVal, err := strconv.ParseBool(vals[0])
Michael Yang's avatar
Michael Yang committed
835
					if err != nil {
836
						return nil, fmt.Errorf("invalid bool value %s", vals)
Michael Yang's avatar
Michael Yang committed
837
838
					}

839
					out[key] = boolVal
Michael Yang's avatar
Michael Yang committed
840
				case reflect.String:
841
					out[key] = vals[0]
842
				case reflect.Slice:
843
844
					// TODO: only string slices are supported right now
					out[key] = vals
Michael Yang's avatar
Michael Yang committed
845
846
847
848
849
850
851
				default:
					return nil, fmt.Errorf("unknown type %s for %s", field.Kind(), key)
				}
			}
		}
	}

852
	return out, nil
853
854
}

Michael Yang's avatar
Michael Yang committed
855
func getLayerDigests(layers []*LayerReader) ([]string, error) {
856
857
858
859
860
861
862
863
864
865
866
	var digests []string
	for _, l := range layers {
		if l.Digest == "" {
			return nil, fmt.Errorf("layer is missing a digest")
		}
		digests = append(digests, l.Digest)
	}
	return digests, nil
}

// CreateLayer creates a Layer object from a given file
Michael Yang's avatar
Michael Yang committed
867
868
func CreateLayer(f io.ReadSeeker) (*LayerReader, error) {
	digest, size := GetSHA256Digest(f)
869
	f.Seek(0, io.SeekStart)
870

Michael Yang's avatar
Michael Yang committed
871
	layer := &LayerReader{
872
873
874
875
876
		Layer: Layer{
			MediaType: "application/vnd.docker.image.rootfs.diff.tar",
			Digest:    digest,
			Size:      size,
		},
Michael Yang's avatar
Michael Yang committed
877
		Reader: f,
878
879
880
881
882
	}

	return layer, nil
}

Patrick Devine's avatar
Patrick Devine committed
883
func CopyModel(src, dest string) error {
884
	srcModelPath := ParseModelPath(src)
885
886
887
888
889
	srcPath, err := srcModelPath.GetManifestPath(false)
	if err != nil {
		return err
	}

890
	destModelPath := ParseModelPath(dest)
891
	destPath, err := destModelPath.GetManifestPath(true)
Patrick Devine's avatar
Patrick Devine committed
892
893
894
895
896
	if err != nil {
		return err
	}

	// copy the file
Michael Yang's avatar
Michael Yang committed
897
	input, err := os.ReadFile(srcPath)
Patrick Devine's avatar
Patrick Devine committed
898
899
900
901
902
	if err != nil {
		fmt.Println("Error reading file:", err)
		return err
	}

Michael Yang's avatar
Michael Yang committed
903
	err = os.WriteFile(destPath, input, 0o644)
Patrick Devine's avatar
Patrick Devine committed
904
905
906
907
908
909
910
911
	if err != nil {
		fmt.Println("Error reading file:", err)
		return err
	}

	return nil
}

912
func deleteUnusedLayers(skipModelPath *ModelPath, deleteMap map[string]bool, dryRun bool) error {
913
914
915
916
	fp, err := GetManifestPath()
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
917
918
919
920

	walkFunc := func(path string, info os.FileInfo, _ error) error {
		if info.IsDir() {
			return nil
921
922
		}

Michael Yang's avatar
Michael Yang committed
923
924
925
926
		dir, file := filepath.Split(path)
		dir = strings.Trim(strings.TrimPrefix(dir, fp), string(os.PathSeparator))
		tag := strings.Join([]string{dir, file}, ":")
		fmp := ParseModelPath(tag)
927

Michael Yang's avatar
Michael Yang committed
928
		// skip the manifest we're trying to delete
929
		if skipModelPath != nil && skipModelPath.GetFullTagname() == fmp.GetFullTagname() {
Michael Yang's avatar
Michael Yang committed
930
			return nil
931
		}
Michael Yang's avatar
Michael Yang committed
932
933
934
935
936
937
938
939
940
941
942
943

		// save (i.e. delete from the deleteMap) any files used in other manifests
		manifest, _, err := GetManifest(fmp)
		if err != nil {
			return nil
		}

		for _, layer := range manifest.Layers {
			delete(deleteMap, layer.Digest)
		}

		delete(deleteMap, manifest.Config.Digest)
944
		return nil
Michael Yang's avatar
Michael Yang committed
945
946
947
	}

	if err := filepath.Walk(fp, walkFunc); err != nil {
Michael Yang's avatar
Michael Yang committed
948
949
		return err
	}
950
951
952
953

	// only delete the files which are still in the deleteMap
	for k, v := range deleteMap {
		if v {
954
			fp, err := GetBlobsPath(k)
955
			if err != nil {
956
957
958
				log.Printf("couldn't get file path for '%s': %v", k, err)
				continue
			}
959
960
961
962
963
964
965
			if !dryRun {
				if err := os.Remove(fp); err != nil {
					log.Printf("couldn't remove file '%s': %v", fp, err)
					continue
				}
			} else {
				log.Printf("wanted to remove: %s", fp)
966
967
968
969
			}
		}
	}

970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
	return nil
}

func PruneLayers() error {
	deleteMap := make(map[string]bool)
	p, err := GetBlobsPath("")
	if err != nil {
		return err
	}

	blobs, err := os.ReadDir(p)
	if err != nil {
		log.Printf("couldn't read dir '%s': %v", p, err)
		return err
	}

	for _, blob := range blobs {
		name := blob.Name()
		if runtime.GOOS == "windows" {
			name = strings.ReplaceAll(name, "-", ":")
		}
		deleteMap[name] = true
	}

	log.Printf("total blobs: %d", len(deleteMap))

	err = deleteUnusedLayers(nil, deleteMap, false)
	if err != nil {
		return err
	}

	log.Printf("total unused blobs removed: %d", len(deleteMap))

	return nil
}

Michael Yang's avatar
Michael Yang committed
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
func PruneDirectory(path string) error {
	info, err := os.Lstat(path)
	if err != nil {
		return err
	}

	if info.IsDir() && info.Mode()&os.ModeSymlink == 0 {
		entries, err := os.ReadDir(path)
		if err != nil {
			return err
		}

		for _, entry := range entries {
			if err := PruneDirectory(filepath.Join(path, entry.Name())); err != nil {
				return err
			}
		}

		entries, err = os.ReadDir(path)
		if err != nil {
			return err
		}

		if len(entries) > 0 {
			return nil
		}

		return os.Remove(path)
	}

	return nil
}

1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
func DeleteModel(name string) error {
	mp := ParseModelPath(name)
	manifest, _, err := GetManifest(mp)
	if err != nil {
		return err
	}

	deleteMap := make(map[string]bool)
	for _, layer := range manifest.Layers {
		deleteMap[layer.Digest] = true
	}
	deleteMap[manifest.Config.Digest] = true

	err = deleteUnusedLayers(&mp, deleteMap, false)
	if err != nil {
		return err
	}

	fp, err := mp.GetManifestPath(false)
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
	if err != nil {
		return err
	}
	err = os.Remove(fp)
	if err != nil {
		log.Printf("couldn't remove manifest file '%s': %v", fp, err)
		return err
	}

	return nil
}

Patrick Devine's avatar
Patrick Devine committed
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
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
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
func ShowModelfile(model *Model) (string, error) {
	type modelTemplate struct {
		*Model
		From   string
		Params string
	}

	var params []string
	for k, v := range model.Options {
		switch val := v.(type) {
		case string:
			params = append(params, fmt.Sprintf("PARAMETER %s %s", k, val))
		case int:
			params = append(params, fmt.Sprintf("PARAMETER %s %s", k, strconv.Itoa(val)))
		case float64:
			params = append(params, fmt.Sprintf("PARAMETER %s %s", k, strconv.FormatFloat(val, 'f', 0, 64)))
		case bool:
			params = append(params, fmt.Sprintf("PARAMETER %s %s", k, strconv.FormatBool(val)))
		case []interface{}:
			for _, nv := range val {
				switch nval := nv.(type) {
				case string:
					params = append(params, fmt.Sprintf("PARAMETER %s %s", k, nval))
				case int:
					params = append(params, fmt.Sprintf("PARAMETER %s %s", k, strconv.Itoa(nval)))
				case float64:
					params = append(params, fmt.Sprintf("PARAMETER %s %s", k, strconv.FormatFloat(nval, 'f', 0, 64)))
				case bool:
					params = append(params, fmt.Sprintf("PARAMETER %s %s", k, strconv.FormatBool(nval)))
				default:
					log.Printf("unknown type: %s", reflect.TypeOf(nv).String())
				}
			}
		default:
			log.Printf("unknown type: %s", reflect.TypeOf(v).String())
		}
	}

	mt := modelTemplate{
		Model:  model,
		From:   model.OriginalModel,
		Params: strings.Join(params, "\n"),
	}

	if mt.From == "" {
		mt.From = model.ModelPath
	}

	modelFile := `# Modelfile generated by "ollama show"
# To build a new Modelfile based on this one, replace the FROM line with:
# FROM {{ .ShortName }}

FROM {{ .From }}
TEMPLATE """{{ .Template }}"""
SYSTEM """{{ .System }}"""
{{ .Params }}
`
	for _, l := range mt.Model.AdapterPaths {
		modelFile += fmt.Sprintf("ADAPTER %s\n", l)
	}

	tmpl, err := template.New("").Parse(modelFile)
	if err != nil {
		log.Printf("error parsing template: %q", err)
		return "", err
	}

	var buf bytes.Buffer

	if err = tmpl.Execute(&buf, mt); err != nil {
		log.Printf("error executing template: %q", err)
		return "", err
	}

	return buf.String(), nil
}

1147
func PushModel(ctx context.Context, name string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
1148
	mp := ParseModelPath(name)
1149
1150
	fn(api.ProgressResponse{Status: "retrieving manifest"})

1151
1152
1153
1154
	if mp.ProtocolScheme == "http" && !regOpts.Insecure {
		return fmt.Errorf("insecure protocol http")
	}

Patrick Devine's avatar
Patrick Devine committed
1155
	manifest, _, err := GetManifest(mp)
1156
	if err != nil {
1157
		fn(api.ProgressResponse{Status: "couldn't retrieve manifest"})
1158
1159
1160
1161
		return err
	}

	var layers []*Layer
Jeffrey Morgan's avatar
Jeffrey Morgan committed
1162
	layers = append(layers, manifest.Layers...)
1163
1164
1165
	layers = append(layers, &manifest.Config)

	for _, layer := range layers {
1166
		exists, err := checkBlobExistence(ctx, mp, layer.Digest, regOpts)
1167
1168
1169
1170
1171
		if err != nil {
			return err
		}

		if exists {
1172
1173
1174
			fn(api.ProgressResponse{
				Status:    "using existing layer",
				Digest:    layer.Digest,
1175
1176
				Total:     layer.Size,
				Completed: layer.Size,
1177
			})
1178
			log.Printf("Layer %s already exists", layer.Digest)
1179
1180
1181
			continue
		}

1182
		fn(api.ProgressResponse{
1183
1184
1185
			Status: "starting upload",
			Digest: layer.Digest,
			Total:  layer.Size,
1186
		})
1187

1188
		location, chunkSize, err := startUpload(ctx, mp, layer, regOpts)
1189
1190
1191
1192
1193
		if err != nil {
			log.Printf("couldn't start upload: %v", err)
			return err
		}

Michael Yang's avatar
Michael Yang committed
1194
1195
		if strings.HasPrefix(filepath.Base(location.Path), "sha256:") {
			layer.Digest = filepath.Base(location.Path)
Michael Yang's avatar
Michael Yang committed
1196
1197
1198
1199
1200
1201
1202
1203
1204
			fn(api.ProgressResponse{
				Status:    "using existing layer",
				Digest:    layer.Digest,
				Total:     layer.Size,
				Completed: layer.Size,
			})
			continue
		}

1205
		if err := uploadBlob(ctx, location, layer, chunkSize, regOpts, fn); err != nil {
1206
1207
1208
			log.Printf("error uploading blob: %v", err)
			return err
		}
1209
1210
	}

1211
	fn(api.ProgressResponse{Status: "pushing manifest"})
Michael Yang's avatar
Michael Yang committed
1212
1213
	requestURL := mp.BaseURL()
	requestURL = requestURL.JoinPath("v2", mp.GetNamespaceRepository(), "manifests", mp.Tag)
1214
1215
1216
1217
1218
1219

	manifestJSON, err := json.Marshal(manifest)
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
1220
1221
	headers := make(http.Header)
	headers.Set("Content-Type", "application/vnd.docker.distribution.manifest.v2+json")
Michael Yang's avatar
Michael Yang committed
1222
	resp, err := makeRequestWithRetry(ctx, "PUT", requestURL, headers, bytes.NewReader(manifestJSON), regOpts)
1223
1224
1225
1226
1227
	if err != nil {
		return err
	}
	defer resp.Body.Close()

1228
	fn(api.ProgressResponse{Status: "success"})
1229
1230
1231
1232

	return nil
}

1233
func PullModel(ctx context.Context, name string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
1234
1235
	mp := ParseModelPath(name)

1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
1255
1256
	var manifest *ManifestV2
	var err error
	var noprune string

	// build deleteMap to prune unused layers
	deleteMap := make(map[string]bool)

	if noprune = os.Getenv("OLLAMA_NOPRUNE"); noprune == "" {
		manifest, _, err = GetManifest(mp)
		if err != nil && !errors.Is(err, os.ErrNotExist) {
			return err
		}

		if manifest != nil {
			for _, l := range manifest.Layers {
				deleteMap[l.Digest] = true
			}
			deleteMap[manifest.Config.Digest] = true
		}
	}

1257
1258
	if mp.ProtocolScheme == "http" && !regOpts.Insecure {
		return fmt.Errorf("insecure protocol http")
1259
	}
1260

1261
	fn(api.ProgressResponse{Status: "pulling manifest"})
1262

1263
	manifest, err = pullModelManifest(ctx, mp, regOpts)
1264
	if err != nil {
1265
		return fmt.Errorf("pull model manifest: %s", err)
1266
1267
1268
	}

	var layers []*Layer
Bruce MacDonald's avatar
Bruce MacDonald committed
1269
	layers = append(layers, manifest.Layers...)
1270
1271
1272
	layers = append(layers, &manifest.Config)

	for _, layer := range layers {
1273
1274
1275
1276
1277
1278
1279
1280
		if err := downloadBlob(
			ctx,
			downloadOpts{
				mp:      mp,
				digest:  layer.Digest,
				regOpts: regOpts,
				fn:      fn,
			}); err != nil {
1281
1282
			return err
		}
1283
		delete(deleteMap, layer.Digest)
1284
	}
1285
	delete(deleteMap, manifest.Config.Digest)
1286

Michael Yang's avatar
Michael Yang committed
1287
1288
1289
	fn(api.ProgressResponse{Status: "verifying sha256 digest"})
	for _, layer := range layers {
		if err := verifyBlob(layer.Digest); err != nil {
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
			if errors.Is(err, errDigestMismatch) {
				// something went wrong, delete the blob
				fp, err := GetBlobsPath(layer.Digest)
				if err != nil {
					return err
				}
				if err := os.Remove(fp); err != nil {
					// log this, but return the original error
					log.Printf("couldn't remove file with digest mismatch '%s': %v", fp, err)
				}
			}
Michael Yang's avatar
Michael Yang committed
1301
1302
1303
1304
			return err
		}
	}

1305
	fn(api.ProgressResponse{Status: "writing manifest"})
1306

1307
	manifestJSON, err := json.Marshal(manifest)
1308
1309
1310
1311
	if err != nil {
		return err
	}

Patrick Devine's avatar
Patrick Devine committed
1312
	fp, err := mp.GetManifestPath(true)
1313
1314
1315
1316
	if err != nil {
		return err
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
1317
	err = os.WriteFile(fp, manifestJSON, 0o644)
1318
1319
1320
1321
1322
	if err != nil {
		log.Printf("couldn't write to %s", fp)
		return err
	}

1323
1324
1325
1326
1327
1328
1329
1330
	if noprune == "" {
		fn(api.ProgressResponse{Status: "removing any unused layers"})
		err = deleteUnusedLayers(nil, deleteMap, false)
		if err != nil {
			return err
		}
	}

1331
	fn(api.ProgressResponse{Status: "success"})
1332
1333
1334
1335

	return nil
}

1336
func pullModelManifest(ctx context.Context, mp ModelPath, regOpts *RegistryOptions) (*ManifestV2, error) {
Michael Yang's avatar
Michael Yang committed
1337
	requestURL := mp.BaseURL().JoinPath("v2", mp.GetNamespaceRepository(), "manifests", mp.Tag)
1338

Michael Yang's avatar
Michael Yang committed
1339
1340
	headers := make(http.Header)
	headers.Set("Accept", "application/vnd.docker.distribution.manifest.v2+json")
Michael Yang's avatar
Michael Yang committed
1341
	resp, err := makeRequest(ctx, "GET", requestURL, headers, nil, regOpts)
1342
1343
1344
1345
1346
1347
	if err != nil {
		log.Printf("couldn't get manifest: %v", err)
		return nil, err
	}
	defer resp.Body.Close()

Michael Yang's avatar
Michael Yang committed
1348
	if resp.StatusCode >= http.StatusBadRequest {
1349
		if resp.StatusCode == http.StatusNotFound {
Bruce MacDonald's avatar
Bruce MacDonald committed
1350
			return nil, fmt.Errorf("model not found")
1351
		}
Michael Yang's avatar
Michael Yang committed
1352

1353
		body, _ := io.ReadAll(resp.Body)
1354
		return nil, fmt.Errorf("on pull registry responded with code %d: %s", resp.StatusCode, body)
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
	}

	var m *ManifestV2
	if err := json.NewDecoder(resp.Body).Decode(&m); err != nil {
		return nil, err
	}

	return m, err
}

1365
1366
1367
1368
func createConfigLayer(config ConfigV2, layers []string) (*LayerReader, error) {
	config.RootFS = RootFS{
		Type:    "layers",
		DiffIDs: layers,
1369
1370
1371
1372
1373
1374
1375
	}

	configJSON, err := json.Marshal(config)
	if err != nil {
		return nil, err
	}

1376
	digest, size := GetSHA256Digest(bytes.NewBuffer(configJSON))
1377

Michael Yang's avatar
Michael Yang committed
1378
	layer := &LayerReader{
1379
1380
1381
1382
1383
		Layer: Layer{
			MediaType: "application/vnd.docker.container.image.v1+json",
			Digest:    digest,
			Size:      size,
		},
1384
		Reader: bytes.NewBuffer(configJSON),
1385
1386
1387
1388
1389
	}
	return layer, nil
}

// GetSHA256Digest returns the SHA256 hash of a given buffer and returns it, and the size of buffer
Michael Yang's avatar
Michael Yang committed
1390
func GetSHA256Digest(r io.Reader) (string, int64) {
Michael Yang's avatar
Michael Yang committed
1391
1392
1393
1394
1395
1396
	h := sha256.New()
	n, err := io.Copy(h, r)
	if err != nil {
		log.Fatal(err)
	}

Michael Yang's avatar
Michael Yang committed
1397
	return fmt.Sprintf("sha256:%x", h.Sum(nil)), n
1398
1399
1400
}

// Function to check if a blob already exists in the Docker registry
1401
func checkBlobExistence(ctx context.Context, mp ModelPath, digest string, regOpts *RegistryOptions) (bool, error) {
Michael Yang's avatar
Michael Yang committed
1402
1403
	requestURL := mp.BaseURL()
	requestURL = requestURL.JoinPath("v2", mp.GetNamespaceRepository(), "blobs", digest)
1404

Michael Yang's avatar
Michael Yang committed
1405
	resp, err := makeRequest(ctx, "HEAD", requestURL, nil, nil, regOpts)
1406
1407
1408
1409
1410
1411
1412
	if err != nil {
		log.Printf("couldn't check for blob: %v", err)
		return false, err
	}
	defer resp.Body.Close()

	// Check for success: If the blob exists, the Docker registry will respond with a 200 OK
Michael Yang's avatar
Michael Yang committed
1413
	return resp.StatusCode < http.StatusBadRequest, nil
1414
1415
}

Michael Yang's avatar
Michael Yang committed
1416
func makeRequestWithRetry(ctx context.Context, method string, requestURL *url.URL, headers http.Header, body io.ReadSeeker, regOpts *RegistryOptions) (*http.Response, error) {
Michael Yang's avatar
Michael Yang committed
1417
	var status string
Michael Yang's avatar
Michael Yang committed
1418
	for try := 0; try < maxRetries; try++ {
Michael Yang's avatar
Michael Yang committed
1419
		resp, err := makeRequest(ctx, method, requestURL, headers, body, regOpts)
Michael Yang's avatar
Michael Yang committed
1420
1421
1422
1423
1424
1425
1426
		if err != nil {
			log.Printf("couldn't start upload: %v", err)
			return nil, err
		}

		status = resp.Status

Michael Yang's avatar
Michael Yang committed
1427
1428
		switch {
		case resp.StatusCode == http.StatusUnauthorized:
Michael Yang's avatar
Michael Yang committed
1429
1430
			auth := resp.Header.Get("www-authenticate")
			authRedir := ParseAuthRedirectString(auth)
Michael Yang's avatar
Michael Yang committed
1431
			token, err := getAuthToken(ctx, authRedir)
Michael Yang's avatar
Michael Yang committed
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
			if err != nil {
				return nil, err
			}

			regOpts.Token = token
			if body != nil {
				if _, err := body.Seek(0, io.SeekStart); err != nil {
					return nil, err
				}
			}

			continue
Michael Yang's avatar
Michael Yang committed
1444
		case resp.StatusCode >= http.StatusBadRequest:
Michael Yang's avatar
Michael Yang committed
1445
1446
			body, _ := io.ReadAll(resp.Body)
			return nil, fmt.Errorf("on upload registry responded with code %d: %s", resp.StatusCode, body)
Michael Yang's avatar
Michael Yang committed
1447
1448
		default:
			return resp, nil
Michael Yang's avatar
Michael Yang committed
1449
1450
1451
1452
1453
1454
		}
	}

	return nil, fmt.Errorf("max retry exceeded: %v", status)
}

Michael Yang's avatar
Michael Yang committed
1455
func makeRequest(ctx context.Context, method string, requestURL *url.URL, headers http.Header, body io.Reader, regOpts *RegistryOptions) (*http.Response, error) {
Michael Yang's avatar
Michael Yang committed
1456
	if requestURL.Scheme != "http" && regOpts != nil && regOpts.Insecure {
Michael Yang's avatar
Michael Yang committed
1457
		requestURL.Scheme = "http"
1458
1459
	}

Michael Yang's avatar
Michael Yang committed
1460
	req, err := http.NewRequestWithContext(ctx, method, requestURL.String(), body)
1461
1462
1463
1464
	if err != nil {
		return nil, err
	}

Michael Yang's avatar
Michael Yang committed
1465
1466
1467
1468
	if headers != nil {
		req.Header = headers
	}

Michael Yang's avatar
Michael Yang committed
1469
1470
1471
1472
1473
1474
	if regOpts != nil {
		if regOpts.Token != "" {
			req.Header.Set("Authorization", "Bearer "+regOpts.Token)
		} else if regOpts.Username != "" && regOpts.Password != "" {
			req.SetBasicAuth(regOpts.Username, regOpts.Password)
		}
1475
1476
	}

Michael Yang's avatar
Michael Yang committed
1477
	req.Header.Set("User-Agent", fmt.Sprintf("ollama/%s (%s %s) Go/%s", version.Version, runtime.GOARCH, runtime.GOOS, runtime.Version()))
1478

Michael Yang's avatar
Michael Yang committed
1479
1480
1481
1482
1483
1484
1485
1486
1487
	if s := req.Header.Get("Content-Length"); s != "" {
		contentLength, err := strconv.ParseInt(s, 10, 64)
		if err != nil {
			return nil, err
		}

		req.ContentLength = contentLength
	}

Michael Yang's avatar
Michael Yang committed
1488
	resp, err := http.DefaultClient.Do(req)
1489
1490
1491
1492
1493
1494
	if err != nil {
		return nil, err
	}

	return resp, nil
}
Michael Yang's avatar
Michael Yang committed
1495

Patrick Devine's avatar
Patrick Devine committed
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
func getValue(header, key string) string {
	startIdx := strings.Index(header, key+"=")
	if startIdx == -1 {
		return ""
	}

	// Move the index to the starting quote after the key.
	startIdx += len(key) + 2
	endIdx := startIdx

	for endIdx < len(header) {
		if header[endIdx] == '"' {
			if endIdx+1 < len(header) && header[endIdx+1] != ',' { // If the next character isn't a comma, continue
				endIdx++
				continue
			}
			break
		}
		endIdx++
	}
	return header[startIdx:endIdx]
}

func ParseAuthRedirectString(authStr string) AuthRedirect {
	authStr = strings.TrimPrefix(authStr, "Bearer ")

	return AuthRedirect{
		Realm:   getValue(authStr, "realm"),
		Service: getValue(authStr, "service"),
		Scope:   getValue(authStr, "scope"),
	}
}

1529
1530
var errDigestMismatch = fmt.Errorf("digest mismatch, file must be downloaded again")

Michael Yang's avatar
Michael Yang committed
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
func verifyBlob(digest string) error {
	fp, err := GetBlobsPath(digest)
	if err != nil {
		return err
	}

	f, err := os.Open(fp)
	if err != nil {
		return err
	}
	defer f.Close()

	fileDigest, _ := GetSHA256Digest(f)
	if digest != fileDigest {
1545
		return fmt.Errorf("%w: want %s, got %s", errDigestMismatch, digest, fileDigest)
Michael Yang's avatar
Michael Yang committed
1546
1547
1548
1549
	}

	return nil
}