images.go 30.3 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
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
13
	"log/slog"
14
	"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
	"runtime"
19
20
	"strconv"
	"strings"
Quinn Slack's avatar
Quinn Slack committed
21
	"text/template"
22
	"text/template/parse"
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"
Michael Yang's avatar
Michael Yang committed
29
	"github.com/jmorganca/ollama/version"
30
31
)

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

39
type Model struct {
Michael Yang's avatar
Michael Yang committed
40
	Name           string `json:"name"`
41
	Config         ConfigV2
Michael Yang's avatar
Michael Yang committed
42
43
	ShortName      string
	ModelPath      string
44
	ParentModel    string
Michael Yang's avatar
Michael Yang committed
45
46
47
48
49
50
	AdapterPaths   []string
	ProjectorPaths []string
	Template       string
	System         string
	License        []string
	Digest         string
Patrick Devine's avatar
Patrick Devine committed
51
	Size           int64
Michael Yang's avatar
Michael Yang committed
52
	Options        map[string]interface{}
53
54
55
56
57
58
	Messages       []Message
}

type Message struct {
	Role    string `json:"role"`
	Content string `json:"content"`
59
60
}

Bruce MacDonald's avatar
Bruce MacDonald committed
61
62
63
64
65
66
type PromptVars struct {
	System   string
	Prompt   string
	Response string
	First    bool
}
67

68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
// extractParts extracts the parts of the template before and after the {{.Response}} node.
func extractParts(tmplStr string) (pre string, post string, err error) {
	tmpl, err := template.New("").Parse(tmplStr)
	if err != nil {
		return "", "", err
	}

	var foundResponse bool

	for _, node := range tmpl.Tree.Root.Nodes {
		if node.Type() == parse.NodeAction && node.String() == "{{.Response}}" {
			foundResponse = true
		}
		if !foundResponse {
			pre += node.String()
		} else {
			post += node.String()
		}
	}

	return pre, post, nil
}

