llmgw/internal/mirror/ollama.go

package mirror

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"regexp"
	"sort"
	"strings"
)

// Ollama keeps a model as a manifest and the blobs it names. The manifest is
// at manifests/<host>/<namespace>/<model>/<tag> under its models directory —
// registry.ollama.ai/library/llama3.2/3b for "llama3.2:3b" — and each layer
// is a file blobs/sha256-<hex>, named by its own digest. The layer whose media
// type is application/vnd.ollama.image.model is the weights as a GGUF, which
// is what llama.cpp runs; a vision model has its projector in a second GGUF.
//
// What an operator already has in Ollama is weights they chose and have
// tried. Importing them puts that file in the mirror, checked against the
// digest Ollama named it by, rather than downloading it again from somewhere
// that may not be the same file.

// Ollama's defaults for the parts of a name left out.
const (
	ollamaHost      = "registry.ollama.ai"
	ollamaNamespace = "library"
	ollamaTag       = "latest"
)

// The layer media types an import reads, or says it leaves.
const (
	ollamaModelLayer     = "application/vnd.ollama.image.model"
	ollamaProjectorLayer = "application/vnd.ollama.image.projector"
	ollamaParamsLayer    = "application/vnd.ollama.image.params"
	ollamaTensorLayer    = "application/vnd.ollama.image.tensor"
)

// OllamaName is a model as Ollama names it.
type OllamaName struct {
	Host, Namespace, Model, Tag string
}

