images.go 27 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
	Name           string `json:"name"`
39
	Config         ConfigV2
Michael Yang's avatar
Michael Yang committed
40
41
42
43
44
45
46
47
48
49
	ShortName      string
	ModelPath      string
	OriginalModel  string
	AdapterPaths   []string
	ProjectorPaths []string
	Template       string
	System         string
	License        []string
	Digest         string
	Options        map[string]interface{}
50
51
}

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

Bruce MacDonald's avatar
Bruce MacDonald committed
59
60
func (m *Model) Prompt(p PromptVars) (string, error) {
	var prompt strings.Builder
61
62
	// Use the "missingkey=zero" option to handle missing variables without panicking
	tmpl, err := template.New("").Option("missingkey=zero").Parse(m.Template)
63
64
65
66
	if err != nil {
		return "", err
	}

67
68
69
70
71
	if p.System == "" {
		// use the default system prompt for this model if one is not specified
		p.System = m.System
	}

72
73
74
75
76
	vars := map[string]any{
		"System":   p.System,
		"Prompt":   p.Prompt,
		"Response": p.Response,
		"First":    p.First,
77
78
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
79
	var sb strings.Builder
80
	if err := tmpl.Execute(&sb, vars); err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
81
82
83
84
85
86
		return "", err
	}
	prompt.WriteString(sb.String())
	prompt.WriteString(p.Response)
	return prompt.String(), nil
}
87

Bruce MacDonald's avatar
Bruce MacDonald committed
88
89
90
91
92
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
93
94
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
95
96
97
98
99
100
101
102
103
104
105
	writePrompt := func() error {
		p, err := m.Prompt(currentVars)
		if err != nil {
			return err
		}
		prompt.WriteString(p)
		currentVars = PromptVars{}
		return nil
	}

	for _, msg := range msgs {
106
		switch strings.ToLower(msg.Role) {
Bruce MacDonald's avatar
Bruce MacDonald committed
107
		case "system":
108
			if currentVars.System != "" {
Bruce MacDonald's avatar
Bruce MacDonald committed
109
110
111
112
113
114
				if err := writePrompt(); err != nil {
					return "", err
				}
			}
			currentVars.System = msg.Content
		case "user":
115
			if currentVars.Prompt != "" {
Bruce MacDonald's avatar
Bruce MacDonald committed
116
117
118
119
120
121
122
123
124
125
126
127
128
				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)
		}
129
130
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
131
132
133
134
135
136
137
138
	// 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
139
140
}

141
142
143
type ManifestV2 struct {
	SchemaVersion int      `json:"schemaVersion"`
	MediaType     string   `json:"mediaType"`
Michael Yang's avatar
Michael Yang committed
144
	Config        *Layer   `json:"config"`
145
146
147
148
	Layers        []*Layer `json:"layers"`
}

type ConfigV2 struct {
149
150
151
152
153
154
	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"`

155
	// required by spec
156
157
	Architecture string `json:"architecture"`
	OS           string `json:"os"`
158
	RootFS       RootFS `json:"rootfs"`
159
160
}

Michael Yang's avatar
Michael Yang committed
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
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
	}
}

191
192
193
194
195
type RootFS struct {
	Type    string   `json:"type"`
	DiffIDs []string `json:"diff_ids"`
}

Michael Yang's avatar
Michael Yang committed
196
func (m *ManifestV2) GetTotalSize() (total int64) {
Patrick Devine's avatar
Patrick Devine committed
197
198
199
	for _, layer := range m.Layers {
		total += layer.Size
	}
Michael Yang's avatar
Michael Yang committed
200

Patrick Devine's avatar
Patrick Devine committed
201
202
203
204
	total += m.Config.Size
	return total
}