func Prompt(promptTemplate string, p PromptVars) (string, error) {
Bruce MacDonald's avatar
Bruce MacDonald committed
92
	var prompt strings.Builder
93
	// Use the "missingkey=zero" option to handle missing variables without panicking
94
	tmpl, err := template.New("").Option("missingkey=zero").Parse(promptTemplate)
95
96
97
98
	if err != nil {
		return "", err
	}

99
100
101
102
103
	vars := map[string]any{
		"System":   p.System,
		"Prompt":   p.Prompt,
		"Response": p.Response,
		"First":    p.First,
104
105
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
106
	var sb strings.Builder
107
	if err := tmpl.Execute(&sb, vars); err != nil {
Bruce MacDonald's avatar
Bruce MacDonald committed
108
109
110
		return "", err
	}
	prompt.WriteString(sb.String())
111
112
113
114
115
116

	if !strings.Contains(prompt.String(), p.Response) {
		// if the response is not in the prompt template, append it to the end
		prompt.WriteString(p.Response)
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
117
118
	return prompt.String(), nil
}
119

120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
// PreResponsePrompt returns the prompt before the response tag
func (m *Model) PreResponsePrompt(p PromptVars) (string, error) {
	if p.System == "" {
		// use the default system prompt for this model if one is not specified
		p.System = m.System
	}
	pre, _, err := extractParts(m.Template)
	if err != nil {
		return "", err
	}

	return Prompt(pre, p)
}

// PostResponseTemplate returns the template after the response tag
func (m *Model) PostResponseTemplate(p PromptVars) (string, error) {
	if p.System == "" {
		// use the default system prompt for this model if one is not specified
		p.System = m.System
	}
	_, post, err := extractParts(m.Template)
	if err != nil {
		return "", err
	}

	if post == "" {
		// if there is no post-response template, return the provided response
		return p.Response, nil
	}

	return Prompt(post, p)
}

153
func (m *Model) ChatPrompt(msgs []api.Message) (string, []api.ImageData, error) {
Bruce MacDonald's avatar
Bruce MacDonald committed
154
155
	// build the prompt from the list of messages
	var prompt strings.Builder
156
	var currentImages []api.ImageData
Bruce MacDonald's avatar
Bruce MacDonald committed
157
	currentVars := PromptVars{
158
159
		First:  true,
		System: m.System,
Bruce MacDonald's avatar
Bruce MacDonald committed
160
161
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
162
	writePrompt := func() error {
163
		p, err := Prompt(m.Template, currentVars)
Bruce MacDonald's avatar
Bruce MacDonald committed
164
165
166
167
168
169
170
171
172
		if err != nil {
			return err
		}
		prompt.WriteString(p)
		currentVars = PromptVars{}
		return nil
	}

	for _, msg := range msgs {
173
		switch strings.ToLower(msg.Role) {
Bruce MacDonald's avatar
Bruce MacDonald committed
174
		case "system":
175
			if currentVars.System != "" {
Bruce MacDonald's avatar
Bruce MacDonald committed
176
				if err := writePrompt(); err != nil {
177
					return "", nil, err
Bruce MacDonald's avatar
Bruce MacDonald committed
178
179
180
181
				}
			}
			currentVars.System = msg.Content
		case "user":
182
			if currentVars.Prompt != "" {
Bruce MacDonald's avatar
Bruce MacDonald committed
183
				if err := writePrompt(); err != nil {
184
					return "", nil, err
Bruce MacDonald's avatar
Bruce MacDonald committed
185
186
187
				}
			}
			currentVars.Prompt = msg.Content
188
			currentImages = msg.Images
Bruce MacDonald's avatar
Bruce MacDonald committed
189
190
191
		case "assistant":
			currentVars.Response = msg.Content
			if err := writePrompt(); err != nil {
192
				return "", nil, err
Bruce MacDonald's avatar
Bruce MacDonald committed
193
194
			}
		default:
195
			return "", nil, fmt.Errorf("invalid role: %s, role must be one of [system, user, assistant]", msg.Role)
Bruce MacDonald's avatar
Bruce MacDonald committed
196
		}
197
198
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
199
200
	// Append the last set of vars if they are non-empty
	if currentVars.Prompt != "" || currentVars.System != "" {
201
202
203
		p, err := m.PreResponsePrompt(currentVars)
		if err != nil {
			return "", nil, fmt.Errorf("pre-response template: %w", err)
Bruce MacDonald's avatar
Bruce MacDonald committed
204
		}
205
		prompt.WriteString(p)
Bruce MacDonald's avatar
Bruce MacDonald committed
206
207
	}

208
	return prompt.String(), currentImages, nil
209
210
}

211
212
213
type ManifestV2 struct {
	SchemaVersion int      `json:"schemaVersion"`
	MediaType     string   `json:"mediaType"`
Michael Yang's avatar
Michael Yang committed
214
	Config        *Layer   `json:"config"`
215
216
217
218
	Layers        []*Layer `json:"layers"`
}

type ConfigV2 struct {
219
220
221
222
223
224
	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"`

225
	// required by spec
226
227
	Architecture string `json:"architecture"`
	OS           string `json:"os"`
228
	RootFS       RootFS `json:"rootfs"`
229
230
}

Michael Yang's avatar
Michael Yang committed
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
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
	}
}

261
262
263
264
265
type RootFS struct {
	Type    string   `json:"type"`
	DiffIDs []string `json:"diff_ids"`
}

Michael Yang's avatar
Michael Yang committed
266
func (m *ManifestV2) GetTotalSize() (total int64) {
Patrick Devine's avatar
Patrick Devine committed
267
268
269
	for _, layer := range m.Layers {
		total += layer.Size
	}
Michael Yang's avatar
Michael Yang committed
270

Patrick Devine's avatar
Patrick Devine committed
271
272
273
274
	total += m.Config.Size
	return total
}

Patrick Devine's avatar
Patrick Devine committed
275
func GetManifest(mp ModelPath) (*ManifestV2, string, error) {
276
	fp, err := mp.GetManifestPath()
277
	if err != nil {
Patrick Devine's avatar
Patrick Devine committed
278
		return nil, "", err
279
	}
280

281
	if _, err = os.Stat(fp); err != nil {
Patrick Devine's avatar
Patrick Devine committed
282
		return nil, "", err
283
284
285
286
	}

	var manifest *ManifestV2

287
	bts, err := os.ReadFile(fp)
288
	if err != nil {
Patrick Devine's avatar
Patrick Devine committed
289
		return nil, "", fmt.Errorf("couldn't open file '%s'", fp)
290
291
	}

Patrick Devine's avatar
Patrick Devine committed
292
293
294
	shaSum := sha256.Sum256(bts)
	shaStr := hex.EncodeToString(shaSum[:])

295
	if err := json.Unmarshal(bts, &manifest); err != nil {
Patrick Devine's avatar
Patrick Devine committed
296
		return nil, "", err
297
298
	}

Patrick Devine's avatar
Patrick Devine committed
299
	return manifest, shaStr, nil
300
301
302
}

func GetModel(name string) (*Model, error) {
303
	mp := ParseModelPath(name)
Patrick Devine's avatar
Patrick Devine committed
304
	manifest, digest, err := GetManifest(mp)
305
306
307
308
309
	if err != nil {
		return nil, err
	}

	model := &Model{
310
311
312
313
314
		Name:      mp.GetFullTagname(),
		ShortName: mp.GetShortTagname(),
		Digest:    digest,
		Template:  "{{ .Prompt }}",
		License:   []string{},
Patrick Devine's avatar
Patrick Devine committed
315
		Size:      manifest.GetTotalSize(),
316
317
	}

318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
	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
	}

333
	for _, layer := range manifest.Layers {
Patrick Devine's avatar
Patrick Devine committed
334
		filename, err := GetBlobsPath(layer.Digest)
335
336
337
338
		if err != nil {
			return nil, err
		}

339
340
341
		switch layer.MediaType {
		case "application/vnd.ollama.image.model":
			model.ModelPath = filename
342
			model.ParentModel = layer.From
343
		case "application/vnd.ollama.image.embed":
344
345
			// Deprecated in versions  > 0.1.2
			// TODO: remove this warning in a future version
346
			slog.Info("WARNING: model contains embeddings, but embeddings in modelfiles have been deprecated and will be ignored.")
347
348
		case "application/vnd.ollama.image.adapter":
			model.AdapterPaths = append(model.AdapterPaths, filename)
Michael Yang's avatar
Michael Yang committed
349
350
		case "application/vnd.ollama.image.projector":
			model.ProjectorPaths = append(model.ProjectorPaths, filename)
351
352
353
354
355
356
357
358
359
		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)
360
361
362
			if err != nil {
				return nil, err
			}
363
364

			model.System = string(bts)
365
366
367
368
369
370
371
		case "application/vnd.ollama.image.prompt":
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}

			model.Template = string(bts)
372
		case "application/vnd.ollama.image.params":
Michael Yang's avatar
Michael Yang committed
373
374
375
376
377
			params, err := os.Open(filename)
			if err != nil {
				return nil, err
			}
			defer params.Close()
378

379
			// parse model options parameters into a map so that we can see which fields have been specified explicitly
380
			if err = json.NewDecoder(params).Decode(&model.Options); err != nil {
381
382
				return nil, err
			}
383
384
385
386
387
388
389
390
391
392
		case "application/vnd.ollama.image.messages":
			msgs, err := os.Open(filename)
			if err != nil {
				return nil, err
			}
			defer msgs.Close()

			if err = json.NewDecoder(msgs).Decode(&model.Messages); err != nil {
				return nil, err
			}
Patrick Devine's avatar
Patrick Devine committed
393
394
395
396
397
398
		case "application/vnd.ollama.image.license":
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}
			model.License = append(model.License, string(bts))
399
400
401
402
403
404
		}
	}

	return model, nil
}

