images.go 26.5 KB
Newer Older
1
2
3
4
package server

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

Michael Yang's avatar
Michael Yang committed
22
23
	"golang.org/x/exp/slices"

24
	"github.com/jmorganca/ollama/api"
25
	"github.com/jmorganca/ollama/llm"
26
	"github.com/jmorganca/ollama/parser"
Michael Yang's avatar
Michael Yang committed
27
	"github.com/jmorganca/ollama/version"
28
29
)

30
31
32
33
type RegistryOptions struct {
	Insecure bool
	Username string
	Password string
Patrick Devine's avatar
Patrick Devine committed
34
	Token    string
35
36
}

37
type Model struct {
Michael Yang's avatar
Michael Yang committed
38
39
40
41
42
43
44
45
46
47
48
	Name           string `json:"name"`
	ShortName      string
	ModelPath      string
	OriginalModel  string
	AdapterPaths   []string
	ProjectorPaths []string
	Template       string
	System         string
	License        []string
	Digest         string
	Options        map[string]interface{}
49
50
}

Bruce MacDonald's avatar
Bruce MacDonald committed
51
52
53
54
55
56
type PromptVars struct {
	System   string
	Prompt   string
	Response string
	First    bool
}
57

Bruce MacDonald's avatar
Bruce MacDonald committed
58
59
60
func (m *Model) Prompt(p PromptVars) (string, error) {
	var prompt strings.Builder
	tmpl, err := template.New("").Parse(m.Template)
61
62
63
64
	if err != nil {
		return "", err
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
65
66
67
	if p.System == "" {
		// use the default system prompt for this model if one is not specified
		p.System = m.System
68
69
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
70
71
72
73
74
75
76
77
	var sb strings.Builder
	if err := tmpl.Execute(&sb, p); err != nil {
		return "", err
	}
	prompt.WriteString(sb.String())
	prompt.WriteString(p.Response)
	return prompt.String(), nil
}
78

Bruce MacDonald's avatar
Bruce MacDonald committed
79
80
81
82
83
func (m *Model) ChatPrompt(msgs []api.Message) (string, error) {
	// build the prompt from the list of messages
	var prompt strings.Builder
	currentVars := PromptVars{
		First: true,
Bruce MacDonald's avatar
Bruce MacDonald committed
84
85
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
	writePrompt := func() error {
		p, err := m.Prompt(currentVars)
		if err != nil {
			return err
		}
		prompt.WriteString(p)
		currentVars = PromptVars{}
		return nil
	}

	for _, msg := range msgs {
		switch msg.Role {
		case "system":
			if currentVars.Prompt != "" || currentVars.System != "" {
				if err := writePrompt(); err != nil {
					return "", err
				}
			}
			currentVars.System = msg.Content
		case "user":
			if currentVars.Prompt != "" || currentVars.System != "" {
				if err := writePrompt(); err != nil {
					return "", err
				}
			}
			currentVars.Prompt = msg.Content
		case "assistant":
			currentVars.Response = msg.Content
			if err := writePrompt(); err != nil {
				return "", err
			}
		default:
			return "", fmt.Errorf("invalid role: %s, role must be one of [system, user, assistant]", msg.Role)
		}
120
121
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
122
123
124
125
126
127
128
129
	// Append the last set of vars if they are non-empty
	if currentVars.Prompt != "" || currentVars.System != "" {
		if err := writePrompt(); err != nil {
			return "", err
		}
	}

	return prompt.String(), nil
130
131
}

132
133
134
type ManifestV2 struct {
	SchemaVersion int      `json:"schemaVersion"`
	MediaType     string   `json:"mediaType"`
Michael Yang's avatar
Michael Yang committed
135
	Config        *Layer   `json:"config"`
136
137
138
139
	Layers        []*Layer `json:"layers"`
}

type ConfigV2 struct {
Michael Yang's avatar
Michael Yang committed
140
141
142
143
144
145
	ModelFormat   string   `json:"model_format"`
	ModelFamily   string   `json:"model_family"`
	ModelFamilies []string `json:"model_families"`
	ModelType     string   `json:"model_type"`
	FileType      string   `json:"file_type"`
	RootFS        RootFS   `json:"rootfs"`
146
147

	// required by spec
148
149
150
151
	Architecture string `json:"architecture"`
	OS           string `json:"os"`
}

Michael Yang's avatar
Michael Yang committed
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
func (c *ConfigV2) SetModelFormat(format string) {
	if c.ModelFormat == "" {
		c.ModelFormat = format
	}
}

func (c *ConfigV2) SetModelFamily(families ...string) {
	for _, family := range families {
		if c.ModelFamily == "" {
			c.ModelFamily = family
		}

		if !slices.Contains(c.ModelFamilies, family) {
			c.ModelFamilies = append(c.ModelFamilies, family)
		}
	}
}

func (c *ConfigV2) SetModelType(modelType string) {
	if c.ModelType == "" {
		c.ModelType = modelType
	}
}

func (c *ConfigV2) SetFileType(fileType string) {
	if c.FileType == "" {
		c.FileType = fileType
	}
}

182
183
184
185
186
type RootFS struct {
	Type    string   `json:"type"`
	DiffIDs []string `json:"diff_ids"`
}

Michael Yang's avatar
Michael Yang committed
187
func (m *ManifestV2) GetTotalSize() (total int64) {
Patrick Devine's avatar
Patrick Devine committed
188
189
190
	for _, layer := range m.Layers {
		total += layer.Size
	}
Michael Yang's avatar
Michael Yang committed
191

Patrick Devine's avatar
Patrick Devine committed
192
193
194
195
	total += m.Config.Size
	return total
}

Patrick Devine's avatar
Patrick Devine committed
196
func GetManifest(mp ModelPath) (*ManifestV2, string, error) {
197
	fp, err := mp.GetManifestPath()
198
	if err != nil {
Patrick Devine's avatar
Patrick Devine committed
199
		return nil, "", err
200
	}
201

202
	if _, err = os.Stat(fp); err != nil {
Patrick Devine's avatar
Patrick Devine committed
203
		return nil, "", err
204
205
206
207
	}

	var manifest *ManifestV2

208
	bts, err := os.ReadFile(fp)
209
	if err != nil {
Patrick Devine's avatar
Patrick Devine committed
210
		return nil, "", fmt.Errorf("couldn't open file '%s'", fp)
211
212
	}

Patrick Devine's avatar
Patrick Devine committed
213
214
215
	shaSum := sha256.Sum256(bts)
	shaStr := hex.EncodeToString(shaSum[:])

216
	if err := json.Unmarshal(bts, &manifest); err != nil {
Patrick Devine's avatar
Patrick Devine committed
217
		return nil, "", err
218
219
	}

Patrick Devine's avatar
Patrick Devine committed
220
	return manifest, shaStr, nil
221
222
223
}

func GetModel(name string) (*Model, error) {
224
	mp := ParseModelPath(name)
Patrick Devine's avatar
Patrick Devine committed
225
	manifest, digest, err := GetManifest(mp)
226
227
228
229
230
	if err != nil {
		return nil, err
	}

	model := &Model{
231
232
233
234
235
		Name:      mp.GetFullTagname(),
		ShortName: mp.GetShortTagname(),
		Digest:    digest,
		Template:  "{{ .Prompt }}",
		License:   []string{},
236
237
238
	}

	for _, layer := range manifest.Layers {
Patrick Devine's avatar
Patrick Devine committed
239
		filename, err := GetBlobsPath(layer.Digest)
240
241
242
243
		if err != nil {
			return nil, err
		}

244
245
246
		switch layer.MediaType {
		case "application/vnd.ollama.image.model":
			model.ModelPath = filename
Patrick Devine's avatar
Patrick Devine committed
247
			model.OriginalModel = layer.From
248
		case "application/vnd.ollama.image.embed":
249
250
251
			// Deprecated in versions  > 0.1.2
			// TODO: remove this warning in a future version
			log.Print("WARNING: model contains embeddings, but embeddings in modelfiles have been deprecated and will be ignored.")
252
253
		case "application/vnd.ollama.image.adapter":
			model.AdapterPaths = append(model.AdapterPaths, filename)
Michael Yang's avatar
Michael Yang committed
254
255
		case "application/vnd.ollama.image.projector":
			model.ProjectorPaths = append(model.ProjectorPaths, filename)
256
257
258
259
260
261
262
263
264
		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)
265
266
267
			if err != nil {
				return nil, err
			}
268
269

			model.System = string(bts)
270
271
272
273
274
275
276
		case "application/vnd.ollama.image.prompt":
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}

			model.Template = string(bts)
277
		case "application/vnd.ollama.image.params":
Michael Yang's avatar
Michael Yang committed
278
279
280
281
282
			params, err := os.Open(filename)
			if err != nil {
				return nil, err
			}
			defer params.Close()
283

284
			// parse model options parameters into a map so that we can see which fields have been specified explicitly
285
			if err = json.NewDecoder(params).Decode(&model.Options); err != nil {
286
287
				return nil, err
			}
Patrick Devine's avatar
Patrick Devine committed
288
289
290
291
292
293
		case "application/vnd.ollama.image.license":
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}
			model.License = append(model.License, string(bts))
294
295
296
297
298
299
		}
	}

	return model, nil
}

300
301
func realpath(mfDir, from string) string {
	abspath, err := filepath.Abs(from)
Michael Yang's avatar
Michael Yang committed
302
	if err != nil {
303
		return from
304
305
	}

Michael Yang's avatar
Michael Yang committed
306
	home, err := os.UserHomeDir()
307
	if err != nil {
Michael Yang's avatar
Michael Yang committed
308
		return abspath
309
310
	}

311
	if from == "~" {
Michael Yang's avatar
Michael Yang committed
312
		return home
313
314
315
316
317
318
319
	} else if strings.HasPrefix(from, "~/") {
		return filepath.Join(home, from[2:])
	}

	if _, err := os.Stat(filepath.Join(mfDir, from)); err == nil {
		// this is a file relative to the Modelfile
		return filepath.Join(mfDir, from)
320
321
	}

Michael Yang's avatar
Michael Yang committed
322
323
324
	return abspath
}

325
func CreateModel(ctx context.Context, name, modelFileDir string, commands []parser.Command, fn func(resp api.ProgressResponse)) error {
326
327
	config := ConfigV2{
		OS:           "linux",
Michael Yang's avatar
Michael Yang committed
328
		Architecture: "amd64",
Michael Yang's avatar
Michael Yang committed
329
330
331
		RootFS: RootFS{
			Type: "layers",
		},
332
333
	}

Michael Yang's avatar
Michael Yang committed
334
335
	deleteMap := make(map[string]struct{})

Michael Yang's avatar
Michael Yang committed
336
	var layers Layers
Michael Yang's avatar
Michael Yang committed
337

338
	params := make(map[string][]string)
Michael Yang's avatar
Michael Yang committed
339
340
	fromParams := make(map[string]any)

341
	for _, c := range commands {
Michael Yang's avatar
Michael Yang committed
342
343
344
		log.Printf("[%s] - %s", c.Name, c.Args)
		mediatype := fmt.Sprintf("application/vnd.ollama.image.%s", c.Name)

345
346
		switch c.Name {
		case "model":
Michael Yang's avatar
Michael Yang committed
347
348
349
350
351
352
353
354
355
			if strings.HasPrefix(c.Args, "@") {
				blobPath, err := GetBlobsPath(strings.TrimPrefix(c.Args, "@"))
				if err != nil {
					return err
				}

				c.Args = blobPath
			}

356
			bin, err := os.Open(realpath(modelFileDir, c.Args))
357
			if err != nil {
Michael Yang's avatar
Michael Yang committed
358
359
360
361
362
363
364
				// not a file on disk so must be a model reference
				modelpath := ParseModelPath(c.Args)
				manifest, _, err := GetManifest(modelpath)
				switch {
				case errors.Is(err, os.ErrNotExist):
					fn(api.ProgressResponse{Status: "pulling model"})
					if err := PullModel(ctx, c.Args, &RegistryOptions{}, fn); err != nil {
365
366
367
						return err
					}

Michael Yang's avatar
Michael Yang committed
368
					manifest, _, err = GetManifest(modelpath)
369
370
371
					if err != nil {
						return err
					}
Michael Yang's avatar
Michael Yang committed
372
373
				case err != nil:
					return err
374
				}
375

376
				fn(api.ProgressResponse{Status: "reading model metadata"})
Michael Yang's avatar
Michael Yang committed
377
				fromConfigPath, err := GetBlobsPath(manifest.Config.Digest)
Michael Yang's avatar
Michael Yang committed
378
379
380
381
				if err != nil {
					return err
				}

Michael Yang's avatar
Michael Yang committed
382
				fromConfigFile, err := os.Open(fromConfigPath)
Michael Yang's avatar
Michael Yang committed
383
384
385
				if err != nil {
					return err
				}
Michael Yang's avatar
Michael Yang committed
386
				defer fromConfigFile.Close()
Michael Yang's avatar
Michael Yang committed
387

Michael Yang's avatar
Michael Yang committed
388
389
				var fromConfig ConfigV2
				if err := json.NewDecoder(fromConfigFile).Decode(&fromConfig); err != nil {
Michael Yang's avatar
Michael Yang committed
390
391
392
					return err
				}

Michael Yang's avatar
Michael Yang committed
393
394
395
396
				config.SetModelFormat(fromConfig.ModelFormat)
				config.SetModelFamily(append(fromConfig.ModelFamilies, fromConfig.ModelFamily)...)
				config.SetModelType(fromConfig.ModelType)
				config.SetFileType(fromConfig.FileType)
Michael Yang's avatar
Michael Yang committed
397

Michael Yang's avatar
Michael Yang committed
398
399
400
401
				for _, layer := range manifest.Layers {
					deleteMap[layer.Digest] = struct{}{}
					if layer.MediaType == "application/vnd.ollama.image.params" {
						fromParamsPath, err := GetBlobsPath(layer.Digest)
Michael Yang's avatar
Michael Yang committed
402
403
404
405
						if err != nil {
							return err
						}

Michael Yang's avatar
Michael Yang committed
406
						fromParamsFile, err := os.Open(fromParamsPath)
Michael Yang's avatar
Michael Yang committed
407
408
409
						if err != nil {
							return err
						}
Michael Yang's avatar
Michael Yang committed
410
						defer fromParamsFile.Close()
Michael Yang's avatar
Michael Yang committed
411

Michael Yang's avatar
Michael Yang committed
412
						if err := json.NewDecoder(fromParamsFile).Decode(&fromParams); err != nil {
Michael Yang's avatar
Michael Yang committed
413
414
415
416
							return err
						}
					}

Michael Yang's avatar
Michael Yang committed
417
					layer, err := NewLayerFromLayer(layer.Digest, layer.MediaType, modelpath.GetShortTagname())
418
419
420
					if err != nil {
						return err
					}
Michael Yang's avatar
Michael Yang committed
421

Michael Yang's avatar
Michael Yang committed
422
					layers.Add(layer)
423
				}
Michael Yang's avatar
Michael Yang committed
424
425
426

				deleteMap[manifest.Config.Digest] = struct{}{}
				continue
427
			}
Michael Yang's avatar
Michael Yang committed
428
			defer bin.Close()
429

430
431
432
			var offset int64
			for {
				fn(api.ProgressResponse{Status: "creating model layer"})
433

434
435
436
437
438
439
440
				bin.Seek(offset, io.SeekStart)
				ggml, err := llm.DecodeGGML(bin)
				if errors.Is(err, io.EOF) {
					break
				} else if err != nil {
					return err
				}
Michael Yang's avatar
Michael Yang committed
441

Michael Yang's avatar
Michael Yang committed
442
443
444
445
				config.SetModelFormat(ggml.Name())
				config.SetModelFamily(ggml.ModelFamily())
				config.SetModelType(ggml.ModelType())
				config.SetFileType(ggml.FileType())
446

447
448
449
450
				mediatype := mediatype
				if ggml.ModelFamily() == "clip" {
					mediatype = "application/vnd.ollama.image.projector"
				}
451

452
453
454
455
456
457
458
459
460
461
				sr := io.NewSectionReader(bin, offset, ggml.Size)
				layer, err := NewLayer(sr, mediatype)
				if err != nil {
					return err
				}

				layers.Add(layer)

				offset += ggml.Size
			}
Michael Yang's avatar
Michael Yang committed
462
		case "adapter":
463
464
465
466
467
468
469
470
			if strings.HasPrefix(c.Args, "@") {
				blobPath, err := GetBlobsPath(strings.TrimPrefix(c.Args, "@"))
				if err != nil {
					return err
				}

				c.Args = blobPath
			}
Bruce MacDonald's avatar
Bruce MacDonald committed
471

Michael Yang's avatar
Michael Yang committed
472
			fn(api.ProgressResponse{Status: "creating adapter layer"})
473
			bin, err := os.Open(realpath(modelFileDir, c.Args))
474
			if err != nil {
Michael Yang's avatar
Michael Yang committed
475
				return err
476
			}
Michael Yang's avatar
Michael Yang committed
477
			defer bin.Close()
478

Michael Yang's avatar
Michael Yang committed
479
			layer, err := NewLayer(bin, mediatype)
480
			if err != nil {
Michael Yang's avatar
Michael Yang committed
481
				return err
482
			}
Bruce MacDonald's avatar
Bruce MacDonald committed
483

Michael Yang's avatar
Michael Yang committed
484
			layers.Add(layer)
Michael Yang's avatar
Michael Yang committed
485
486
		case "license":
			fn(api.ProgressResponse{Status: "creating license layer"})
Michael Yang's avatar
Michael Yang committed
487
488
489

			bin := strings.NewReader(c.Args)
			layer, err := NewLayer(bin, mediatype)
Bruce MacDonald's avatar
Bruce MacDonald committed
490
491
492
493
			if err != nil {
				return err
			}

Michael Yang's avatar
Michael Yang committed
494
			layers.Add(layer)
Michael Yang's avatar
Michael Yang committed
495
496
497
		case "template", "system":
			fn(api.ProgressResponse{Status: fmt.Sprintf("creating %s layer", c.Name)})

Michael Yang's avatar
Michael Yang committed
498
499
			bin := strings.NewReader(c.Args)
			layer, err := NewLayer(bin, mediatype)
500
			if err != nil {
501
				return err
502
			}
503

Michael Yang's avatar
Michael Yang committed
504
			layers.Replace(layer)
505
		default:
506
			params[c.Name] = append(params[c.Name], c.Args)
507
508
509
		}
	}

Michael Yang's avatar
Michael Yang committed
510
	if len(params) > 0 {
Michael Yang's avatar
Michael Yang committed
511
		fn(api.ProgressResponse{Status: "creating parameters layer"})
Michael Yang's avatar
Michael Yang committed
512

513
		formattedParams, err := api.FormatParams(params)
514
		if err != nil {
Michael Yang's avatar
Michael Yang committed
515
			return err
516
		}
517

Michael Yang's avatar
Michael Yang committed
518
		for k, v := range fromParams {
Michael Yang's avatar
Michael Yang committed
519
520
521
522
523
			if _, ok := formattedParams[k]; !ok {
				formattedParams[k] = v
			}
		}

Michael Yang's avatar
Michael Yang committed
524
		if config.ModelType == "65B" {
Michael Yang's avatar
Michael Yang committed
525
			if gqa, ok := formattedParams["gqa"].(int); ok && gqa == 8 {
Michael Yang's avatar
Michael Yang committed
526
527
528
529
				config.ModelType = "70B"
			}
		}

Michael Yang's avatar
Michael Yang committed
530
531
		var b bytes.Buffer
		if err := json.NewEncoder(&b).Encode(formattedParams); err != nil {
532
533
534
			return err
		}

Michael Yang's avatar
Michael Yang committed
535
		fn(api.ProgressResponse{Status: "creating config layer"})
Michael Yang's avatar
Michael Yang committed
536
		layer, err := NewLayer(&b, "application/vnd.ollama.image.params")
537
		if err != nil {
Michael Yang's avatar
Michael Yang committed
538
			return err
539
		}
Michael Yang's avatar
Michael Yang committed
540

Michael Yang's avatar
Michael Yang committed
541
		layers.Replace(layer)
542
543
	}

Michael Yang's avatar
Michael Yang committed
544
545
546
	digests := make([]string, len(layers.items))
	for i, layer := range layers.items {
		digests[i] = layer.Digest
547
548
	}

Michael Yang's avatar
Michael Yang committed
549
	config.RootFS.DiffIDs = digests
Michael Yang's avatar
Michael Yang committed
550

Michael Yang's avatar
Michael Yang committed
551
552
	var b bytes.Buffer
	if err := json.NewEncoder(&b).Encode(config); err != nil {
553
554
555
		return err
	}

Michael Yang's avatar
Michael Yang committed
556
557
	configLayer, err := NewLayer(&b, "application/vnd.docker.container.image.v1+json")
	if err != nil {
558
559
560
		return err
	}

Michael Yang's avatar
Michael Yang committed
561
	delete(deleteMap, configLayer.Digest)
562

Michael Yang's avatar
Michael Yang committed
563
564
	for _, layer := range append(layers.items, configLayer) {
		committed, err := layer.Commit()
565
566
567
		if err != nil {
			return err
		}
568

Michael Yang's avatar
Michael Yang committed
569
570
571
		status := "writing layer"
		if !committed {
			status = "using already created layer"
572
573
		}

Michael Yang's avatar
Michael Yang committed
574
		fn(api.ProgressResponse{Status: fmt.Sprintf("%s %s", status, layer.Digest)})
575

Michael Yang's avatar
Michael Yang committed
576
		delete(deleteMap, layer.Digest)
577
578
	}

Michael Yang's avatar
Michael Yang committed
579
580
	fn(api.ProgressResponse{Status: "writing manifest"})
	if err := WriteManifest(name, configLayer, layers.items); err != nil {
581
582
		return err
	}
583

Michael Yang's avatar
Michael Yang committed
584
585
586
	if noprune := os.Getenv("OLLAMA_NOPRUNE"); noprune == "" {
		if err := deleteUnusedLayers(nil, deleteMap, false); err != nil {
			return err
587
588
589
		}
	}

Michael Yang's avatar
Michael Yang committed
590
591
	fn(api.ProgressResponse{Status: "success"})
	return nil
592
593
}

Patrick Devine's avatar
Patrick Devine committed
594
func CopyModel(src, dest string) error {
595
	srcModelPath := ParseModelPath(src)
596
	srcPath, err := srcModelPath.GetManifestPath()
597
598
599
600
	if err != nil {
		return err
	}

601
	destModelPath := ParseModelPath(dest)
602
	destPath, err := destModelPath.GetManifestPath()
Patrick Devine's avatar
Patrick Devine committed
603
604
605
	if err != nil {
		return err
	}
606
607
608
	if err := os.MkdirAll(filepath.Dir(destPath), 0o755); err != nil {
		return err
	}
Patrick Devine's avatar
Patrick Devine committed
609
610

	// copy the file
Michael Yang's avatar
Michael Yang committed
611
	input, err := os.ReadFile(srcPath)
Patrick Devine's avatar
Patrick Devine committed
612
613
614
615
616
	if err != nil {
		fmt.Println("Error reading file:", err)
		return err
	}

Michael Yang's avatar
Michael Yang committed
617
	err = os.WriteFile(destPath, input, 0o644)
Patrick Devine's avatar
Patrick Devine committed
618
619
620
621
622
623
624
625
	if err != nil {
		fmt.Println("Error reading file:", err)
		return err
	}

	return nil
}

Michael Yang's avatar
Michael Yang committed
626
func deleteUnusedLayers(skipModelPath *ModelPath, deleteMap map[string]struct{}, dryRun bool) error {
627
628
629
630
	fp, err := GetManifestPath()
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
631
632
633
634

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

Michael Yang's avatar
Michael Yang committed
637
638
639
640
		dir, file := filepath.Split(path)
		dir = strings.Trim(strings.TrimPrefix(dir, fp), string(os.PathSeparator))
		tag := strings.Join([]string{dir, file}, ":")
		fmp := ParseModelPath(tag)
641

Michael Yang's avatar
Michael Yang committed
642
		// skip the manifest we're trying to delete
643
		if skipModelPath != nil && skipModelPath.GetFullTagname() == fmp.GetFullTagname() {
Michael Yang's avatar
Michael Yang committed
644
			return nil
645
		}
Michael Yang's avatar
Michael Yang committed
646
647
648
649
650
651
652
653
654
655
656
657

		// 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)
658
		return nil
Michael Yang's avatar
Michael Yang committed
659
660
661
	}

	if err := filepath.Walk(fp, walkFunc); err != nil {
Michael Yang's avatar
Michael Yang committed
662
663
		return err
	}
664
665

	// only delete the files which are still in the deleteMap
Michael Yang's avatar
Michael Yang committed
666
667
668
669
670
671
672
673
674
	for k := range deleteMap {
		fp, err := GetBlobsPath(k)
		if err != nil {
			log.Printf("couldn't get file path for '%s': %v", k, err)
			continue
		}
		if !dryRun {
			if err := os.Remove(fp); err != nil {
				log.Printf("couldn't remove file '%s': %v", fp, err)
675
676
				continue
			}
Michael Yang's avatar
Michael Yang committed
677
678
		} else {
			log.Printf("wanted to remove: %s", fp)
679
680
681
		}
	}

682
683
684
685
	return nil
}