Patrick Devine's avatar
Patrick Devine committed
205
func GetManifest(mp ModelPath) (*ManifestV2, string, error) {
206
	fp, err := mp.GetManifestPath()
207
	if err != nil {
Patrick Devine's avatar
Patrick Devine committed
208
		return nil, "", err
209
	}
210

211
	if _, err = os.Stat(fp); err != nil {
Patrick Devine's avatar
Patrick Devine committed
212
		return nil, "", err
213
214
215
216
	}

	var manifest *ManifestV2

217
	bts, err := os.ReadFile(fp)
218
	if err != nil {
Patrick Devine's avatar
Patrick Devine committed
219
		return nil, "", fmt.Errorf("couldn't open file '%s'", fp)
220
221
	}

Patrick Devine's avatar
Patrick Devine committed
222
223
224
	shaSum := sha256.Sum256(bts)
	shaStr := hex.EncodeToString(shaSum[:])

225
	if err := json.Unmarshal(bts, &manifest); err != nil {
Patrick Devine's avatar
Patrick Devine committed
226
		return nil, "", err
227
228
	}

Patrick Devine's avatar
Patrick Devine committed
229
	return manifest, shaStr, nil
230
231
232
}

func GetModel(name string) (*Model, error) {
233
	mp := ParseModelPath(name)
Patrick Devine's avatar
Patrick Devine committed
234
	manifest, digest, err := GetManifest(mp)
235
236
237
238
239
	if err != nil {
		return nil, err
	}

	model := &Model{
240
241
242
243
244
		Name:      mp.GetFullTagname(),
		ShortName: mp.GetShortTagname(),
		Digest:    digest,
		Template:  "{{ .Prompt }}",
		License:   []string{},
245
246
	}

247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
	filename, err := GetBlobsPath(manifest.Config.Digest)
	if err != nil {
		return nil, err
	}

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

	if err := json.NewDecoder(configFile).Decode(&model.Config); err != nil {
		return nil, err
	}

262
	for _, layer := range manifest.Layers {
Patrick Devine's avatar
Patrick Devine committed
263
		filename, err := GetBlobsPath(layer.Digest)
264
265
266
267
		if err != nil {
			return nil, err
		}

268
269
270
		switch layer.MediaType {
		case "application/vnd.ollama.image.model":
			model.ModelPath = filename
Patrick Devine's avatar
Patrick Devine committed
271
			model.OriginalModel = layer.From
272
		case "application/vnd.ollama.image.embed":
273
274
275
			// 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.")
276
277
		case "application/vnd.ollama.image.adapter":
			model.AdapterPaths = append(model.AdapterPaths, filename)
Michael Yang's avatar
Michael Yang committed
278
279
		case "application/vnd.ollama.image.projector":
			model.ProjectorPaths = append(model.ProjectorPaths, filename)
280
281
282
283
284
285
286
287
288
		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)
289
290
291
			if err != nil {
				return nil, err
			}
292
293

			model.System = string(bts)
294
295
296
297
298
299
300
		case "application/vnd.ollama.image.prompt":
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}

			model.Template = string(bts)
301
		case "application/vnd.ollama.image.params":
Michael Yang's avatar
Michael Yang committed
302
303
304
305
306
			params, err := os.Open(filename)
			if err != nil {
				return nil, err
			}
			defer params.Close()
307

308
			// parse model options parameters into a map so that we can see which fields have been specified explicitly
309
			if err = json.NewDecoder(params).Decode(&model.Options); err != nil {
310
311
				return nil, err
			}
Patrick Devine's avatar
Patrick Devine committed
312
313
314
315
316
317
		case "application/vnd.ollama.image.license":
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}
			model.License = append(model.License, string(bts))
318
319
320
321
322
323
		}
	}

	return model, nil
}

324
325
func realpath(mfDir, from string) string {
	abspath, err := filepath.Abs(from)
Michael Yang's avatar
Michael Yang committed
326
	if err != nil {
327
		return from
328
329
	}

Michael Yang's avatar
Michael Yang committed
330
	home, err := os.UserHomeDir()
331
	if err != nil {
Michael Yang's avatar
Michael Yang committed
332
		return abspath
333
334
	}

335
	if from == "~" {
Michael Yang's avatar
Michael Yang committed
336
		return home
337
338
339
340
341
342
343
	} 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)
344
345
	}

Michael Yang's avatar
Michael Yang committed
346
347
348
	return abspath
}

