modelpath.go 3.9 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
	"strings"
11
12

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

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

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

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

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

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

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

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

71
	return mp
Patrick Devine's avatar
Patrick Devine committed
72
73
}

74
75
76
77
78
79
80
81
82
83
84
85
86
87
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
88
89
90
91
92
93
94
95
96
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 {
97
98
99
100
101
	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
102
	}
103
	return fmt.Sprintf("%s/%s/%s:%s", mp.Registry, mp.Namespace, mp.Repository, mp.Tag)
Patrick Devine's avatar
Patrick Devine committed
104
105
}

106
107
108
// 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) {
109
	return envconfig.ModelsDir, nil
110
}
Patrick Devine's avatar
Patrick Devine committed
111

112
113
114
115
116
// 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
117
118
	}

119
	return filepath.Join(dir, "manifests", mp.Registry, mp.Namespace, mp.Repository, mp.Tag), nil
Patrick Devine's avatar
Patrick Devine committed
120
121
}

Michael Yang's avatar
Michael Yang committed
122
123
124
125
126
127
128
func (mp ModelPath) BaseURL() *url.URL {
	return &url.URL{
		Scheme: mp.ProtocolScheme,
		Host:   mp.Registry,
	}
}

Patrick Devine's avatar
Patrick Devine committed
129
func GetManifestPath() (string, error) {
130
	dir, err := modelsDir()
Patrick Devine's avatar
Patrick Devine committed
131
132
133
134
	if err != nil {
		return "", err
	}

135
	path := filepath.Join(dir, "manifests")
Michael Yang's avatar
Michael Yang committed
136
	if err := os.MkdirAll(path, 0o755); err != nil {
Michael Yang's avatar
Michael Yang committed
137
138
139
140
		return "", err
	}

	return path, nil
Patrick Devine's avatar
Patrick Devine committed
141
142
}

Patrick Devine's avatar
Patrick Devine committed
143
func GetBlobsPath(digest string) (string, error) {
144
	dir, err := modelsDir()
Patrick Devine's avatar
Patrick Devine committed
145
146
147
148
	if err != nil {
		return "", err
	}

149
150
151
152
153
154
155
156
	// only accept actual sha256 digests
	pattern := "^sha256[:-][0-9a-fA-F]{64}$"
	re := regexp.MustCompile(pattern)

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

157
	digest = strings.ReplaceAll(digest, ":", "-")
158
	path := filepath.Join(dir, "blobs", digest)
159
160
161
162
163
164
	dirPath := filepath.Dir(path)
	if digest == "" {
		dirPath = path
	}

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

Michael Yang's avatar
Michael Yang committed
168
	return path, nil
Patrick Devine's avatar
Patrick Devine committed
169
}