images.go 30.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
	"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
type PromptVars struct {
	System   string
	Prompt   string
	Response string
	First    bool
Michael Yang's avatar
Michael Yang committed
66
	Images   []llm.ImageData
Bruce MacDonald's avatar
Bruce MacDonald committed
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
92
// 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
93
	var prompt strings.Builder
94
	// Use the "missingkey=zero" option to handle missing variables without panicking
95
	tmpl, err := template.New("").Option("missingkey=zero").Parse(promptTemplate)
96
97
98
99
	if err != nil {
		return "", err
	}

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

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

	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
118
119
	return prompt.String(), nil
}
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
// PreResponsePrompt returns the prompt before the response tag
func (m *Model) PreResponsePrompt(p PromptVars) (string, error) {
	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)
}

150
type ChatHistory struct {
Michael Yang's avatar
Michael Yang committed
151
152
	Prompts    []PromptVars
	LastSystem string
153
154
155
156
}

// ChatPrompts returns a list of formatted chat prompts from a list of messages
func (m *Model) ChatPrompts(msgs []api.Message) (*ChatHistory, error) {
Bruce MacDonald's avatar
Bruce MacDonald committed
157
	// build the prompt from the list of messages
158
	lastSystem := m.System
Bruce MacDonald's avatar
Bruce MacDonald committed
159
	currentVars := PromptVars{
160
161
		First:  true,
		System: m.System,
Bruce MacDonald's avatar
Bruce MacDonald committed
162
163
	}

164
	prompts := []PromptVars{}
Michael Yang's avatar
Michael Yang committed
165
	var images []llm.ImageData
Bruce MacDonald's avatar
Bruce MacDonald committed
166
167

	for _, msg := range msgs {
168
		switch strings.ToLower(msg.Role) {
Bruce MacDonald's avatar
Bruce MacDonald committed
169
		case "system":
170
171
			// if this is the first message it overrides the system prompt in the modelfile
			if !currentVars.First && currentVars.System != "" {
172
173
				prompts = append(prompts, currentVars)
				currentVars = PromptVars{}
Bruce MacDonald's avatar
Bruce MacDonald committed
174
175
			}
			currentVars.System = msg.Content
176
			lastSystem = msg.Content
Bruce MacDonald's avatar
Bruce MacDonald committed
177
		case "user":
178
			if currentVars.Prompt != "" {
179
180
				prompts = append(prompts, currentVars)
				currentVars = PromptVars{}
Bruce MacDonald's avatar
Bruce MacDonald committed
181
			}
182

Bruce MacDonald's avatar
Bruce MacDonald committed
183
			currentVars.Prompt = msg.Content
184
			for i := range msg.Images {
Michael Yang's avatar
Michael Yang committed
185
186
				id := len(images) + i
				currentVars.Prompt += fmt.Sprintf(" [img-%d]", id)
Michael Yang's avatar
Michael Yang committed
187
				currentVars.Images = append(currentVars.Images, llm.ImageData{
Michael Yang's avatar
Michael Yang committed
188
					ID:   id,
Michael Yang's avatar
Michael Yang committed
189
					Data: msg.Images[i],
Michael Yang's avatar
Michael Yang committed
190
				})
191
192
			}

Michael Yang's avatar
Michael Yang committed
193
			images = append(images, currentVars.Images...)
Bruce MacDonald's avatar
Bruce MacDonald committed
194
195
		case "assistant":
			currentVars.Response = msg.Content
196
197
			prompts = append(prompts, currentVars)
			currentVars = PromptVars{}
Bruce MacDonald's avatar
Bruce MacDonald committed
198
		default:
199
			return nil, fmt.Errorf("invalid role: %s, role must be one of [system, user, assistant]", msg.Role)
Bruce MacDonald's avatar
Bruce MacDonald committed
200
		}
201
202
	}

Bruce MacDonald's avatar
Bruce MacDonald committed
203
204
	// Append the last set of vars if they are non-empty
	if currentVars.Prompt != "" || currentVars.System != "" {
205
		prompts = append(prompts, currentVars)
Bruce MacDonald's avatar
Bruce MacDonald committed
206
207
	}

208
	return &ChatHistory{
Michael Yang's avatar
Michael Yang committed
209
210
		Prompts:    prompts,
		LastSystem: lastSystem,
211
	}, nil
212
213
}

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