349
func CreateModel(ctx context.Context, name, modelFileDir string, commands []parser.Command, fn func(resp api.ProgressResponse)) error {
350
351
	config := ConfigV2{
		OS:           "linux",
Michael Yang's avatar
Michael Yang committed
352
		Architecture: "amd64",
Michael Yang's avatar
Michael Yang committed
353
354
355
		RootFS: RootFS{
			Type: "layers",
		},
356
357
	}

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

Michael Yang's avatar
Michael Yang committed
360
	var layers Layers
Michael Yang's avatar
Michael Yang committed
361

362
	params := make(map[string][]string)
Michael Yang's avatar
Michael Yang committed
363
364
	fromParams := make(map[string]any)

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

369
370
		switch c.Name {
		case "model":
Michael Yang's avatar
Michael Yang committed
371
372
373
374
375
376
377
378
379
			if strings.HasPrefix(c.Args, "@") {
				blobPath, err := GetBlobsPath(strings.TrimPrefix(c.Args, "@"))
				if err != nil {
					return err
				}

				c.Args = blobPath
			}

380
			bin, err := os.Open(realpath(modelFileDir, c.Args))
381
			if err != nil {
Michael Yang's avatar
Michael Yang committed
382
383
384
385
386
387
388
				// 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 {
389
390
391
						return err
					}

Michael Yang's avatar
Michael Yang committed
392
					manifest, _, err = GetManifest(modelpath)
393
394
395
					if err != nil {
						return err
					}
Michael Yang's avatar
Michael Yang committed
396
397
				case err != nil:
					return err
398
				}
399

400
				fn(api.ProgressResponse{Status: "reading model metadata"})
Michael Yang's avatar
Michael Yang committed
401
				fromConfigPath, err := GetBlobsPath(manifest.Config.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
				fromConfigFile, err := os.Open(fromConfigPath)
Michael Yang's avatar
Michael Yang committed
407
408
409
				if err != nil {
					return err
				}
Michael Yang's avatar
Michael Yang committed
410
				defer fromConfigFile.Close()
Michael Yang's avatar
Michael Yang committed
411

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

Michael Yang's avatar
Michael Yang committed
417
418
419
420
				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
421

Michael Yang's avatar
Michael Yang committed
422
423
424
425
				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
426
427
428
429
						if err != nil {
							return err
						}

Michael Yang's avatar
Michael Yang committed
430
						fromParamsFile, err := os.Open(fromParamsPath)
Michael Yang's avatar
Michael Yang committed
431
432
433
						if err != nil {
							return err
						}
Michael Yang's avatar
Michael Yang committed
434
						defer fromParamsFile.Close()
Michael Yang's avatar
Michael Yang committed
435

Michael Yang's avatar
Michael Yang committed
436
						if err := json.NewDecoder(fromParamsFile).Decode(&fromParams); err != nil {
Michael Yang's avatar
Michael Yang committed
437
438
439
440
							return err
						}
					}

Michael Yang's avatar
Michael Yang committed
441
					layer, err := NewLayerFromLayer(layer.Digest, layer.MediaType, modelpath.GetShortTagname())
442
443
444
					if err != nil {
						return err
					}
Michael Yang's avatar
Michael Yang committed
445

Michael Yang's avatar
Michael Yang committed
446
					layers.Add(layer)
447
				}
Michael Yang's avatar
Michael Yang committed
448
449
450

				deleteMap[manifest.Config.Digest] = struct{}{}
				continue
451
			}
Michael Yang's avatar
Michael Yang committed
452
			defer bin.Close()
453