405
406
func realpath(mfDir, from string) string {
	abspath, err := filepath.Abs(from)
Michael Yang's avatar
Michael Yang committed
407
	if err != nil {
408
		return from
409
410
	}

Michael Yang's avatar
Michael Yang committed
411
	home, err := os.UserHomeDir()
412
	if err != nil {
Michael Yang's avatar
Michael Yang committed
413
		return abspath
414
415
	}

416
	if from == "~" {
Michael Yang's avatar
Michael Yang committed
417
		return home
418
419
420
421
422
423
424
	} 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)
425
426
	}

Michael Yang's avatar
Michael Yang committed
427
428
429
	return abspath
}

430
func CreateModel(ctx context.Context, name, modelFileDir string, commands []parser.Command, fn func(resp api.ProgressResponse)) error {
431
432
433
434
435
436
437
	deleteMap := make(map[string]struct{})
	if manifest, _, err := GetManifest(ParseModelPath(name)); err == nil {
		for _, layer := range append(manifest.Layers, manifest.Config) {
			deleteMap[layer.Digest] = struct{}{}
		}
	}

438
439
	config := ConfigV2{
		OS:           "linux",
Michael Yang's avatar
Michael Yang committed
440
		Architecture: "amd64",
Michael Yang's avatar
Michael Yang committed
441
442
443
		RootFS: RootFS{
			Type: "layers",
		},
444
445
	}

Michael Yang's avatar
Michael Yang committed
446
	var layers Layers
447
	messages := []string{}
Michael Yang's avatar
Michael Yang committed
448

449
	params := make(map[string][]string)
Michael Yang's avatar
Michael Yang committed
450
451
	fromParams := make(map[string]any)

452
	for _, c := range commands {
Michael Yang's avatar
Michael Yang committed
453
454
		mediatype := fmt.Sprintf("application/vnd.ollama.image.%s", c.Name)

455
456
		switch c.Name {
		case "model":
Michael Yang's avatar
Michael Yang committed
457
458
459
460
461
462
463
464
465
			if strings.HasPrefix(c.Args, "@") {
				blobPath, err := GetBlobsPath(strings.TrimPrefix(c.Args, "@"))
				if err != nil {
					return err
				}

				c.Args = blobPath
			}

466
			bin, err := os.Open(realpath(modelFileDir, c.Args))
467
			if err != nil {
Michael Yang's avatar
Michael Yang committed
468
469
470
471
472
473
474
				// 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 {
475
476
477
						return err
					}

Michael Yang's avatar
Michael Yang committed
478
					manifest, _, err = GetManifest(modelpath)
479
480
481
					if err != nil {
						return err
					}
Michael Yang's avatar
Michael Yang committed
482
483
				case err != nil:
					return err
484
				}
485

486
				fn(api.ProgressResponse{Status: "reading model metadata"})
Michael Yang's avatar
Michael Yang committed
487
				fromConfigPath, err := GetBlobsPath(manifest.Config.Digest)
Michael Yang's avatar
Michael Yang committed
488
489
490
491
				if err != nil {
					return err
				}

Michael Yang's avatar
Michael Yang committed
492
				fromConfigFile, err := os.Open(fromConfigPath)
Michael Yang's avatar
Michael Yang committed
493
494
495
				if err != nil {
					return err
				}
Michael Yang's avatar
Michael Yang committed
496
				defer fromConfigFile.Close()
Michael Yang's avatar
Michael Yang committed
497

Michael Yang's avatar
Michael Yang committed
498
499
				var fromConfig ConfigV2
				if err := json.NewDecoder(fromConfigFile).Decode(&fromConfig); err != nil {
Michael Yang's avatar
Michael Yang committed
500
501
502
					return err
				}

Bruce MacDonald's avatar
Bruce MacDonald committed
503
504
505
506
507
				// if the model is still not in gguf format, error out
				if fromConfig.ModelFormat != "gguf" {
					return fmt.Errorf("%s is not in gguf format, this base model is not compatible with this version of ollama", c.Args)
				}

Michael Yang's avatar
Michael Yang committed
508
509
510
511
				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
512

Michael Yang's avatar
Michael Yang committed
513
514
515
516
				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
517
518
519
520
						if err != nil {
							return err
						}

Michael Yang's avatar
Michael Yang committed
521
						fromParamsFile, err := os.Open(fromParamsPath)
Michael Yang's avatar
Michael Yang committed
522
523
524
						if err != nil {
							return err
						}
Michael Yang's avatar
Michael Yang committed
525
						defer fromParamsFile.Close()
Michael Yang's avatar
Michael Yang committed
526

Michael Yang's avatar
Michael Yang committed
527
						if err := json.NewDecoder(fromParamsFile).Decode(&fromParams); err != nil {
Michael Yang's avatar
Michael Yang committed
528
529
530
531
							return err
						}
					}

Michael Yang's avatar
Michael Yang committed
532
					layer, err := NewLayerFromLayer(layer.Digest, layer.MediaType, modelpath.GetShortTagname())
533
534
535
					if err != nil {
						return err
					}
Michael Yang's avatar
Michael Yang committed
536

Michael Yang's avatar
Michael Yang committed
537
					layers.Add(layer)
538
				}
