modelpath.go 4.03 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"
Michael Yang's avatar
Michael Yang committed
6
	"net/url"
Patrick Devine's avatar
Patrick Devine committed
7
8
	"os"
	"path/filepath"
9
	"regexp"
Patrick Devine's avatar
Patrick Devine committed
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
	"strings"
)

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

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

28
var (
29
30
31
32
	ErrInvalidImageFormat  = errors.New("invalid image format")
	ErrInvalidProtocol     = errors.New("invalid protocol scheme")
	ErrInsecureProtocol    = errors.New("insecure protocol http")
	ErrInvalidDigestFormat = errors.New("invalid digest format")
33
34
)

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

Michael Yang's avatar
Michael Yang committed
44
45
46
47
	before, after, found := strings.Cut(name, "://")
	if found {
		mp.ProtocolScheme = before
		name = after
48
49
	}

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

64
	if repo, tag, found := strings.Cut(mp.Repository, ":"); found {
65
66
		mp.Repository = repo
		mp.Tag = tag
Patrick Devine's avatar
Patrick Devine committed
67
68
	}

69
	return mp
Patrick Devine's avatar
Patrick Devine committed
70
71
}

72
73
74
75
76
77
78
79
80
81
82
83
84
85
var errModelPathInvalid = errors.New("invalid model path")

func (mp ModelPath) Validate() error {
	if mp.Repository == "" {
		return fmt.Errorf("%w: model repository name is required", errModelPathInvalid)
	}

	if strings.Contains(mp.Tag, ":") {
		return fmt.Errorf("%w: ':' (colon) is not allowed in tag names", errModelPathInvalid)
	}

	return nil
}

Patrick Devine's avatar
Patrick Devine committed
86
87
88
89
90
91
92
93
94
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 {
95
96
97
98
99
	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
100
	}
101
	return fmt.Sprintf("%s/%s/%s:%s", mp.Registry, mp.Namespace, mp.Repository, mp.Tag)
Patrick Devine's avatar
Patrick Devine committed
102
103
}

104
105
106
107
108
109
// modelsDir returns the value of the OLLAMA_MODELS environment variable or the user's home directory if OLLAMA_MODELS is not set.
// The models directory is where Ollama stores its model files and manifests.
func modelsDir() (string, error) {
	if models, exists := os.LookupEnv("OLLAMA_MODELS"); exists {
		return models, nil
	}
Patrick Devine's avatar
Patrick Devine committed
110
111
112
113
	home, err := os.UserHomeDir()
	if err != nil {
		return "", err
	}
114
115
	return filepath.Join(home, ".ollama", "models"), nil
}
Patrick Devine's avatar
Patrick Devine committed
116

117
118
119
120
121
// 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) {
	dir, err := modelsDir()
	if err != nil {
		return "", err
Patrick Devine's avatar
Patrick Devine committed
122
123
	}

124
	return filepath.Join(dir, "manifests", mp.Registry, mp.Namespace, mp.Repository, mp.Tag), nil
Patrick Devine's avatar
Patrick Devine committed
125
126
}

Michael Yang's avatar
Michael Yang committed
127
128
129
130
131
132
133
func (mp ModelPath) BaseURL() *url.URL {
	return &url.URL{
		Scheme: mp.ProtocolScheme,
		Host:   mp.Registry,
	}
}

Patrick Devine's avatar
Patrick Devine committed
134
func GetManifestPath() (string, error) {
135
	dir, err := modelsDir()
Patrick Devine's avatar
Patrick Devine committed
136
137
138
139
	if err != nil {
		return "", err
	}

140
	path := filepath.Join(dir, "manifests")
Michael Yang's avatar
Michael Yang committed
141
	if err := os.MkdirAll(path, 0o755); err != nil {
Michael Yang's avatar
Michael Yang committed
142
143
144
145
		return "", err
	}

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

Patrick Devine's avatar
Patrick Devine committed
148
func GetBlobsPath(digest string) (string, error) {
149
	dir, err := modelsDir()
Patrick Devine's avatar
Patrick Devine committed
150
151
152
153
	if err != nil {
		return "", err
	}

154
155
156
157
158
159
160
161
	// only accept actual sha256 digests
	pattern := "^sha256[:-][0-9a-fA-F]{64}$"
	re := regexp.MustCompile(pattern)

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

162
	digest = strings.ReplaceAll(digest, ":", "-")
163
	path := filepath.Join(dir, "blobs", digest)
164
165
166
167
168
169
	dirPath := filepath.Dir(path)
	if digest == "" {
		dirPath = path
	}

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

Michael Yang's avatar
Michael Yang committed
173
	return path, nil
Patrick Devine's avatar
Patrick Devine committed
174
}