454
455
456
			var offset int64
			for {
				fn(api.ProgressResponse{Status: "creating model layer"})
457

458
459
460
461
462
463
464
				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
465

Michael Yang's avatar
Michael Yang committed
466
467
468
469
				config.SetModelFormat(ggml.Name())
				config.SetModelFamily(ggml.ModelFamily())
				config.SetModelType(ggml.ModelType())
				config.SetFileType(ggml.FileType())
470

471
472
473
474
				mediatype := mediatype
				if ggml.ModelFamily() == "clip" {
					mediatype = "application/vnd.ollama.image.projector"
				}
475

476
477
478
479
480
481
482
483
484
485
				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
486
		case "adapter":
487
488
489
490
491
492
493
494
			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
495

Michael Yang's avatar
Michael Yang committed
496
			fn(api.ProgressResponse{Status: "creating adapter layer"})
497
			bin, err := os.Open(realpath(modelFileDir, c.Args))
498
			if err != nil {
Michael Yang's avatar
Michael Yang committed
499
				return err
500
			}
Michael Yang's avatar
Michael Yang committed
501
			defer bin.Close()
502

Michael Yang's avatar
Michael Yang committed
503
			layer, err := NewLayer(bin, mediatype)
504
			if err != nil {
Michael Yang's avatar
Michael Yang committed
505
				return err
506
			}
Bruce MacDonald's avatar
Bruce MacDonald committed
507

Michael Yang's avatar
Michael Yang committed
508
			layers.Add(layer)
Michael Yang's avatar
Michael Yang committed
509
510
		case "license":
			fn(api.ProgressResponse{Status: "creating license layer"})
Michael Yang's avatar
Michael Yang committed
511
512
513

			bin := strings.NewReader(c.Args)
			layer, err := NewLayer(bin, mediatype)
Bruce MacDonald's avatar
Bruce MacDonald committed
514
515
516
517
			if err != nil {
				return err
			}

Michael Yang's avatar
Michael Yang committed
518
			layers.Add(layer)
Michael Yang's avatar
Michael Yang committed
519
520
521
		case "template", "system":
			fn(api.ProgressResponse{Status: fmt.Sprintf("creating %s layer", c.Name)})

Michael Yang's avatar
Michael Yang committed
522
523
			bin := strings.NewReader(c.Args)
			layer, err := NewLayer(bin, mediatype)
524
			if err != nil {
525
				return err
526
			}
527

Michael Yang's avatar
Michael Yang committed
528
			layers.Replace(layer)
529
		default:
530
			params[c.Name] = append(params[c.Name], c.Args)
531
532
533
		}
	}

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

537
		formattedParams, err := api.FormatParams(params)
538
		if err != nil {
Michael Yang's avatar
Michael Yang committed
539
			return err
540
		}
541

Michael Yang's avatar
Michael Yang committed
542
		for k, v := range fromParams {
Michael Yang's avatar
Michael Yang committed
543
544
545
546
547
			if _, ok := formattedParams[k]; !ok {
				formattedParams[k] = v
			}
		}

Michael Yang's avatar
Michael Yang committed
548
		if config.ModelType == "65B" {
Michael Yang's avatar
Michael Yang committed
549
			if gqa, ok := formattedParams["gqa"].(int); ok && gqa == 8 {
Michael Yang's avatar
Michael Yang committed
550
551
552
553
				config.ModelType = "70B"
			}
		}

Michael Yang's avatar
Michael Yang committed
554
555
		var b bytes.Buffer
		if err := json.NewEncoder(&b).Encode(formattedParams); err != nil {
556
557
558
			return err
		}

Michael Yang's avatar
Michael Yang committed
559
		fn(api.ProgressResponse{Status: "creating config layer"})
Michael Yang's avatar
Michael Yang committed
560
		layer, err := NewLayer(&b, "application/vnd.ollama.image.params")
561
		if err != nil {
Michael Yang's avatar
Michael Yang committed
562
			return err
563
		}
Michael Yang's avatar
Michael Yang committed
564

Michael Yang's avatar
Michael Yang committed
565
		layers.Replace(layer)
566
567
	}

Michael Yang's avatar
Michael Yang committed
568
569
570
	digests := make([]string, len(layers.items))
	for i, layer := range layers.items {
		digests[i] = layer.Digest
571
572
	}

Michael Yang's avatar
Michael Yang committed
573
	config.RootFS.DiffIDs = digests
Michael Yang's avatar
Michael Yang committed
574

