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
				currentVars.Prompt += fmt.Sprintf(" [img-%d]", len(images)+i)
Michael Yang's avatar
Michael Yang committed
186
				currentVars.Images = append(currentVars.Images, llm.ImageData{
Michael Yang's avatar
Michael Yang committed
187
					ID:   len(images) + i,
Michael Yang's avatar
Michael Yang committed
188
					Data: msg.Images[i],
Michael Yang's avatar
Michael Yang committed
189
				})
190
191
			}

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

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

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

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

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

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

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

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

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

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

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

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

	var manifest *ManifestV2

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

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

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

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

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

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

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

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

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

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

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

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

	return model, nil
}

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

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

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

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

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

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

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

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

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

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

				c.Args = blobPath
			}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

Michael Yang's avatar
Michael Yang committed
627
			layers.Replace(layer)
628
629
		case "message":
			messages = append(messages, c.Args)
630
		default:
631
			params[c.Name] = append(params[c.Name], c.Args)
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
657
658
	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
659
	if len(params) > 0 {
Michael Yang's avatar
Michael Yang committed
660
		fn(api.ProgressResponse{Status: "creating parameters layer"})
Michael Yang's avatar
Michael Yang committed
661

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

	return nil
}

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

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

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

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

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

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

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

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

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

833
834
835
836
	return nil
}

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

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

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

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

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

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

	return nil
}

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

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

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

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

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

	return nil
}

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

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

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

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

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

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

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

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

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

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

	var buf bytes.Buffer

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

	return buf.String(), nil
}

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

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

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

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

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

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

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

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

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

	return nil
}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

	return nil
}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

	return nil
}