// ParseOllamaName reads a name as Ollama does: model, namespace/model or
// host/namespace/model, with an optional :tag, and its defaults for whatever
// is left out — "llama3.2" is registry.ollama.ai/library/llama3.2:latest.
func ParseOllamaName(s string) (OllamaName, error) {
	s = strings.TrimSpace(s)
	var n OllamaName
	// A digest pins a name in Ollama; the manifest on disk is by tag.
	s, _, _ = strings.Cut(s, "@")
	// "/" cannot be in a tag, so a colon after the last slash is the tag's.
	if i := strings.LastIndex(s, ":"); i > strings.LastIndex(s, "/") {
		s, n.Tag = s[:i], s[i+1:]
	}
	parts := strings.Split(s, "/")
	switch len(parts) {
	case 1:
		n.Model = parts[0]
	case 2:
		n.Namespace, n.Model = parts[0], parts[1]
	case 3:
		n.Host, n.Namespace, n.Model = parts[0], parts[1], parts[2]
	default:
		return OllamaName{}, fmt.Errorf("%q is not an Ollama model name: model[:tag], namespace/model[:tag] "+
			"or host/namespace/model[:tag]", s)
	}
	if n.Host == "" {
		n.Host = ollamaHost
	}
	if n.Namespace == "" {
		n.Namespace = ollamaNamespace
	}
	if n.Tag == "" {
		n.Tag = ollamaTag
	}
	for _, p := range []string{n.Host, n.Namespace, n.Model, n.Tag} {
		// Each part is a directory of Ollama's; one that is not a single path
		// element would read a manifest from somewhere else.
		if p == "" || p == "." || p == ".." || strings.ContainsAny(p, `\`+"\x00") {
			return OllamaName{}, fmt.Errorf("%q is not an Ollama model name", s)
		}
	}
	return n, nil
}

// String is the name the way Ollama shows it: without the parts that are its
// defaults.
func (n OllamaName) String() string {
	name := n.Model + ":" + n.Tag
	switch {
	case !strings.EqualFold(n.Host, ollamaHost):
		return n.Host + "/" + n.Namespace + "/" + name
	case !strings.EqualFold(n.Namespace, ollamaNamespace):
		return n.Namespace + "/" + name
	}
	return name
}

func (n OllamaName) manifestPath(dir string) string {
	return filepath.Join(dir, "manifests", n.Host, n.Namespace, n.Model, n.Tag)
}

// OllamaDir finds Ollama's models directory: the one given (or its models
// subdirectory, for a ~/.ollama given whole), else OLLAMA_MODELS, else the
// places Ollama keeps it — this user's ~/.ollama/models, and on Linux the
// home of the ollama user its installer makes.
func OllamaDir(given string) (string, error) {
	var candidates []string
	if given != "" {
		candidates = []string{given, filepath.Join(given, "models")}
	} else {
		if env := os.Getenv("OLLAMA_MODELS"); env != "" {
			candidates = append(candidates, env)
		}
		if home, err := os.UserHomeDir(); err == nil {
			candidates = append(candidates, filepath.Join(home, ".ollama", "models"))
		}
		candidates = append(candidates,
			"/usr/share/ollama/.ollama/models", // Ollama's Linux installer
			"/var/lib/ollama/.ollama/models",   // some distributions' packages
		)
	}
	for _, c := range candidates {
		if st, err := os.Stat(filepath.Join(c, "manifests")); err == nil && st.IsDir() {
			return c, nil
		}
	}
	if given != "" {
		return "", fmt.Errorf("%s holds no Ollama models: there is no manifests directory in it or in its models", given)
	}
	return "", fmt.Errorf("no Ollama models found in %s; name the directory with --ollama-dir",
		strings.Join(candidates, ", "))
}

// ListOllama lists the models in an Ollama models directory.
func ListOllama(dir string) ([]OllamaName, error) {
	root := filepath.Join(dir, "manifests")
	var out []OllamaName
	err := filepath.WalkDir(root, func(path string, d os.DirEntry, err error) error {
		if err != nil || d.IsDir() {
			return err
		}
		rel, err := filepath.Rel(root, path)
		if err != nil {
			return err
		}
		parts := strings.Split(rel, string(filepath.Separator))
		if len(parts) != 4 || strings.HasPrefix(parts[3], ".") {
			return nil // not a manifest: something else left in Ollama's directory
		}
		out = append(out, OllamaName{Host: parts[0], Namespace: parts[1], Model: parts[2], Tag: parts[3]})
		return nil
	})
	if err != nil {
		return nil, fmt.Errorf("reading Ollama's models in %s: %w", dir, err)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].String() < out[j].String() })
	return out, nil
}

// OllamaLayer is one layer of an Ollama manifest.
type OllamaLayer struct {
	MediaType string `json:"mediaType"`
	Digest    string `json:"digest"`
	Size      int64  `json:"size"`
}

// ollamaManifest is the manifest's shape: Docker's image manifest, with
// Ollama's media types for the layers.
type ollamaManifest struct {
	SchemaVersion int           `json:"schemaVersion"`
	Layers        []OllamaLayer `json:"layers"`
}

// OllamaModel is what an import takes from one of Ollama's models.
type OllamaModel struct {
	Name OllamaName
	// Model is the weights, a GGUF.
	Model OllamaLayer
	// Projectors are a vision model's image encoders, GGUFs beside it.
	Projectors []OllamaLayer
	// ContextLength is the num_ctx Ollama was told to run it with, 0 when
	// it was not: worth carrying over, since llama.cpp's own default differs.
	ContextLength int
	// Left says what of the Ollama model an import does not carry.
	Left []string
}

var digestPattern = regexp.MustCompile(`^sha256:([0-9a-f]{64})$`)

// ollamaBlob is where a layer's file is, and the sha256 its name promises.
func ollamaBlob(dir, digest string) (path, sum string, err error) {
	m := digestPattern.FindStringSubmatch(digest)
	if m == nil {
		return "", "", fmt.Errorf("the manifest names a layer by %q, which is not a sha256 digest", digest)
	}
	return filepath.Join(dir, "blobs", "sha256-"+m[1]), m[1], nil
}

// ReadOllama reads one model's manifest and the layers an import takes.
func ReadOllama(dir string, name OllamaName) (OllamaModel, error) {
	b, err := os.ReadFile(name.manifestPath(dir)) //nolint:gosec // a manifest under Ollama's directory
	if errors.Is(err, os.ErrNotExist) {
		return OllamaModel{}, fmt.Errorf("no %s among the Ollama models in %s; `llmgw mirror import-ollama` with no name lists them",
			name, dir)
	}
	if err != nil {
		return OllamaModel{}, err
	}
	var mf ollamaManifest
	if err := json.Unmarshal(b, &mf); err != nil {
		return OllamaModel{}, fmt.Errorf("%s's manifest cannot be read: %w", name, err)
	}

	m := OllamaModel{Name: name}
	var models []OllamaLayer
	for _, l := range mf.Layers {
		switch l.MediaType {
		case ollamaModelLayer:
			models = append(models, l)
		case ollamaProjectorLayer:
			m.Projectors = append(m.Projectors, l)
		case ollamaTensorLayer:
			return OllamaModel{}, fmt.Errorf("%s is kept by Ollama as separate tensors rather than a GGUF; "+
				"llama.cpp cannot run that form: mirror the model's GGUF from its publisher instead", name)
		case ollamaParamsLayer:
			m.ContextLength = ollamaContext(dir, l)
		case "application/vnd.ollama.image.template", "application/vnd.ollama.image.prompt":
			m.Left = append(m.Left, "its prompt template: llama.cpp renders the chat template inside the GGUF")
		case "application/vnd.ollama.image.system":
			m.Left = append(m.Left, "its system prompt")
		case "application/vnd.ollama.image.adapter":
			m.Left = append(m.Left, "a LoRA adapter: give it to the model as one of its adapters")
		case "application/vnd.ollama.image.draft":
			m.Left = append(m.Left, "its draft model for speculative decoding")
		}
	}
	switch len(models) {
	case 0:
		return OllamaModel{}, fmt.Errorf("%s has no GGUF layer (%s) to import", name, ollamaModelLayer)
	case 1:
		m.Model = models[0]
	default:
		return OllamaModel{}, fmt.Errorf("%s is made of %d weight layers, which an import does not join yet; "+
			"mirror the model's GGUF from its publisher instead", name, len(models))
	}
	return m, nil
}

// ollamaContext is the num_ctx in a params layer, 0 when there is none or it
// cannot be read: a hint, not worth failing an import over.
func ollamaContext(dir string, l OllamaLayer) int {
	path, _, err := ollamaBlob(dir, l.Digest)
	if err != nil {
		return 0
	}
	b, err := os.ReadFile(path) //nolint:gosec // a blob under Ollama's directory
	if err != nil {
		return 0
	}
	var p struct {
		NumCtx int `json:"num_ctx"`
	}
	_ = json.Unmarshal(b, &p)
	return p.NumCtx
}

// OllamaEntries are where an Ollama model's files go in the mirror, by
// default: the model's own name, and a version that is its tag with the start
// of its digest — a tag such as "latest" moves in Ollama, and a mirror
// version must not.
func OllamaEntries(m OllamaModel, model, version string) (weights Entry, projectors []Entry) {
	if model == "" {
		model = m.Name.Model
	}
	_, sum, _ := ollamaBlob("", m.Model.Digest)
	if version == "" {
		version = m.Name.Tag + "-" + sum[:min(12, len(sum))]
	}
	base := m.Name.Model + "-" + m.Name.Tag
	weights = Entry{Model: model, Version: version, File: base + ".gguf"}
	for i := range m.Projectors {
		// "mmproj" in the name is what llama.cpp's engine knows a projector by,
		// beside the model in one directory.
		file := "mmproj-" + base + ".gguf"
		if i > 0 {
			file = fmt.Sprintf("mmproj-%s-%d.gguf", base, i+1)
		}
		projectors = append(projectors, Entry{Model: model, Version: version, File: file})
	}
	return weights, projectors
}

// ImportOllama copies an Ollama model's GGUF, and its projectors, into one
// version of the mirror, each checked against the digest Ollama named it by.
// progress is told which file it is about.
//
// All or nothing: a version left holding a projector and no model is one a
// node would fetch and fail to start, so the files already copied are taken
// out again when a later one fails.
func ImportOllama(ctx context.Context, root, dir string, m OllamaModel, model, version string,
	progress func(file string, done, total int64),
) ([]Added, error) {
	weights, projectors := OllamaEntries(m, model, version)
	type job struct {
		layer OllamaLayer
		entry Entry
	}
	jobs := []job{{m.Model, weights}}
	for i, l := range m.Projectors {
		jobs = append(jobs, job{l, projectors[i]})
	}
	// Refused before anything is copied, rather than after the first file.
	if _, err := os.Stat(filepath.Join(root, weights.Model, weights.Version)); err == nil {
		return nil, fmt.Errorf("%s/%s: %w", weights.Model, weights.Version, ErrExists)
	}

	source := "ollama://" + m.Name.Host + "/" + m.Name.Namespace + "/" + m.Name.Model + ":" + m.Name.Tag
	var done []Added
	for _, j := range jobs {
		a, err := importLayer(ctx, root, dir, j.layer, j.entry, source, progress)
		if err != nil {
			for _, prev := range done {
				_ = Remove(root, prev.Entry)
			}
			return nil, fmt.Errorf("%s: %w", m.Name, err)
		}
		done = append(done, a)
	}
	return done, nil
}

func importLayer(ctx context.Context, root, dir string, l OllamaLayer, e Entry, source string,
	progress func(file string, done, total int64),
) (Added, error) {
	src, sum, err := ollamaBlob(dir, l.Digest)
	if err != nil {
		return Added{}, err
	}
	if _, err := os.Stat(src); err != nil {
		return Added{}, fmt.Errorf("its layer %s is not in %s: %w", l.Digest, filepath.Join(dir, "blobs"), err)
	}
	var report func(int64, int64)
	if progress != nil {
		report = func(d, t int64) { progress(e.File, d, t) }
	}
	return Import(ctx, root, src, e, sum, source, report)
}