Michael Yang's avatar
Michael Yang committed
575
576
	var b bytes.Buffer
	if err := json.NewEncoder(&b).Encode(config); err != nil {
577
578
579
		return err
	}

Michael Yang's avatar
Michael Yang committed
580
581
	configLayer, err := NewLayer(&b, "application/vnd.docker.container.image.v1+json")
	if err != nil {
582
583
584
		return err
	}

Michael Yang's avatar
Michael Yang committed
585
	delete(deleteMap, configLayer.Digest)
586

Michael Yang's avatar
Michael Yang committed
587
588
	for _, layer := range append(layers.items, configLayer) {
		committed, err := layer.Commit()
589
590
591
		if err != nil {
			return err
		}
592

Michael Yang's avatar
Michael Yang committed
593
594
595
		status := "writing layer"
		if !committed {
			status = "using already created layer"
596
597
		}

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

Michael Yang's avatar
Michael Yang committed
600
		delete(deleteMap, layer.Digest)
601
602
	}

Michael Yang's avatar
Michael Yang committed
603
604
	fn(api.ProgressResponse{Status: "writing manifest"})
	if err := WriteManifest(name, configLayer, layers.items); err != nil {
605
606
		return err
	}
607

Michael Yang's avatar
Michael Yang committed
608
609
610
	if noprune := os.Getenv("OLLAMA_NOPRUNE"); noprune == "" {
		if err := deleteUnusedLayers(nil, deleteMap, false); err != nil {
			return err
611
612
613
		}
	}

Michael Yang's avatar
Michael Yang committed
614
615
	fn(api.ProgressResponse{Status: "success"})
	return nil
616
617
}

Patrick Devine's avatar
Patrick Devine committed
618
func CopyModel(src, dest string) error {
619
	srcModelPath := ParseModelPath(src)
620
	srcPath, err := srcModelPath.GetManifestPath()
621
622
623
624
	if err != nil {
		return err
	}

625
	destModelPath := ParseModelPath(dest)
626
	destPath, err := destModelPath.GetManifestPath()
Patrick Devine's avatar
Patrick Devine committed
627
628
629
	if err != nil {
		return err
	}
630
631
632
	if err := os.MkdirAll(filepath.Dir(destPath), 0o755); err != nil {
		return err
	}
Patrick Devine's avatar
Patrick Devine committed
633
634

	// copy the file
Michael Yang's avatar
Michael Yang committed
635
	input, err := os.ReadFile(srcPath)
Patrick Devine's avatar
Patrick Devine committed
636
637
638
639
640
	if err != nil {
		fmt.Println("Error reading file:", err)
		return err
	}

Michael Yang's avatar
Michael Yang committed
641
	err = os.WriteFile(destPath, input, 0o644)
Patrick Devine's avatar
Patrick Devine committed
642
643
644
645
646
647
648
649
	if err != nil {
		fmt.Println("Error reading file:", err)
		return err
	}

	return nil
}

Michael Yang's avatar
Michael Yang committed
650
func deleteUnusedLayers(skipModelPath *ModelPath, deleteMap map[string]struct{}, dryRun bool) error {
651
652
653
654
	fp, err := GetManifestPath()
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
655
656
657
658

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

Michael Yang's avatar
Michael Yang committed
661
662
663
664
		dir, file := filepath.Split(path)
		dir = strings.Trim(strings.TrimPrefix(dir, fp), string(os.PathSeparator))
		tag := strings.Join([]string{dir, file}, ":")
		fmp := ParseModelPath(tag)
665

Michael Yang's avatar
Michael Yang committed
666
		// skip the manifest we're trying to delete
667
		if skipModelPath != nil && skipModelPath.GetFullTagname() == fmp.GetFullTagname() {
Michael Yang's avatar
Michael Yang committed
668
			return nil
669
		}
Michael Yang's avatar
Michael Yang committed
670
671
672
673
674
675
676
677
678
679
680
681

		// 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)
682
		return nil
Michael Yang's avatar
Michael Yang committed
683
684
685
	}

	if err := filepath.Walk(fp, walkFunc); err != nil {
Michael Yang's avatar
Michael Yang committed
686
687
		return err
	}