type ConfigV2 struct {
222
223
224
225
226
227
	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"`

228
	// required by spec
229
230
	Architecture string `json:"architecture"`
	OS           string `json:"os"`
231
	RootFS       RootFS `json:"rootfs"`
232
233
}

Michael Yang's avatar
Michael Yang committed
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
261
262
263
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
	}
}

264
265
266
267
268
type RootFS struct {
	Type    string   `json:"type"`
	DiffIDs []string `json:"diff_ids"`
}

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

Patrick Devine's avatar
Patrick Devine committed
274
275
276
277
	total += m.Config.Size
	return total
}

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

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

	var manifest *ManifestV2

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

Patrick Devine's avatar
Patrick Devine committed
295
296
297
	shaSum := sha256.Sum256(bts)
	shaStr := hex.EncodeToString(shaSum[:])

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

Patrick Devine's avatar
Patrick Devine committed
302
	return manifest, shaStr, nil
303
304
305
}

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

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

321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
	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
	}

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

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

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

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

382
			// parse model options parameters into a map so that we can see which fields have been specified explicitly
383
			if err = json.NewDecoder(params).Decode(&model.Options); err != nil {
384
385
				return nil, err
			}
386
387
388
389
390
391
392
393
394
395
		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
396
397
398
399
400
401
		case "application/vnd.ollama.image.license":
			bts, err := os.ReadFile(filename)
			if err != nil {
				return nil, err
			}
			model.License = append(model.License, string(bts))
402
403
404
405
406
407
		}
	}

	return model, nil
}

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

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

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

Michael Yang's avatar
Michael Yang committed
430
431
432
	return abspath
}

433
func CreateModel(ctx context.Context, name, modelFileDir string, commands []parser.Command, fn func(resp api.ProgressResponse)) error {
434
435
436
437
438
439
440
	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{}{}
		}
	}

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

Michael Yang's avatar
Michael Yang committed
449
	var layers Layers
450
	messages := []string{}
Michael Yang's avatar
Michael Yang committed
451

452
	params := make(map[string][]string)
Michael Yang's avatar
Michael Yang committed
453
454
	fromParams := make(map[string]any)

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

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

				c.Args = blobPath
			}

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

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

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

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

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

Bruce MacDonald's avatar
Bruce MacDonald committed
506
507
508
509
510
				// 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
511
512
513
514
				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
515

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

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

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

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

Michael Yang's avatar
Michael Yang committed
540
					layers.Add(layer)
541
				}
Michael Yang's avatar
Michael Yang committed
542
543
544

				deleteMap[manifest.Config.Digest] = struct{}{}
				continue
545
			}
Michael Yang's avatar
Michael Yang committed
546
			defer bin.Close()
547

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

553
554
				bin.Seek(offset, io.SeekStart)
				ggml, err := llm.DecodeGGML(bin)
Bruce MacDonald's avatar
Bruce MacDonald committed
555
556
557
558
559
560
561
562
563
				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
					}
564
				}
Michael Yang's avatar
Michael Yang committed
565

Michael Yang's avatar
Michael Yang committed
566
567
568
569
				config.SetModelFormat(ggml.Name())
				config.SetModelFamily(ggml.ModelFamily())
				config.SetModelType(ggml.ModelType())
				config.SetFileType(ggml.FileType())
570

571
572
573
574
				mediatype := mediatype
				if ggml.ModelFamily() == "clip" {
					mediatype = "application/vnd.ollama.image.projector"
				}
575

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

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

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

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

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

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

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

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

636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
	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
660
	if len(params) > 0 {
Michael Yang's avatar
Michael Yang committed
661
		fn(api.ProgressResponse{Status: "creating parameters layer"})
Michael Yang's avatar
Michael Yang committed
662

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

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

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

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

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

Michael Yang's avatar
Michael Yang committed
692
		layers.Replace(layer)
693
694
	}

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

Michael Yang's avatar
Michael Yang committed
700
	config.RootFS.DiffIDs = digests
Michael Yang's avatar
Michael Yang committed
701

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

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

Michael Yang's avatar
Michael Yang committed
712
	delete(deleteMap, configLayer.Digest)
713

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