Michael Yang's avatar
Michael Yang committed
539
540
541

				deleteMap[manifest.Config.Digest] = struct{}{}
				continue
542
			}
Michael Yang's avatar
Michael Yang committed
543
			defer bin.Close()
544

545
			var offset int64
Bruce MacDonald's avatar
Bruce MacDonald committed
546
		CREATE:
547
548
			for {
				fn(api.ProgressResponse{Status: "creating model layer"})
549

550
551
				bin.Seek(offset, io.SeekStart)
				ggml, err := llm.DecodeGGML(bin)
Bruce MacDonald's avatar
Bruce MacDonald committed
552
553
554
555
556
557
558
559
560
				if err != nil {
					switch {
					case errors.Is(err, io.EOF):
						break CREATE
					case errors.Is(err, llm.ErrUnsupportedFormat):
						return fmt.Errorf("model binary specified in FROM field is not a valid gguf format model, %w", err)
					default:
						return err
					}
561
				}
Michael Yang's avatar
Michael Yang committed
562

Michael Yang's avatar
Michael Yang committed
563
564
565
566
				config.SetModelFormat(ggml.Name())
				config.SetModelFamily(ggml.ModelFamily())
				config.SetModelType(ggml.ModelType())
				config.SetFileType(ggml.FileType())
567

568
569
570
571
				mediatype := mediatype
				if ggml.ModelFamily() == "clip" {
					mediatype = "application/vnd.ollama.image.projector"
				}
572

573
574
575
576
577
578
579
580
581
582
				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
583
		case "adapter":
584
585
586
587
588
589
590
591
			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
592

Michael Yang's avatar
Michael Yang committed
593
			fn(api.ProgressResponse{Status: "creating adapter layer"})
594
			bin, err := os.Open(realpath(modelFileDir, c.Args))
595
			if err != nil {
Michael Yang's avatar
Michael Yang committed
596
				return err
597
			}
Michael Yang's avatar
Michael Yang committed
598
			defer bin.Close()
599

Michael Yang's avatar
Michael Yang committed
600
			layer, err := NewLayer(bin, mediatype)
601
			if err != nil {
Michael Yang's avatar
Michael Yang committed
602
				return err
603
			}
Bruce MacDonald's avatar
Bruce MacDonald committed
604

Michael Yang's avatar
Michael Yang committed
605
			layers.Add(layer)
Michael Yang's avatar
Michael Yang committed
606
607
		case "license":
			fn(api.ProgressResponse{Status: "creating license layer"})
Michael Yang's avatar
Michael Yang committed
608
609
610

			bin := strings.NewReader(c.Args)
			layer, err := NewLayer(bin, mediatype)
Bruce MacDonald's avatar
Bruce MacDonald committed
611
612
613
614
			if err != nil {
				return err
			}

Michael Yang's avatar
Michael Yang committed
615
			layers.Add(layer)
Michael Yang's avatar
Michael Yang committed
616
617
618
		case "template", "system":
			fn(api.ProgressResponse{Status: fmt.Sprintf("creating %s layer", c.Name)})

Michael Yang's avatar
Michael Yang committed
619
620
			bin := strings.NewReader(c.Args)
			layer, err := NewLayer(bin, mediatype)
621
			if err != nil {
622
				return err
623
			}
624

Michael Yang's avatar
Michael Yang committed
625
			layers.Replace(layer)
626
627
		case "message":
			messages = append(messages, c.Args)
628
		default:
629
			params[c.Name] = append(params[c.Name], c.Args)
630
631
632
		}
	}

633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
	if len(messages) > 0 {
		fn(api.ProgressResponse{Status: "creating parameters layer"})

		msgs := make([]api.Message, 0)

		for _, m := range messages {
			// todo: handle images
			msg := strings.SplitN(m, ": ", 2)
			msgs = append(msgs, api.Message{Role: msg[0], Content: msg[1]})
		}

		var b bytes.Buffer
		if err := json.NewEncoder(&b).Encode(msgs); err != nil {
			return err
		}

		layer, err := NewLayer(&b, "application/vnd.ollama.image.messages")
		if err != nil {
			return err
		}

		layers.Replace(layer)
	}

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

660
		formattedParams, err := api.FormatParams(params)