688
689

	// only delete the files which are still in the deleteMap
Michael Yang's avatar
Michael Yang committed
690
691
692
693
694
695
696
697
698
	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)
699
700
				continue
			}
Michael Yang's avatar
Michael Yang committed
701
702
		} else {
			log.Printf("wanted to remove: %s", fp)
703
704
705
		}
	}

706
707
708
709
	return nil
}

func PruneLayers() error {
Michael Yang's avatar
Michael Yang committed
710
	deleteMap := make(map[string]struct{})
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
	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
727
728
729
		if strings.HasPrefix(name, "sha256:") {
			deleteMap[name] = struct{}{}
		}
730
731
732
733
734
735
736
737
738
739
740
741
742
743
	}

	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
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
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
}

777
778
779
780
781
782
783
func DeleteModel(name string) error {
	mp := ParseModelPath(name)
	manifest, _, err := GetManifest(mp)
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
784
	deleteMap := make(map[string]struct{})
785
	for _, layer := range manifest.Layers {
Michael Yang's avatar
Michael Yang committed
786
		deleteMap[layer.Digest] = struct{}{}
787
	}
Michael Yang's avatar
Michael Yang committed
788
	deleteMap[manifest.Config.Digest] = struct{}{}
789
790
791
792
793
794

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

795
	fp, err := mp.GetManifestPath()
796
797
798
799
800
801
802
803
804
805
806
807
	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
808
func ShowModelfile(model *Model) (string, error) {
Michael Yang's avatar
Michael Yang committed
809
	var mt struct {
Patrick Devine's avatar
Patrick Devine committed
810
		*Model
Michael Yang's avatar
Michael Yang committed
811
		From       string
Michael Yang's avatar
Michael Yang committed
812
		Parameters map[string][]any
Patrick Devine's avatar
Patrick Devine committed
813
814
	}

Michael Yang's avatar
Michael Yang committed
815
	mt.Parameters = make(map[string][]any)
Patrick Devine's avatar
Patrick Devine committed
816
	for k, v := range model.Options {
Michael Yang's avatar
Michael Yang committed
817
818
819
		if s, ok := v.([]any); ok {
			mt.Parameters[k] = s
			continue
Patrick Devine's avatar
Patrick Devine committed
820
821
		}

Michael Yang's avatar
Michael Yang committed
822
		mt.Parameters[k] = []any{v}
Patrick Devine's avatar
Patrick Devine committed
823
824
	}

Michael Yang's avatar
Michael Yang committed
825
826
	mt.Model = model
	mt.From = model.ModelPath
Patrick Devine's avatar
Patrick Devine committed
827

Michael Yang's avatar
Michael Yang committed
828
	if model.OriginalModel != "" {
Daniel Reis's avatar
Daniel Reis committed
829
		mt.From = model.OriginalModel
Patrick Devine's avatar
Patrick Devine committed
830
831
832
833
834
835
836
837
	}

	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 }}"""
838
839

{{- if .System }}
Patrick Devine's avatar
Patrick Devine committed
840
SYSTEM """{{ .System }}"""
841
{{- end }}
842
843
844
845

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

Michael Yang's avatar
Michael Yang committed
847
848
849
850
{{- range $k, $v := .Parameters }}
{{- range $parameter := $v }}
PARAMETER {{ $k }} {{ printf "%#v" $parameter }}
{{- end }}
Michael Yang's avatar
Michael Yang committed
851
{{- end }}`
Patrick Devine's avatar
Patrick Devine committed
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868

	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
}

869
func PushModel(ctx context.Context, name string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
870
	mp := ParseModelPath(name)
871
872
	fn(api.ProgressResponse{Status: "retrieving manifest"})

873
874
875
876
	if mp.ProtocolScheme == "http" && !regOpts.Insecure {
		return fmt.Errorf("insecure protocol http")
	}

Patrick Devine's avatar
Patrick Devine committed
877
	manifest, _, err := GetManifest(mp)
878
	if err != nil {
879
		fn(api.ProgressResponse{Status: "couldn't retrieve manifest"})
880
881
882
883
		return err
	}

	var layers []*Layer
Jeffrey Morgan's avatar
Jeffrey Morgan committed
884
	layers = append(layers, manifest.Layers...)
Michael Yang's avatar
Michael Yang committed
885
	layers = append(layers, manifest.Config)
886
887

	for _, layer := range layers {
Michael Yang's avatar
Michael Yang committed
888
		if err := uploadBlob(ctx, mp, layer, regOpts, fn); err != nil {
889
			log.Printf("error uploading blob: %v", err)
890
891
892
			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())
			}
893
894
			return err
		}
895
896
	}

897
	fn(api.ProgressResponse{Status: "pushing manifest"})
Michael Yang's avatar
Michael Yang committed
898
899
	requestURL := mp.BaseURL()
	requestURL = requestURL.JoinPath("v2", mp.GetNamespaceRepository(), "manifests", mp.Tag)
900
901
902
903
904
905

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

Michael Yang's avatar
Michael Yang committed
906
907
	headers := make(http.Header)
	headers.Set("Content-Type", "application/vnd.docker.distribution.manifest.v2+json")
Michael Yang's avatar
Michael Yang committed
908
	resp, err := makeRequestWithRetry(ctx, http.MethodPut, requestURL, headers, bytes.NewReader(manifestJSON), regOpts)
909
910
911
912
913
	if err != nil {
		return err
	}
	defer resp.Body.Close()

914
	fn(api.ProgressResponse{Status: "success"})
915
916
917
918

	return nil
}

919
func PullModel(ctx context.Context, name string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
920
921
	mp := ParseModelPath(name)

922
923
924
925
926
	var manifest *ManifestV2
	var err error
	var noprune string

	// build deleteMap to prune unused layers
Michael Yang's avatar
Michael Yang committed
927
	deleteMap := make(map[string]struct{})
928
929
930
931
932
933
934
935
936

	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
937
				deleteMap[l.Digest] = struct{}{}
938
			}
Michael Yang's avatar
Michael Yang committed
939
			deleteMap[manifest.Config.Digest] = struct{}{}
940
941
942
		}
	}

943
944
	if mp.ProtocolScheme == "http" && !regOpts.Insecure {
		return fmt.Errorf("insecure protocol http")
945
	}
946

947
	fn(api.ProgressResponse{Status: "pulling manifest"})
948

949
	manifest, err = pullModelManifest(ctx, mp, regOpts)
950
	if err != nil {
951
		return fmt.Errorf("pull model manifest: %s", err)
952
953
954
	}

	var layers []*Layer
Bruce MacDonald's avatar
Bruce MacDonald committed
955
	layers = append(layers, manifest.Layers...)
Michael Yang's avatar
Michael Yang committed
956
	layers = append(layers, manifest.Config)
957
958

	for _, layer := range layers {
959
960
961
962
963
964
965
966
		if err := downloadBlob(
			ctx,
			downloadOpts{
				mp:      mp,
				digest:  layer.Digest,
				regOpts: regOpts,
				fn:      fn,
			}); err != nil {
967
968
			return err
		}
969
		delete(deleteMap, layer.Digest)
970
	}
971
	delete(deleteMap, manifest.Config.Digest)
972

Michael Yang's avatar
Michael Yang committed
973
974
975
	fn(api.ProgressResponse{Status: "verifying sha256 digest"})
	for _, layer := range layers {
		if err := verifyBlob(layer.Digest); err != nil {
976
977
978
979
980
981
982
983
984
985
986
			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
987
988
989
990
			return err
		}
	}

991
	fn(api.ProgressResponse{Status: "writing manifest"})
992

993
	manifestJSON, err := json.Marshal(manifest)
994
995
996
997
	if err != nil {
		return err
	}

998
	fp, err := mp.GetManifestPath()
999
1000
1001
	if err != nil {
		return err
	}
1002
1003
1004
	if err := os.MkdirAll(filepath.Dir(fp), 0o755); err != nil {
		return err
	}
1005

Bruce MacDonald's avatar
Bruce MacDonald committed
1006
	err = os.WriteFile(fp, manifestJSON, 0o644)
1007
1008
1009
1010
1011
	if err != nil {
		log.Printf("couldn't write to %s", fp)
		return err
	}

1012
1013
1014
1015
1016
1017
1018
1019
	if noprune == "" {
		fn(api.ProgressResponse{Status: "removing any unused layers"})
		err = deleteUnusedLayers(nil, deleteMap, false)
		if err != nil {
			return err
		}
	}

1020
	fn(api.ProgressResponse{Status: "success"})
1021
1022
1023
1024

	return nil
}

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

Michael Yang's avatar
Michael Yang committed
1028
1029
	headers := make(http.Header)
	headers.Set("Accept", "application/vnd.docker.distribution.manifest.v2+json")
Michael Yang's avatar
Michael Yang committed
1030
	resp, err := makeRequestWithRetry(ctx, http.MethodGet, requestURL, headers, nil, regOpts)
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
	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
1045
func GetSHA256Digest(r io.Reader) (string, int64) {
Michael Yang's avatar
Michael Yang committed
1046
1047
1048
1049
1050
1051
	h := sha256.New()
	n, err := io.Copy(h, r)
	if err != nil {
		log.Fatal(err)
	}

Michael Yang's avatar
Michael Yang committed
1052
	return fmt.Sprintf("sha256:%x", h.Sum(nil)), n
1053
1054
}

1055
1056
var errUnauthorized = fmt.Errorf("unauthorized")

Michael Yang's avatar
Michael Yang committed
1057
func makeRequestWithRetry(ctx context.Context, method string, requestURL *url.URL, headers http.Header, body io.ReadSeeker, regOpts *RegistryOptions) (*http.Response, error) {
1058
1059
1060
1061
1062
	resp, err := makeRequest(ctx, method, requestURL, headers, body, regOpts)
	if err != nil {
		if !errors.Is(err, context.Canceled) {
			log.Printf("request failed: %v", err)
		}
1063

1064
1065
1066
1067
1068
1069
1070
1071
1072
		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
1073
1074
1075
		if err != nil {
			return nil, err
		}
1076
1077
1078
		regOpts.Token = token
		if body != nil {
			_, err = body.Seek(0, io.SeekStart)
Michael Yang's avatar
Michael Yang committed
1079
1080
1081
1082
			if err != nil {
				return nil, err
			}
		}
1083
1084
1085
1086
1087
1088
1089

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

		return resp, err
1090
1091
1092
1093
1094
1095
1096
1097
	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
1098
1099
	}

1100
	return resp, nil
Michael Yang's avatar
Michael Yang committed
1101
1102
}

Michael Yang's avatar
Michael Yang committed
1103
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
1104
	if requestURL.Scheme != "http" && regOpts != nil && regOpts.Insecure {
Michael Yang's avatar
Michael Yang committed
1105
		requestURL.Scheme = "http"
1106
1107
	}

Michael Yang's avatar
Michael Yang committed
1108
	req, err := http.NewRequestWithContext(ctx, method, requestURL.String(), body)
1109
1110
1111
1112
	if err != nil {
		return nil, err
	}

Michael Yang's avatar
Michael Yang committed
1113
1114
1115
1116
	if headers != nil {
		req.Header = headers
	}

Michael Yang's avatar
Michael Yang committed
1117
1118
1119
1120
1121
1122
	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)
		}
1123
1124
	}

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

Michael Yang's avatar
Michael Yang committed
1127
1128
1129
1130
1131
1132
1133
1134
1135
	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
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
	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)
1148
1149
1150
1151
1152
1153
	if err != nil {
		return nil, err
	}

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

Patrick Devine's avatar
Patrick Devine committed
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
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"),
	}
}

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

Michael Yang's avatar
Michael Yang committed
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
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 {
1204
		return fmt.Errorf("%w: want %s, got %s", errDigestMismatch, digest, fileDigest)
Michael Yang's avatar
Michael Yang committed
1205
1206
1207
1208
	}

	return nil
}