modelpath.go 3.39 KB
Newer Older
Patrick Devine's avatar
Patrick Devine committed
1
2
3
package server

import (
4
	"errors"
Patrick Devine's avatar
Patrick Devine committed
5
	"fmt"
6
	"io/fs"
Michael Yang's avatar
Michael Yang committed
7
	"net/url"
Patrick Devine's avatar
Patrick Devine committed
8
9
	"os"
	"path/filepath"
10
	"regexp"
Patrick Devine's avatar
Patrick Devine committed
11
	"strings"
12
13

	"github.com/ollama/ollama/envconfig"
14
	"github.com/ollama/ollama/types/model"
Patrick Devine's avatar
Patrick Devine committed
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
)

type ModelPath struct {
	ProtocolScheme string
	Registry       string
	Namespace      string
	Repository     string
	Tag            string
}

const (
	DefaultRegistry       = "registry.ollama.ai"
	DefaultNamespace      = "library"
	DefaultTag            = "latest"
	DefaultProtocolScheme = "https"
)

32
var (
33
34
35
36
	ErrInvalidImageFormat  = errors.New("invalid image format")
	ErrInvalidProtocol     = errors.New("invalid protocol scheme")
	ErrInsecureProtocol    = errors.New("insecure protocol http")
	ErrInvalidDigestFormat = errors.New("invalid digest format")
37
38
)

39
func ParseModelPath(name string) ModelPath {
40
41
42
43
44
45
46
	mp := ModelPath{
		ProtocolScheme: DefaultProtocolScheme,
		Registry:       DefaultRegistry,
		Namespace:      DefaultNamespace,
		Repository:     "",
		Tag:            DefaultTag,
	}
Patrick Devine's avatar
Patrick Devine committed
47

Michael Yang's avatar
Michael Yang committed
48
49
50
51
	before, after, found := strings.Cut(name, "://")
	if found {
		mp.ProtocolScheme = before
		name = after
52
53
	}

54
	name = strings.ReplaceAll(name, string(os.PathSeparator), "/")
55
	parts := strings.Split(name, "/")
56
	switch len(parts) {
Patrick Devine's avatar
Patrick Devine committed
57
	case 3:
58
59
60
		mp.Registry = parts[0]
		mp.Namespace = parts[1]
		mp.Repository = parts[2]
Patrick Devine's avatar
Patrick Devine committed
61
	case 2:
62
63
		mp.Namespace = parts[0]
		mp.Repository = parts[1]
Patrick Devine's avatar
Patrick Devine committed
64
	case 1:
65
		mp.Repository = parts[0]
Patrick Devine's avatar
Patrick Devine committed
66
67
	}

68
	if repo, tag, found := strings.Cut(mp.Repository, ":"); found {
69
70
		mp.Repository = repo
		mp.Tag = tag
Patrick Devine's avatar
Patrick Devine committed
71
72
	}

73
	return mp
Patrick Devine's avatar
Patrick Devine committed
74
75
}

76
77
var errModelPathInvalid = errors.New("invalid model path")

Patrick Devine's avatar
Patrick Devine committed
78
79
80
81
82
83
84
85
86
func (mp ModelPath) GetNamespaceRepository() string {
	return fmt.Sprintf("%s/%s", mp.Namespace, mp.Repository)
}

func (mp ModelPath) GetFullTagname() string {
	return fmt.Sprintf("%s/%s/%s:%s", mp.Registry, mp.Namespace, mp.Repository, mp.Tag)
}

func (mp ModelPath) GetShortTagname() string {
87
88
89
90
91
	if mp.Registry == DefaultRegistry {
		if mp.Namespace == DefaultNamespace {
			return fmt.Sprintf("%s:%s", mp.Repository, mp.Tag)
		}
		return fmt.Sprintf("%s/%s:%s", mp.Namespace, mp.Repository, mp.Tag)
Patrick Devine's avatar
Patrick Devine committed
92
	}
93
	return fmt.Sprintf("%s/%s/%s:%s", mp.Registry, mp.Namespace, mp.Repository, mp.Tag)
Patrick Devine's avatar
Patrick Devine committed
94
95
}

96
97
// GetManifestPath returns the path to the manifest file for the given model path, it is up to the caller to create the directory if it does not exist.
func (mp ModelPath) GetManifestPath() (string, error) {
98
99
100
101
102
	name := model.Name{
		Host:      mp.Registry,
		Namespace: mp.Namespace,
		Model:     mp.Repository,
		Tag:       mp.Tag,
Michael Yang's avatar
Michael Yang committed
103
	}
104
105
106
107
	if !name.IsValid() {
		return "", fs.ErrNotExist
	}
	return filepath.Join(envconfig.Models(), "manifests", name.Filepath()), nil
Patrick Devine's avatar
Patrick Devine committed
108
109
}

Michael Yang's avatar
Michael Yang committed
110
111
112
113
114
115
116
func (mp ModelPath) BaseURL() *url.URL {
	return &url.URL{
		Scheme: mp.ProtocolScheme,
		Host:   mp.Registry,
	}
}

Patrick Devine's avatar
Patrick Devine committed
117
func GetManifestPath() (string, error) {
Michael Yang's avatar
models  
Michael Yang committed
118
	path := filepath.Join(envconfig.Models(), "manifests")
Michael Yang's avatar
Michael Yang committed
119
	if err := os.MkdirAll(path, 0o755); err != nil {
Michael Yang's avatar
Michael Yang committed
120
121
122
123
		return "", err
	}

	return path, nil
Patrick Devine's avatar
Patrick Devine committed
124
125
}

Patrick Devine's avatar
Patrick Devine committed
126
func GetBlobsPath(digest string) (string, error) {
127
128
129
130
131
132
133
134
	// only accept actual sha256 digests
	pattern := "^sha256[:-][0-9a-fA-F]{64}$"
	re := regexp.MustCompile(pattern)

	if digest != "" && !re.MatchString(digest) {
		return "", ErrInvalidDigestFormat
	}

135
	digest = strings.ReplaceAll(digest, ":", "-")
Michael Yang's avatar
models  
Michael Yang committed
136
	path := filepath.Join(envconfig.Models(), "blobs", digest)
137
138
139
140
141
142
	dirPath := filepath.Dir(path)
	if digest == "" {
		dirPath = path
	}

	if err := os.MkdirAll(dirPath, 0o755); err != nil {
Patrick Devine's avatar
Patrick Devine committed
143
144
145
		return "", err
	}

Michael Yang's avatar
Michael Yang committed
146
	return path, nil
Patrick Devine's avatar
Patrick Devine committed
147
}