modelpath.go 2.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"
Michael Yang's avatar
Michael Yang committed
9
	"runtime"
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
29
30
31
32
33
var (
	ErrInvalidImageFormat = errors.New("invalid image format")
	ErrInvalidProtocol    = errors.New("invalid protocol scheme")
	ErrInsecureProtocol   = errors.New("insecure protocol http")
)

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

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

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

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

67
	return mp
Patrick Devine's avatar
Patrick Devine committed
68
69
70
71
72
73
74
75
76
77
78
}

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

func (mp ModelPath) GetManifestPath(createDir bool) (string, error) {
	home, err := os.UserHomeDir()
	if err != nil {
		return "", err
	}

	path := filepath.Join(home, ".ollama", "models", "manifests", mp.Registry, mp.Namespace, mp.Repository, mp.Tag)
	if createDir {
		if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
			return "", err
		}
	}

	return path, nil
}

Michael Yang's avatar
Michael Yang committed
104
105
106
107
108
109
110
func (mp ModelPath) BaseURL() *url.URL {
	return &url.URL{
		Scheme: mp.ProtocolScheme,
		Host:   mp.Registry,
	}
}

Patrick Devine's avatar
Patrick Devine committed
111
112
113
114
115
116
117
118
119
func GetManifestPath() (string, error) {
	home, err := os.UserHomeDir()
	if err != nil {
		return "", err
	}

	return filepath.Join(home, ".ollama", "models", "manifests"), nil
}

Patrick Devine's avatar
Patrick Devine committed
120
121
122
123
124
125
func GetBlobsPath(digest string) (string, error) {
	home, err := os.UserHomeDir()
	if err != nil {
		return "", err
	}

Michael Yang's avatar
Michael Yang committed
126
127
128
129
	if runtime.GOOS == "windows" {
		digest = strings.ReplaceAll(digest, ":", "-")
	}

Michael Yang's avatar
Michael Yang committed
130
	path := filepath.Join(home, ".ollama", "models", "blobs", digest)
Patrick Devine's avatar
Patrick Devine committed
131
132
133
134
	if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
		return "", err
	}

Michael Yang's avatar
Michael Yang committed
135
	return path, nil
Patrick Devine's avatar
Patrick Devine committed
136
}