661
		if err != nil {
Michael Yang's avatar
Michael Yang committed
662
			return err
663
		}
664

Michael Yang's avatar
Michael Yang committed
665
		for k, v := range fromParams {
Michael Yang's avatar
Michael Yang committed
666
667
668
669
670
			if _, ok := formattedParams[k]; !ok {
				formattedParams[k] = v
			}
		}

Patrick Devine's avatar
Patrick Devine committed
671
		// xxx - can this be removed?
Michael Yang's avatar
Michael Yang committed
672
		if config.ModelType == "65B" {
Michael Yang's avatar
Michael Yang committed
673
			if gqa, ok := formattedParams["gqa"].(int); ok && gqa == 8 {
Michael Yang's avatar
Michael Yang committed
674
675
676
677
				config.ModelType = "70B"
			}
		}

Michael Yang's avatar
Michael Yang committed
678
679
		var b bytes.Buffer
		if err := json.NewEncoder(&b).Encode(formattedParams); err != nil {
680
681
682
			return err
		}

Michael Yang's avatar
Michael Yang committed
683
		fn(api.ProgressResponse{Status: "creating config layer"})
Michael Yang's avatar
Michael Yang committed
684
		layer, err := NewLayer(&b, "application/vnd.ollama.image.params")
685
		if err != nil {
Michael Yang's avatar
Michael Yang committed
686
			return err
687
		}
Michael Yang's avatar
Michael Yang committed
688

Michael Yang's avatar
Michael Yang committed
689
		layers.Replace(layer)
690
691
	}

Michael Yang's avatar
Michael Yang committed
692
693
694
	digests := make([]string, len(layers.items))
	for i, layer := range layers.items {
		digests[i] = layer.Digest
695
696
	}

Michael Yang's avatar
Michael Yang committed
697
	config.RootFS.DiffIDs = digests
Michael Yang's avatar
Michael Yang committed
698

Michael Yang's avatar
Michael Yang committed
699
700
	var b bytes.Buffer
	if err := json.NewEncoder(&b).Encode(config); err != nil {
701
702
703
		return err
	}

Michael Yang's avatar
Michael Yang committed
704
705
	configLayer, err := NewLayer(&b, "application/vnd.docker.container.image.v1+json")
	if err != nil {
706
707
708
		return err
	}

Michael Yang's avatar
Michael Yang committed
709
	delete(deleteMap, configLayer.Digest)
710

Michael Yang's avatar
Michael Yang committed
711
712
	for _, layer := range append(layers.items, configLayer) {
		committed, err := layer.Commit()
713
714
715
		if err != nil {
			return err
		}
716

Michael Yang's avatar
Michael Yang committed
717
718
719
		status := "writing layer"
		if !committed {
			status = "using already created layer"
720
721
		}

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

Michael Yang's avatar
Michael Yang committed
724
		delete(deleteMap, layer.Digest)
725
726
	}

Michael Yang's avatar
Michael Yang committed
727
728
	fn(api.ProgressResponse{Status: "writing manifest"})
	if err := WriteManifest(name, configLayer, layers.items); err != nil {
729
730
		return err
	}
731

Michael Yang's avatar
Michael Yang committed
732
733
734
	if noprune := os.Getenv("OLLAMA_NOPRUNE"); noprune == "" {
		if err := deleteUnusedLayers(nil, deleteMap, false); err != nil {
			return err
735
736
737
		}
	}

Michael Yang's avatar
Michael Yang committed
738
739
	fn(api.ProgressResponse{Status: "success"})
	return nil
740
741
}

Patrick Devine's avatar
Patrick Devine committed
742
func CopyModel(src, dest string) error {
743
	srcModelPath := ParseModelPath(src)
744
	srcPath, err := srcModelPath.GetManifestPath()
745
746
747
748
	if err != nil {
		return err
	}

749
	destModelPath := ParseModelPath(dest)
750
	destPath, err := destModelPath.GetManifestPath()
Patrick Devine's avatar
Patrick Devine committed
751
752
753
	if err != nil {
		return err
	}
754
755
756
	if err := os.MkdirAll(filepath.Dir(destPath), 0o755); err != nil {
		return err
	}
Patrick Devine's avatar
Patrick Devine committed
757
758

	// copy the file
Michael Yang's avatar
Michael Yang committed
759
	input, err := os.ReadFile(srcPath)
Patrick Devine's avatar
Patrick Devine committed
760
761
762
763
764
	if err != nil {
		fmt.Println("Error reading file:", err)
		return err
	}

Michael Yang's avatar
Michael Yang committed
765
	err = os.WriteFile(destPath, input, 0o644)
Patrick Devine's avatar
Patrick Devine committed
766
767
768
769
770
771
772
773
	if err != nil {
		fmt.Println("Error reading file:", err)
		return err
	}

	return nil
}