Michael Yang's avatar
Michael Yang committed
720
721
722
		status := "writing layer"
		if !committed {
			status = "using already created layer"
723
724
		}

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

Michael Yang's avatar
Michael Yang committed
727
		delete(deleteMap, layer.Digest)
728
729
	}

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

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

Michael Yang's avatar
Michael Yang committed
741
742
	fn(api.ProgressResponse{Status: "success"})
	return nil
743
744
}

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

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

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

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

	return nil
}

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

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

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

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

		// 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
801
			// nolint: nilerr
Michael Yang's avatar
Michael Yang committed
802
803
804
805
806
807
808
809
			return nil
		}

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

		delete(deleteMap, manifest.Config.Digest)
810
		return nil
Michael Yang's avatar
Michael Yang committed
811
812
813
	}

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

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

834
835
836
837
	return nil
}

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

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

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

860
	slog.Info(fmt.Sprintf("total blobs: %d", len(deleteMap)))
861
862
863
864
865
866

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

867
	slog.Info(fmt.Sprintf("total unused blobs removed: %d", len(deleteMap)))
868
869
870
871

	return nil
}

Michael Yang's avatar
Michael Yang committed
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
902
903
904
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
}

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

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

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

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

	return nil
}

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

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

Michael Yang's avatar
Michael Yang committed
950
		mt.Parameters[k] = []any{v}
Patrick Devine's avatar
Patrick Devine committed
951
952
	}

Michael Yang's avatar
Michael Yang committed
953
954
	mt.Model = model
	mt.From = model.ModelPath
Patrick Devine's avatar
Patrick Devine committed
955

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

	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 }}"""
966
967

{{- if .System }}
Patrick Devine's avatar
Patrick Devine committed
968
SYSTEM """{{ .System }}"""
969
{{- end }}
970
971
972
973

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

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

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

	var buf bytes.Buffer

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

	return buf.String(), nil
}

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

1001
1002
1003
1004
	if mp.ProtocolScheme == "http" && !regOpts.Insecure {
		return fmt.Errorf("insecure protocol http")
	}

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

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

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

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

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

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

1042
	fn(api.ProgressResponse{Status: "success"})
1043
1044
1045
1046

	return nil
}

1047
func PullModel(ctx context.Context, name string, regOpts *RegistryOptions, fn func(api.ProgressResponse)) error {
1048
1049
	mp := ParseModelPath(name)

1050
1051
1052
1053
1054
	var manifest *ManifestV2
	var err error
	var noprune string

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

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

1071
1072
	if mp.ProtocolScheme == "http" && !regOpts.Insecure {
		return fmt.Errorf("insecure protocol http")
1073
	}
1074

1075
	fn(api.ProgressResponse{Status: "pulling manifest"})
1076

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

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

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

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

1119
	fn(api.ProgressResponse{Status: "writing manifest"})
1120

1121
	manifestJSON, err := json.Marshal(manifest)
1122
1123
1124
1125
	if err != nil {
		return err
	}

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

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

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

1148
	fn(api.ProgressResponse{Status: "success"})
1149
1150
1151
1152

	return nil
}

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

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

Michael Yang's avatar
Michael Yang committed
1180
	return fmt.Sprintf("sha256:%x", h.Sum(nil)), n
1181
1182
}

1183
1184
var errUnauthorized = fmt.Errorf("unauthorized")

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

Michael Yang's avatar
Michael Yang committed
1193
1194
			return nil, err
		}
Michael Yang's avatar
Michael Yang committed
1195
1196
1197
1198
1199
1200
1201

		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
1202
1203
1204
			if err != nil {
				return nil, err
			}
Michael Yang's avatar
Michael Yang committed
1205
1206
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
			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
1222
1223
1224
		}
	}

Michael Yang's avatar
Michael Yang committed
1225
	return nil, errUnauthorized
Michael Yang's avatar
Michael Yang committed
1226
1227
}

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

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

Michael Yang's avatar
Michael Yang committed
1238
1239
1240
1241
	if headers != nil {
		req.Header = headers
	}

Michael Yang's avatar
Michael Yang committed
1242
1243
1244
1245
1246
1247
	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)
		}
1248
1249
	}

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

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

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

Patrick Devine's avatar
Patrick Devine committed
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
1310
1311
1312
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"),
	}
}

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

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

	return nil
}