func PruneLayers() error {
Michael Yang's avatar
Michael Yang committed
686
	deleteMap := make(map[string]struct{})
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
	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, "-", ":")
		}
Michael Yang's avatar
Michael Yang committed
703
704
705
		if strings.HasPrefix(name, "sha256:") {
			deleteMap[name] = struct{}{}
		}
706
707
708
709
710
711
712
713
714
715
716
717
718
719
	}

	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
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
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
}

753
754
755
756
757
758
759
func DeleteModel(name string) error {
	mp := ParseModelPath(name)
	manifest, _, err := GetManifest(mp)
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
760
	deleteMap := make(map[string]struct{})
761
	for _, layer := range manifest.Layers {
Michael Yang's avatar
Michael Yang committed
762
		deleteMap[layer.Digest] = struct{}{}
763
	}
Michael Yang's avatar
Michael Yang committed
764
	deleteMap[manifest.Config.Digest] = struct{}{}
765
766
767
768
769
770

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

771
	fp, err := mp.GetManifestPath()
772
773
774
775
776
777
778
779
780
781
782
783
	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
784
func ShowModelfile(model *Model) (string, error) {
Michael Yang's avatar
Michael Yang committed
785
	var mt struct {
Patrick Devine's avatar
Patrick Devine committed
786
		*Model
Michael Yang's avatar
Michael Yang committed
787
		From       string
Michael Yang's avatar
Michael Yang committed
788
		Parameters map[string][]any
Patrick Devine's avatar
Patrick Devine committed
789
790
	}

Michael Yang's avatar
Michael Yang committed
791
	mt.Parameters = make(map[string][]any)
Patrick Devine's avatar
Patrick Devine committed
792
	for k, v := range model.Options {
Michael Yang's avatar
Michael Yang committed
793
794
795
		if s, ok := v.([]any); ok {
			mt.Parameters[k] = s
			continue
Patrick Devine's avatar
Patrick Devine committed
796
797
		}

Michael Yang's avatar
Michael Yang committed
798
		mt.Parameters[k] = []any{v}
Patrick Devine's avatar
Patrick Devine committed
799
800
	}

Michael Yang's avatar
Michael Yang committed
801
802
	mt.Model = model
	mt.From = model.ModelPath
Patrick Devine's avatar
Patrick Devine committed
803

Michael Yang's avatar
Michael Yang committed
804
	if model.OriginalModel != "" {
Daniel Reis's avatar
Daniel Reis committed
805
		mt.From = model.OriginalModel
Patrick Devine's avatar
Patrick Devine committed
806
807
808
809
810
811
812
813
	}

	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 }}"""
814
815

{{- if .System }}
Patrick Devine's avatar
Patrick Devine committed
816
SYSTEM """{{ .System }}"""
817
{{- end }}
818
819
820
821

{{- range $adapter := .AdapterPaths }}
ADAPTER {{ $adapter }}
{{- end }}
Michael Yang's avatar
Michael Yang committed
822

Michael Yang's avatar
Michael Yang committed
823
824
825
826
{{- range $k, $v := .Parameters }}
{{- range $parameter := $v }}
PARAMETER {{ $k }} {{ printf "%#v" $parameter }}
{{- end }}
Michael Yang's avatar
Michael Yang committed
827
{{- end }}`
Patrick Devine's avatar
Patrick Devine committed
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844

	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
}

845
func PushModel(ctx context.Context, name string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
846
	mp := ParseModelPath(name)
847
848
	fn(api.ProgressResponse{Status: "retrieving manifest"})

849
850
851
852
	if mp.ProtocolScheme == "http" && !regOpts.Insecure {
		return fmt.Errorf("insecure protocol http")
	}

Patrick Devine's avatar
Patrick Devine committed
853
	manifest, _, err := GetManifest(mp)
854
	if err != nil {
855
		fn(api.ProgressResponse{Status: "couldn't retrieve manifest"})
856
857
858
859
		return err
	}

	var layers []*Layer
Jeffrey Morgan's avatar
Jeffrey Morgan committed
860
	layers = append(layers, manifest.Layers...)
Michael Yang's avatar
Michael Yang committed
861
	layers = append(layers, manifest.Config)
862
863

	for _, layer := range layers {
Michael Yang's avatar
Michael Yang committed
864
		if err := uploadBlob(ctx, mp, layer, regOpts, fn); err != nil {
865
			log.Printf("error uploading blob: %v", err)
866
867
868
			if errors.Is(err, errUnauthorized) {
				return fmt.Errorf("unable to push %s, make sure this namespace exists and you are authorized to push to it", ParseModelPath(name).GetNamespaceRepository())
			}
869
870
			return err
		}
871
872
	}

873
	fn(api.ProgressResponse{Status: "pushing manifest"})
Michael Yang's avatar
Michael Yang committed
874
875
	requestURL := mp.BaseURL()
	requestURL = requestURL.JoinPath("v2", mp.GetNamespaceRepository(), "manifests", mp.Tag)
876
877
878
879
880
881

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

Michael Yang's avatar
Michael Yang committed
882
883
	headers := make(http.Header)
	headers.Set("Content-Type", "application/vnd.docker.distribution.manifest.v2+json")
Michael Yang's avatar
Michael Yang committed
884
	resp, err := makeRequestWithRetry(ctx, http.MethodPut, requestURL, headers, bytes.NewReader(manifestJSON), regOpts)
885
886
887
888
889
	if err != nil {
		return err
	}
	defer resp.Body.Close()

890
	fn(api.ProgressResponse{Status: "success"})
891
892
893
894

	return nil
}

895
func PullModel(ctx context.Context, name string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
896
897
	mp := ParseModelPath(name)

898
899
900
901
902
	var manifest *ManifestV2
	var err error
	var noprune string

	// build deleteMap to prune unused layers
Michael Yang's avatar
Michael Yang committed
903
	deleteMap := make(map[string]struct{})
904
905
906
907
908
909
910
911
912

	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 {
Michael Yang's avatar
Michael Yang committed
913
				deleteMap[l.Digest] = struct{}{}
914
			}
Michael Yang's avatar
Michael Yang committed
915
			deleteMap[manifest.Config.Digest] = struct{}{}
916
917
918
		}
	}

919
920
	if mp.ProtocolScheme == "http" && !regOpts.Insecure {
		return fmt.Errorf("insecure protocol http")
921
	}
922

923
	fn(api.ProgressResponse{Status: "pulling manifest"})
924

925
	manifest, err = pullModelManifest(ctx, mp, regOpts)
926
	if err != nil {
927
		return fmt.Errorf("pull model manifest: %s", err)
928
929
930
	}

	var layers []*Layer
Bruce MacDonald's avatar
Bruce MacDonald committed
931
	layers = append(layers, manifest.Layers...)
Michael Yang's avatar
Michael Yang committed
932
	layers = append(layers, manifest.Config)
933
934

	for _, layer := range layers {
935
936
937
938
939
940
941
942
		if err := downloadBlob(
			ctx,
			downloadOpts{
				mp:      mp,
				digest:  layer.Digest,
				regOpts: regOpts,
				fn:      fn,
			}); err != nil {
943
944
			return err
		}
945
		delete(deleteMap, layer.Digest)
946
	}
947
	delete(deleteMap, manifest.Config.Digest)
948

Michael Yang's avatar
Michael Yang committed
949
950
951
	fn(api.ProgressResponse{Status: "verifying sha256 digest"})
	for _, layer := range layers {
		if err := verifyBlob(layer.Digest); err != nil {
952
953
954
955
956
957
958
959
960
961
962
			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
963
964
965
966
			return err
		}
	}

967
	fn(api.ProgressResponse{Status: "writing manifest"})
968

969
	manifestJSON, err := json.Marshal(manifest)
970
971
972
973
	if err != nil {
		return err
	}

974
	fp, err := mp.GetManifestPath()
975
976
977
	if err != nil {
		return err
	}
978
979
980
	if err := os.MkdirAll(filepath.Dir(fp), 0o755); err != nil {
		return err
	}
981

Bruce MacDonald's avatar
Bruce MacDonald committed
982
	err = os.WriteFile(fp, manifestJSON, 0o644)
983
984
985
986
987
	if err != nil {
		log.Printf("couldn't write to %s", fp)
		return err
	}

988
989
990
991
992
993
994
995
	if noprune == "" {
		fn(api.ProgressResponse{Status: "removing any unused layers"})
		err = deleteUnusedLayers(nil, deleteMap, false)
		if err != nil {
			return err
		}
	}

996
	fn(api.ProgressResponse{Status: "success"})
997
998
999
1000

	return nil
}

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

Michael Yang's avatar
Michael Yang committed
1004
1005
	headers := make(http.Header)
	headers.Set("Accept", "application/vnd.docker.distribution.manifest.v2+json")
Michael Yang's avatar
Michael Yang committed
1006
	resp, err := makeRequestWithRetry(ctx, http.MethodGet, requestURL, headers, nil, regOpts)
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

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

	return m, err
}

// GetSHA256Digest returns the SHA256 hash of a given buffer and returns it, and the size of buffer
Michael Yang's avatar
Michael Yang committed
1021
func GetSHA256Digest(r io.Reader) (string, int64) {
Michael Yang's avatar
Michael Yang committed
1022
1023
1024
1025
1026
1027
	h := sha256.New()
	n, err := io.Copy(h, r)
	if err != nil {
		log.Fatal(err)
	}

Michael Yang's avatar
Michael Yang committed
1028
	return fmt.Sprintf("sha256:%x", h.Sum(nil)), n
1029
1030
}

1031
1032
var errUnauthorized = fmt.Errorf("unauthorized")

Michael Yang's avatar
Michael Yang committed
1033
func makeRequestWithRetry(ctx context.Context, method string, requestURL *url.URL, headers http.Header, body io.ReadSeeker, regOpts *RegistryOptions) (*http.Response, error) {
1034
1035
1036
1037
1038
	resp, err := makeRequest(ctx, method, requestURL, headers, body, regOpts)
	if err != nil {
		if !errors.Is(err, context.Canceled) {
			log.Printf("request failed: %v", err)
		}
1039

1040
1041
1042
1043
1044
1045
1046
1047
1048
		return nil, err
	}

	switch {
	case resp.StatusCode == http.StatusUnauthorized:
		// Handle authentication error with one retry
		auth := resp.Header.Get("www-authenticate")
		authRedir := ParseAuthRedirectString(auth)
		token, err := getAuthToken(ctx, authRedir)
Michael Yang's avatar
Michael Yang committed
1049
1050
1051
		if err != nil {
			return nil, err
		}
1052
1053
1054
		regOpts.Token = token
		if body != nil {
			_, err = body.Seek(0, io.SeekStart)
Michael Yang's avatar
Michael Yang committed
1055
1056
1057
1058
			if err != nil {
				return nil, err
			}
		}
1059
1060
1061
1062
1063
1064
1065

		resp, err := makeRequest(ctx, method, requestURL, headers, body, regOpts)
		if resp.StatusCode == http.StatusUnauthorized {
			return nil, errUnauthorized
		}

		return resp, err
1066
1067
1068
1069
1070
1071
1072
1073
	case resp.StatusCode == http.StatusNotFound:
		return nil, os.ErrNotExist
	case resp.StatusCode >= http.StatusBadRequest:
		responseBody, err := io.ReadAll(resp.Body)
		if err != nil {
			return nil, fmt.Errorf("%d: %s", resp.StatusCode, err)
		}
		return nil, fmt.Errorf("%d: %s", resp.StatusCode, responseBody)
Michael Yang's avatar
Michael Yang committed
1074
1075
	}

1076
	return resp, nil
Michael Yang's avatar
Michael Yang committed
1077
1078
}

Michael Yang's avatar
Michael Yang committed
1079
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
1080
	if requestURL.Scheme != "http" && regOpts != nil && regOpts.Insecure {
Michael Yang's avatar
Michael Yang committed
1081
		requestURL.Scheme = "http"
1082
1083
	}

Michael Yang's avatar
Michael Yang committed
1084
	req, err := http.NewRequestWithContext(ctx, method, requestURL.String(), body)
1085
1086
1087
1088
	if err != nil {
		return nil, err
	}

Michael Yang's avatar
Michael Yang committed
1089
1090
1091
1092
	if headers != nil {
		req.Header = headers
	}

Michael Yang's avatar
Michael Yang committed
1093
1094
1095
1096
1097
1098
	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)
		}
1099
1100
	}

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

Michael Yang's avatar
Michael Yang committed
1103
1104
1105
1106
1107
1108
1109
1110
1111
	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
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
	proxyURL, err := http.ProxyFromEnvironment(req)
	if err != nil {
		return nil, err
	}

	client := http.Client{
		Transport: &http.Transport{
			Proxy: http.ProxyURL(proxyURL),
		},
	}

	resp, err := client.Do(req)
1124
1125
1126
1127
1128
1129
	if err != nil {
		return nil, err
	}

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

Patrick Devine's avatar
Patrick Devine committed
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
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"),
	}
}

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

Michael Yang's avatar
Michael Yang committed
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
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 {
1180
		return fmt.Errorf("%w: want %s, got %s", errDigestMismatch, digest, fileDigest)
Michael Yang's avatar
Michael Yang committed
1181
1182
1183
1184
	}

	return nil
}