Michael Yang's avatar
Michael Yang committed
774
func deleteUnusedLayers(skipModelPath *ModelPath, deleteMap map[string]struct{}, dryRun bool) error {
775
776
777
778
	fp, err := GetManifestPath()
	if err != nil {
		return err
	}
Michael Yang's avatar
Michael Yang committed
779
780
781
782

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

Michael Yang's avatar
Michael Yang committed
785
786
787
788
		dir, file := filepath.Split(path)
		dir = strings.Trim(strings.TrimPrefix(dir, fp), string(os.PathSeparator))
		tag := strings.Join([]string{dir, file}, ":")
		fmp := ParseModelPath(tag)
789

Michael Yang's avatar
Michael Yang committed
790
		// skip the manifest we're trying to delete
791
		if skipModelPath != nil && skipModelPath.GetFullTagname() == fmp.GetFullTagname() {
Michael Yang's avatar
Michael Yang committed
792
			return nil
793
		}
Michael Yang's avatar
Michael Yang committed
794
795
796
797

		// save (i.e. delete from the deleteMap) any files used in other manifests
		manifest, _, err := GetManifest(fmp)
		if err != nil {
Michael Yang's avatar
Michael Yang committed
798
			// nolint: nilerr
Michael Yang's avatar
Michael Yang committed
799
800
801
802
803
804
805
806
			return nil
		}

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

		delete(deleteMap, manifest.Config.Digest)
807
		return nil
Michael Yang's avatar
Michael Yang committed
808
809
810
	}

	if err := filepath.Walk(fp, walkFunc); err != nil {
Michael Yang's avatar
Michael Yang committed
811
812
		return err
	}
813
814

	// only delete the files which are still in the deleteMap
Michael Yang's avatar
Michael Yang committed
815
816
817
	for k := range deleteMap {
		fp, err := GetBlobsPath(k)
		if err != nil {
818
			slog.Info(fmt.Sprintf("couldn't get file path for '%s': %v", k, err))
Michael Yang's avatar
Michael Yang committed
819
820
821
822
			continue
		}
		if !dryRun {
			if err := os.Remove(fp); err != nil {
823
				slog.Info(fmt.Sprintf("couldn't remove file '%s': %v", fp, err))
824
825
				continue
			}
Michael Yang's avatar
Michael Yang committed
826
		} else {
827
			slog.Info(fmt.Sprintf("wanted to remove: %s", fp))
828
829
830
		}
	}

831
832
833
834
	return nil
}

func PruneLayers() error {
Michael Yang's avatar
Michael Yang committed
835
	deleteMap := make(map[string]struct{})
836
837
838
839
840
841
842
	p, err := GetBlobsPath("")
	if err != nil {
		return err
	}

	blobs, err := os.ReadDir(p)
	if err != nil {
843
		slog.Info(fmt.Sprintf("couldn't read dir '%s': %v", p, err))
844
845
846
847
848
849
850
851
		return err
	}

	for _, blob := range blobs {
		name := blob.Name()
		if runtime.GOOS == "windows" {
			name = strings.ReplaceAll(name, "-", ":")
		}
Michael Yang's avatar
Michael Yang committed
852
853
854
		if strings.HasPrefix(name, "sha256:") {
			deleteMap[name] = struct{}{}
		}
855
856
	}

857
	slog.Info(fmt.Sprintf("total blobs: %d", len(deleteMap)))
858
859
860
861
862
863

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

864
	slog.Info(fmt.Sprintf("total unused blobs removed: %d", len(deleteMap)))
865
866
867
868

	return nil
}

Michael Yang's avatar
Michael Yang committed
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
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
}

902
903
904
905
906
907
908
func DeleteModel(name string) error {
	mp := ParseModelPath(name)
	manifest, _, err := GetManifest(mp)
	if err != nil {
		return err
	}

Michael Yang's avatar
Michael Yang committed
909
	deleteMap := make(map[string]struct{})
910
	for _, layer := range manifest.Layers {
Michael Yang's avatar
Michael Yang committed
911
		deleteMap[layer.Digest] = struct{}{}
912
	}
Michael Yang's avatar
Michael Yang committed
913
	deleteMap[manifest.Config.Digest] = struct{}{}
914
915
916
917
918
919

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

920
	fp, err := mp.GetManifestPath()
921
922
923
924
925
	if err != nil {
		return err
	}
	err = os.Remove(fp)
	if err != nil {
926
		slog.Info(fmt.Sprintf("couldn't remove manifest file '%s': %v", fp, err))
927
928
929
930
931
932
		return err
	}

	return nil
}

Patrick Devine's avatar
Patrick Devine committed
933
func ShowModelfile(model *Model) (string, error) {
Michael Yang's avatar
Michael Yang committed
934
	var mt struct {
Patrick Devine's avatar
Patrick Devine committed
935
		*Model
Michael Yang's avatar
Michael Yang committed
936
		From       string
Michael Yang's avatar
Michael Yang committed
937
		Parameters map[string][]any
Patrick Devine's avatar
Patrick Devine committed
938
939
	}

Michael Yang's avatar
Michael Yang committed
940
	mt.Parameters = make(map[string][]any)
Patrick Devine's avatar
Patrick Devine committed
941
	for k, v := range model.Options {
Michael Yang's avatar
Michael Yang committed
942
943
944
		if s, ok := v.([]any); ok {
			mt.Parameters[k] = s
			continue
Patrick Devine's avatar
Patrick Devine committed
945
946
		}

Michael Yang's avatar
Michael Yang committed
947
		mt.Parameters[k] = []any{v}
Patrick Devine's avatar
Patrick Devine committed
948
949
	}

Michael Yang's avatar
Michael Yang committed
950
951
	mt.Model = model
	mt.From = model.ModelPath
Patrick Devine's avatar
Patrick Devine committed
952

953
954
	if model.ParentModel != "" {
		mt.From = model.ParentModel
Patrick Devine's avatar
Patrick Devine committed
955
956
957
958
959
960
961
962
	}

	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 }}"""
963
964

{{- if .System }}
Patrick Devine's avatar
Patrick Devine committed
965
SYSTEM """{{ .System }}"""
966
{{- end }}
967
968
969
970

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

Michael Yang's avatar
Michael Yang committed
972
973
974
975
{{- range $k, $v := .Parameters }}
{{- range $parameter := $v }}
PARAMETER {{ $k }} {{ printf "%#v" $parameter }}
{{- end }}
Michael Yang's avatar
Michael Yang committed
976
{{- end }}`
Patrick Devine's avatar
Patrick Devine committed
977
978
979

	tmpl, err := template.New("").Parse(modelFile)
	if err != nil {
980
		slog.Info(fmt.Sprintf("error parsing template: %q", err))
Patrick Devine's avatar
Patrick Devine committed
981
982
983
984
985
986
		return "", err
	}

	var buf bytes.Buffer

	if err = tmpl.Execute(&buf, mt); err != nil {
987
		slog.Info(fmt.Sprintf("error executing template: %q", err))
Patrick Devine's avatar
Patrick Devine committed
988
989
990
991
992
993
		return "", err
	}

	return buf.String(), nil
}

994
func PushModel(ctx context.Context, name string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
995
	mp := ParseModelPath(name)
996
997
	fn(api.ProgressResponse{Status: "retrieving manifest"})

998
999
1000
1001
	if mp.ProtocolScheme == "http" && !regOpts.Insecure {
		return fmt.Errorf("insecure protocol http")
	}

Patrick Devine's avatar
Patrick Devine committed
1002
	manifest, _, err := GetManifest(mp)
1003
	if err != nil {
1004
		fn(api.ProgressResponse{Status: "couldn't retrieve manifest"})
1005
1006
1007
1008
		return err
	}

	var layers []*Layer
Jeffrey Morgan's avatar
Jeffrey Morgan committed
1009
	layers = append(layers, manifest.Layers...)
Michael Yang's avatar
Michael Yang committed
1010
	layers = append(layers, manifest.Config)
1011
1012

	for _, layer := range layers {
Michael Yang's avatar
Michael Yang committed
1013
		if err := uploadBlob(ctx, mp, layer, regOpts, fn); err != nil {
1014
			slog.Info(fmt.Sprintf("error uploading blob: %v", err))
1015
1016
1017
			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())
			}
1018
1019
			return err
		}
1020
1021
	}

1022
	fn(api.ProgressResponse{Status: "pushing manifest"})
Michael Yang's avatar
Michael Yang committed
1023
1024
	requestURL := mp.BaseURL()
	requestURL = requestURL.JoinPath("v2", mp.GetNamespaceRepository(), "manifests", mp.Tag)
1025
1026
1027
1028
1029
1030

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

Michael Yang's avatar
Michael Yang committed
1031
1032
	headers := make(http.Header)
	headers.Set("Content-Type", "application/vnd.docker.distribution.manifest.v2+json")
Michael Yang's avatar
Michael Yang committed
1033
	resp, err := makeRequestWithRetry(ctx, http.MethodPut, requestURL, headers, bytes.NewReader(manifestJSON), regOpts)
1034
1035
1036
1037
1038
	if err != nil {
		return err
	}
	defer resp.Body.Close()

1039
	fn(api.ProgressResponse{Status: "success"})
1040
1041
1042
1043

	return nil
}

1044
func PullModel(ctx context.Context, name string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
1045
1046
	mp := ParseModelPath(name)

1047
1048
1049
1050
1051
	var manifest *ManifestV2
	var err error
	var noprune string

	// build deleteMap to prune unused layers
Michael Yang's avatar
Michael Yang committed
1052
	deleteMap := make(map[string]struct{})
1053
1054
1055
1056
1057
1058
1059
1060
1061

	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
1062
				deleteMap[l.Digest] = struct{}{}
1063
			}
Michael Yang's avatar
Michael Yang committed
1064
			deleteMap[manifest.Config.Digest] = struct{}{}
1065
1066
1067
		}
	}

1068
1069
	if mp.ProtocolScheme == "http" && !regOpts.Insecure {
		return fmt.Errorf("insecure protocol http")
1070
	}
1071

1072
	fn(api.ProgressResponse{Status: "pulling manifest"})
1073

1074
	manifest, err = pullModelManifest(ctx, mp, regOpts)
1075
	if err != nil {
1076
		return fmt.Errorf("pull model manifest: %s", err)
1077
1078
1079
	}

	var layers []*Layer
Bruce MacDonald's avatar
Bruce MacDonald committed
1080
	layers = append(layers, manifest.Layers...)
Michael Yang's avatar
Michael Yang committed
1081
	layers = append(layers, manifest.Config)
1082
1083

	for _, layer := range layers {
1084
1085
1086
1087
1088
1089
1090
1091
		if err := downloadBlob(
			ctx,
			downloadOpts{
				mp:      mp,
				digest:  layer.Digest,
				regOpts: regOpts,
				fn:      fn,
			}); err != nil {
1092
1093
			return err
		}
1094
		delete(deleteMap, layer.Digest)
1095
	}
1096
	delete(deleteMap, manifest.Config.Digest)
1097

Michael Yang's avatar
Michael Yang committed
1098
1099
1100
	fn(api.ProgressResponse{Status: "verifying sha256 digest"})
	for _, layer := range layers {
		if err := verifyBlob(layer.Digest); err != nil {
1101
1102
1103
1104
1105
1106
1107
1108
			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
1109
					slog.Info(fmt.Sprintf("couldn't remove file with digest mismatch '%s': %v", fp, err))
1110
1111
				}
			}
Michael Yang's avatar
Michael Yang committed
1112
1113
1114
1115
			return err
		}
	}

1116
	fn(api.ProgressResponse{Status: "writing manifest"})
1117

1118
	manifestJSON, err := json.Marshal(manifest)
1119
1120
1121
1122
	if err != nil {
		return err
	}

1123
	fp, err := mp.GetManifestPath()
1124
1125
1126
	if err != nil {
		return err
	}
1127
1128
1129
	if err := os.MkdirAll(filepath.Dir(fp), 0o755); err != nil {
		return err
	}
1130

Bruce MacDonald's avatar
Bruce MacDonald committed
1131
	err = os.WriteFile(fp, manifestJSON, 0o644)
1132
	if err != nil {
1133
		slog.Info(fmt.Sprintf("couldn't write to %s", fp))
1134
1135
1136
		return err
	}

1137
1138
1139
1140
1141
1142
1143
1144
	if noprune == "" {
		fn(api.ProgressResponse{Status: "removing any unused layers"})
		err = deleteUnusedLayers(nil, deleteMap, false)
		if err != nil {
			return err
		}
	}

1145
	fn(api.ProgressResponse{Status: "success"})
1146
1147
1148
1149

	return nil
}

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

Michael Yang's avatar
Michael Yang committed
1153
1154
	headers := make(http.Header)
	headers.Set("Accept", "application/vnd.docker.distribution.manifest.v2+json")
Michael Yang's avatar
Michael Yang committed
1155
	resp, err := makeRequestWithRetry(ctx, http.MethodGet, requestURL, headers, nil, regOpts)
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
	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
1170
func GetSHA256Digest(r io.Reader) (string, int64) {
Michael Yang's avatar
Michael Yang committed
1171
1172
1173
1174
1175
1176
	h := sha256.New()
	n, err := io.Copy(h, r)
	if err != nil {
		log.Fatal(err)
	}

Michael Yang's avatar
Michael Yang committed
1177
	return fmt.Sprintf("sha256:%x", h.Sum(nil)), n
1178
1179
}

1180
1181
var errUnauthorized = fmt.Errorf("unauthorized")

Michael Yang's avatar
Michael Yang committed
1182
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
1183
1184
	for i := 0; i < 2; i++ {
		resp, err := makeRequest(ctx, method, requestURL, headers, body, regOpts)
Michael Yang's avatar
Michael Yang committed
1185
		if err != nil {
Michael Yang's avatar
Michael Yang committed
1186
			if !errors.Is(err, context.Canceled) {
1187
				slog.Info(fmt.Sprintf("request failed: %v", err))
Michael Yang's avatar
Michael Yang committed
1188
1189
			}

Michael Yang's avatar
Michael Yang committed
1190
1191
			return nil, err
		}
Michael Yang's avatar
Michael Yang committed
1192
1193
1194
1195
1196
1197
1198

		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
1199
1200
1201
			if err != nil {
				return nil, err
			}
Michael Yang's avatar
Michael Yang committed
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
			regOpts.Token = token
			if body != nil {
				_, err = body.Seek(0, io.SeekStart)
				if err != nil {
					return nil, err
				}
			}
		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)
		default:
			return resp, nil
Michael Yang's avatar
Michael Yang committed
1219
1220
1221
		}
	}

Michael Yang's avatar
Michael Yang committed
1222
	return nil, errUnauthorized
Michael Yang's avatar
Michael Yang committed
1223
1224
}

Michael Yang's avatar
Michael Yang committed
1225
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
1226
	if requestURL.Scheme != "http" && regOpts != nil && regOpts.Insecure {
Michael Yang's avatar
Michael Yang committed
1227
		requestURL.Scheme = "http"
1228
1229
	}

Michael Yang's avatar
Michael Yang committed
1230
	req, err := http.NewRequestWithContext(ctx, method, requestURL.String(), body)
1231
1232
1233
1234
	if err != nil {
		return nil, err
	}

Michael Yang's avatar
Michael Yang committed
1235
1236
1237
1238
	if headers != nil {
		req.Header = headers
	}

Michael Yang's avatar
Michael Yang committed
1239
1240
1241
1242
1243
1244
	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)
		}
1245
1246
	}

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

Michael Yang's avatar
Michael Yang committed
1249
1250
1251
1252
1253
1254
1255
1256
1257
	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
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
1269
	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)
1270
1271
1272
1273
1274
1275
	if err != nil {
		return nil, err
	}

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

Patrick Devine's avatar
Patrick Devine committed
1277
1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
1296
1297
1298
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
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"),
	}
}

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

Michael Yang's avatar
Michael Yang committed
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
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 {
1326
		return fmt.Errorf("%w: want %s, got %s", errDigestMismatch, digest, fileDigest)
Michael Yang's avatar
Michael Yang committed
1327
1328
1329
1330
	}

	return nil
}