feat(oci): support OCI images and Ollama models (#2628)

* Support specifying oci:// and ollama:// for model URLs

Fixes: https://github.com/mudler/LocalAI/issues/2527
Fixes: https://github.com/mudler/LocalAI/issues/1028

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* Lower watcher warnings

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* Allow to install ollama models from CLI

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* fixup tests

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* Do not keep file ownership

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

* Skip test on darwin

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

---------

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

pkg/downloader/uri.go

@@ -11,19 +11,24 @@ import (
 	"strconv"
 	"strings"
 
+	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+
+	"github.com/go-skynet/LocalAI/pkg/oci"
 	"github.com/go-skynet/LocalAI/pkg/utils"
 	"github.com/rs/zerolog/log"
 )
 
 const (
 	HuggingFacePrefix = "huggingface://"
+	OCIPrefix         = "oci://"
+	OllamaPrefix      = "ollama://"
 	HTTPPrefix        = "http://"
 	HTTPSPrefix       = "https://"
 	GithubURI         = "github:"
 	GithubURI2        = "github://"
 )
 
-func GetURI(url string, basePath string, f func(url string, i []byte) error) error {
+func DownloadAndUnmarshal(url string, basePath string, f func(url string, i []byte) error) error {
 	url = ConvertURL(url)
 
 	if strings.HasPrefix(url, "file://") {
@@ -76,9 +81,15 @@ func LooksLikeURL(s string) bool {
 		strings.HasPrefix(s, HTTPSPrefix) ||
 		strings.HasPrefix(s, HuggingFacePrefix) ||
 		strings.HasPrefix(s, GithubURI) ||
+		strings.HasPrefix(s, OllamaPrefix) ||
+		strings.HasPrefix(s, OCIPrefix) ||
 		strings.HasPrefix(s, GithubURI2)
 }
 
+func LooksLikeOCI(s string) bool {
+	return strings.HasPrefix(s, OCIPrefix) || strings.HasPrefix(s, OllamaPrefix)
+}
+
 func ConvertURL(s string) string {
 	switch {
 	case strings.HasPrefix(s, GithubURI2):
@@ -149,6 +160,32 @@ func removePartialFile(tmpFilePath string) error {
 
 func DownloadFile(url string, filePath, sha string, fileN, total int, downloadStatus func(string, string, string, float64)) error {
 	url = ConvertURL(url)
+	if LooksLikeOCI(url) {
+		progressStatus := func(desc ocispec.Descriptor) io.Writer {
+			return &progressWriter{
+				fileName:       filePath,
+				total:          desc.Size,
+				hash:           sha256.New(),
+				fileNo:         fileN,
+				totalFiles:     total,
+				downloadStatus: downloadStatus,
+			}
+		}
+
+		if strings.HasPrefix(url, OllamaPrefix) {
+			url = strings.TrimPrefix(url, OllamaPrefix)
+			return oci.OllamaFetchModel(url, filePath, progressStatus)
+		}
+
+		url = strings.TrimPrefix(url, OCIPrefix)
+		img, err := oci.GetImage(url, "", nil, nil)
+		if err != nil {
+			return fmt.Errorf("failed to get image %q: %v", url, err)
+		}
+
+		return oci.ExtractOCIImage(img, filepath.Dir(filePath))
+	}
+
 	// Check if the file already exists
 	_, err := os.Stat(filePath)
 	if err == nil {

pkg/downloader/uri_test.go

@@ -10,7 +10,7 @@ var _ = Describe("Gallery API tests", func() {
 	Context("URI", func() {
 		It("parses github with a branch", func() {
 			Expect(
-				GetURI("github:go-skynet/model-gallery/gpt4all-j.yaml", "", func(url string, i []byte) error {
+				DownloadAndUnmarshal("github:go-skynet/model-gallery/gpt4all-j.yaml", "", func(url string, i []byte) error {
 					Expect(url).To(Equal("https://raw.githubusercontent.com/go-skynet/model-gallery/main/gpt4all-j.yaml"))
 					return nil
 				}),
@@ -18,7 +18,7 @@ var _ = Describe("Gallery API tests", func() {
 		})
 		It("parses github without a branch", func() {
 			Expect(
-				GetURI("github:go-skynet/model-gallery/gpt4all-j.yaml@main", "", func(url string, i []byte) error {
+				DownloadAndUnmarshal("github:go-skynet/model-gallery/gpt4all-j.yaml@main", "", func(url string, i []byte) error {
 					Expect(url).To(Equal("https://raw.githubusercontent.com/go-skynet/model-gallery/main/gpt4all-j.yaml"))
 					return nil
 				}),
@@ -26,7 +26,7 @@ var _ = Describe("Gallery API tests", func() {
 		})
 		It("parses github with urls", func() {
 			Expect(
-				GetURI("https://raw.githubusercontent.com/go-skynet/model-gallery/main/gpt4all-j.yaml", "", func(url string, i []byte) error {
+				DownloadAndUnmarshal("https://raw.githubusercontent.com/go-skynet/model-gallery/main/gpt4all-j.yaml", "", func(url string, i []byte) error {
 					Expect(url).To(Equal("https://raw.githubusercontent.com/go-skynet/model-gallery/main/gpt4all-j.yaml"))
 					return nil
 				}),

pkg/gallery/gallery.go

@@ -134,7 +134,7 @@ func AvailableGalleryModels(galleries []Gallery, basePath string) ([]*GalleryMod
 
 func findGalleryURLFromReferenceURL(url string, basePath string) (string, error) {
 	var refFile string
-	err := downloader.GetURI(url, basePath, func(url string, d []byte) error {
+	err := downloader.DownloadAndUnmarshal(url, basePath, func(url string, d []byte) error {
 		refFile = string(d)
 		if len(refFile) == 0 {
 			return fmt.Errorf("invalid reference file at url %s: %s", url, d)
@@ -157,7 +157,7 @@ func getGalleryModels(gallery Gallery, basePath string) ([]*GalleryModel, error)
 		}
 	}
 
-	err := downloader.GetURI(gallery.URL, basePath, func(url string, d []byte) error {
+	err := downloader.DownloadAndUnmarshal(gallery.URL, basePath, func(url string, d []byte) error {
 		return yaml.Unmarshal(d, &models)
 	})
 	if err != nil {

pkg/gallery/models.go

@@ -65,7 +65,7 @@ type PromptTemplate struct {
 
 func GetGalleryConfigFromURL(url string, basePath string) (Config, error) {
 	var config Config
-	err := downloader.GetURI(url, basePath, func(url string, d []byte) error {
+	err := downloader.DownloadAndUnmarshal(url, basePath, func(url string, d []byte) error {
 		return yaml.Unmarshal(d, &config)
 	})
 	if err != nil {

pkg/oci/blob.go

@@ -0,0 +1,52 @@
+package oci
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"os"
+
+	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+
+	oras "oras.land/oras-go/v2"
+	"oras.land/oras-go/v2/registry/remote"
+)
+
+func FetchImageBlob(r, reference, dst string, statusReader func(ocispec.Descriptor) io.Writer) error {
+	// 0. Create a file store for the output
+	fs, err := os.Create(dst)
+	if err != nil {
+		return err
+	}
+	defer fs.Close()
+
+	// 1. Connect to a remote repository
+	ctx := context.Background()
+	repo, err := remote.NewRepository(r)
+	if err != nil {
+		return fmt.Errorf("failed to create repository: %v", err)
+	}
+	repo.SkipReferrersGC = true
+
+	// https://github.com/oras-project/oras/blob/main/cmd/oras/internal/option/remote.go#L364
+	// https://github.com/oras-project/oras/blob/main/cmd/oras/root/blob/fetch.go#L136
+	desc, reader, err := oras.Fetch(ctx, repo.Blobs(), reference, oras.DefaultFetchOptions)
+	if err != nil {
+		return fmt.Errorf("failed to fetch image: %v", err)
+	}
+
+	if statusReader != nil {
+		// 3. Write the file to the file store
+		_, err = io.Copy(io.MultiWriter(fs, statusReader(desc)), reader)
+		if err != nil {
+			return err
+		}
+	} else {
+		_, err = io.Copy(fs, reader)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}

pkg/oci/blob_test.go

@@ -0,0 +1,21 @@
+package oci_test
+
+import (
+	"os"
+
+	. "github.com/go-skynet/LocalAI/pkg/oci" // Update with your module path
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("OCI", func() {
+	Context("pulling images", func() {
+		It("should fetch blobs correctly", func() {
+			f, err := os.CreateTemp("", "ollama")
+			Expect(err).NotTo(HaveOccurred())
+			defer os.RemoveAll(f.Name())
+			err = FetchImageBlob("registry.ollama.ai/library/gemma", "sha256:c1864a5eb19305c40519da12cc543519e48a0697ecd30e15d5ac228644957d12", f.Name(), nil)
+			Expect(err).NotTo(HaveOccurred())
+		})
+	})
+})

pkg/oci/image.go

@@ -0,0 +1,153 @@
+package oci
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"io"
+	"net/http"
+	"runtime"
+	"strings"
+	"syscall"
+	"time"
+
+	"github.com/containerd/containerd/archive"
+	registrytypes "github.com/docker/docker/api/types/registry"
+	"github.com/google/go-containerregistry/pkg/authn"
+	"github.com/google/go-containerregistry/pkg/logs"
+	"github.com/google/go-containerregistry/pkg/name"
+	v1 "github.com/google/go-containerregistry/pkg/v1"
+	"github.com/google/go-containerregistry/pkg/v1/mutate"
+	"github.com/google/go-containerregistry/pkg/v1/remote"
+	"github.com/google/go-containerregistry/pkg/v1/remote/transport"
+)
+
+// ref: https://github.com/mudler/luet/blob/master/pkg/helpers/docker/docker.go#L117
+type staticAuth struct {
+	auth *registrytypes.AuthConfig
+}
+
+func (s staticAuth) Authorization() (*authn.AuthConfig, error) {
+	if s.auth == nil {
+		return nil, nil
+	}
+	return &authn.AuthConfig{
+		Username:      s.auth.Username,
+		Password:      s.auth.Password,
+		Auth:          s.auth.Auth,
+		IdentityToken: s.auth.IdentityToken,
+		RegistryToken: s.auth.RegistryToken,
+	}, nil
+}
+
+var defaultRetryBackoff = remote.Backoff{
+	Duration: 1.0 * time.Second,
+	Factor:   3.0,
+	Jitter:   0.1,
+	Steps:    3,
+}
+
+var defaultRetryPredicate = func(err error) bool {
+	if err == nil {
+		return false
+	}
+
+	if errors.Is(err, io.ErrUnexpectedEOF) || errors.Is(err, io.EOF) || errors.Is(err, syscall.EPIPE) || errors.Is(err, syscall.ECONNRESET) || strings.Contains(err.Error(), "connection refused") {
+		logs.Warn.Printf("retrying %v", err)
+		return true
+	}
+	return false
+}
+
+// ExtractOCIImage will extract a given targetImage into a given targetDestination
+func ExtractOCIImage(img v1.Image, targetDestination string) error {
+	reader := mutate.Extract(img)
+
+	_, err := archive.Apply(context.Background(), targetDestination, reader, archive.WithNoSameOwner())
+
+	return err
+}
+
+func ParseImageParts(image string) (tag, repository, dstimage string) {
+	tag = "latest"
+	repository = "library"
+	if strings.Contains(image, ":") {
+		parts := strings.Split(image, ":")
+		image = parts[0]
+		tag = parts[1]
+	}
+	if strings.Contains("/", image) {
+		parts := strings.Split(image, "/")
+		repository = parts[0]
+		image = parts[1]
+	}
+	dstimage = image
+	return tag, repository, image
+}
+
+// GetImage if returns the proper image to pull with transport and auth
+// tries local daemon first and then fallbacks into remote
+// if auth is nil, it will try to use the default keychain https://github.com/google/go-containerregistry/tree/main/pkg/authn#tldr-for-consumers-of-this-package
+func GetImage(targetImage, targetPlatform string, auth *registrytypes.AuthConfig, t http.RoundTripper) (v1.Image, error) {
+	var platform *v1.Platform
+	var image v1.Image
+	var err error
+
+	if targetPlatform != "" {
+		platform, err = v1.ParsePlatform(targetPlatform)
+		if err != nil {
+			return image, err
+		}
+	} else {
+		platform, err = v1.ParsePlatform(fmt.Sprintf("%s/%s", runtime.GOOS, runtime.GOARCH))
+		if err != nil {
+			return image, err
+		}
+	}
+
+	ref, err := name.ParseReference(targetImage)
+	if err != nil {
+		return image, err
+	}
+
+	if t == nil {
+		t = http.DefaultTransport
+	}
+
+	tr := transport.NewRetry(t,
+		transport.WithRetryBackoff(defaultRetryBackoff),
+		transport.WithRetryPredicate(defaultRetryPredicate),
+	)
+
+	opts := []remote.Option{
+		remote.WithTransport(tr),
+		remote.WithPlatform(*platform),
+	}
+	if auth != nil {
+		opts = append(opts, remote.WithAuth(staticAuth{auth}))
+	} else {
+		opts = append(opts, remote.WithAuthFromKeychain(authn.DefaultKeychain))
+	}
+
+	image, err = remote.Image(ref, opts...)
+
+	return image, err
+}
+
+func GetOCIImageSize(targetImage, targetPlatform string, auth *registrytypes.AuthConfig, t http.RoundTripper) (int64, error) {
+	var size int64
+	var img v1.Image
+	var err error
+
+	img, err = GetImage(targetImage, targetPlatform, auth, t)
+	if err != nil {
+		return size, err
+	}
+	layers, _ := img.Layers()
+	for _, layer := range layers {
+		s, _ := layer.Size()
+		size += s
+	}
+
+	return size, nil
+}

pkg/oci/image_test.go

@@ -0,0 +1,37 @@
+package oci_test
+
+import (
+	"os"
+	"runtime"
+
+	. "github.com/go-skynet/LocalAI/pkg/oci" // Update with your module path
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("OCI", func() {
+
+	Context("when template is loaded successfully", func() {
+		It("should evaluate the template correctly", func() {
+			if runtime.GOOS == "darwin" {
+				Skip("Skipping test on darwin")
+			}
+			imageName := "alpine"
+			img, err := GetImage(imageName, "", nil, nil)
+			Expect(err).NotTo(HaveOccurred())
+
+			size, err := GetOCIImageSize(imageName, "", nil, nil)
+			Expect(err).NotTo(HaveOccurred())
+
+			Expect(size).ToNot(Equal(int64(0)))
+
+			// Create tempdir
+			dir, err := os.MkdirTemp("", "example")
+			Expect(err).NotTo(HaveOccurred())
+			defer os.RemoveAll(dir)
+
+			err = ExtractOCIImage(img, dir)
+			Expect(err).NotTo(HaveOccurred())
+		})
+	})
+})

pkg/oci/oci_suite_test.go

@@ -0,0 +1,13 @@
+package oci_test
+
+import (
+	"testing"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+func TestOCI(t *testing.T) {
+	RegisterFailHandler(Fail)
+	RunSpecs(t, "OCI test suite")
+}

pkg/oci/ollama.go

@@ -0,0 +1,88 @@
+package oci
+
+import (
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+
+	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
+)
+
+// Define the main struct for the JSON data
+type Manifest struct {
+	SchemaVersion int           `json:"schemaVersion"`
+	MediaType     string        `json:"mediaType"`
+	Config        Config        `json:"config"`
+	Layers        []LayerDetail `json:"layers"`
+}
+
+// Define the struct for the "config" section
+type Config struct {
+	Digest    string `json:"digest"`
+	MediaType string `json:"mediaType"`
+	Size      int    `json:"size"`
+}
+
+// Define the struct for each item in the "layers" array
+type LayerDetail struct {
+	Digest    string `json:"digest"`
+	MediaType string `json:"mediaType"`
+	Size      int    `json:"size"`
+}
+
+func OllamaModelManifest(image string) (*Manifest, error) {
+	// parse the repository and tag from `image`. `image` should be for e.g. gemma:2b, or foobar/gemma:2b
+
+	// if there is a : in the image, then split it
+	// if there is no : in the image, then assume it is the latest tag
+	tag, repository, image := ParseImageParts(image)
+
+	// get e.g. https://registry.ollama.ai/v2/library/llama3/manifests/latest
+	req, err := http.NewRequest("GET", "https://registry.ollama.ai/v2/"+repository+"/"+image+"/manifests/"+tag, nil)
+	if err != nil {
+		return nil, err
+	}
+	req.Header.Set("Accept", "application/vnd.docker.distribution.manifest.v2+json")
+	client := &http.Client{}
+	resp, err := client.Do(req)
+	if err != nil {
+		return nil, err
+	}
+
+	// parse the JSON response
+	var manifest Manifest
+	err = json.NewDecoder(resp.Body).Decode(&manifest)
+	if err != nil {
+		return nil, err
+	}
+
+	return &manifest, nil
+}
+
+func OllamaModelBlob(image string) (string, error) {
+	manifest, err := OllamaModelManifest(image)
+	if err != nil {
+		return "", err
+	}
+	// find a application/vnd.ollama.image.model in the mediaType
+
+	for _, layer := range manifest.Layers {
+		if layer.MediaType == "application/vnd.ollama.image.model" {
+			return layer.Digest, nil
+		}
+	}
+
+	return "", nil
+}
+
+func OllamaFetchModel(image string, output string, statusWriter func(ocispec.Descriptor) io.Writer) error {
+	_, repository, imageNoTag := ParseImageParts(image)
+
+	blobID, err := OllamaModelBlob(image)
+	if err != nil {
+		return err
+	}
+
+	return FetchImageBlob(fmt.Sprintf("registry.ollama.ai/%s/%s", repository, imageNoTag), blobID, output, statusWriter)
+}

pkg/oci/ollama_test.go

@@ -0,0 +1,21 @@
+package oci_test
+
+import (
+	"os"
+
+	. "github.com/go-skynet/LocalAI/pkg/oci" // Update with your module path
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("OCI", func() {
+	Context("ollama", func() {
+		It("pulls model files", func() {
+			f, err := os.CreateTemp("", "ollama")
+			Expect(err).NotTo(HaveOccurred())
+			defer os.RemoveAll(f.Name())
+			err = OllamaFetchModel("gemma:2b", f.Name(), nil)
+			Expect(err).NotTo(HaveOccurred())
+		})
+	})
+})

pkg/startup/model_preload.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"os"
 	"path/filepath"
+	"strings"
 
 	"github.com/go-skynet/LocalAI/embedded"
 	"github.com/go-skynet/LocalAI/pkg/downloader"
@@ -52,6 +53,28 @@ func InstallModels(galleries []gallery.Gallery, modelLibraryURL string, modelPat
 				log.Error().Err(e).Str("filepath", modelDefinitionFilePath).Msg("error writing model definition")
 				err = errors.Join(err, e)
 			}
+		case downloader.LooksLikeOCI(url):
+			log.Debug().Msgf("[startup] resolved OCI model to download: %s", url)
+
+			// convert OCI image name to a file name.
+			ociName := strings.TrimPrefix(url, downloader.OCIPrefix)
+			ociName = strings.TrimPrefix(ociName, downloader.OllamaPrefix)
+			ociName = strings.ReplaceAll(ociName, "/", "__")
+			ociName = strings.ReplaceAll(ociName, ":", "__")
+
+			// check if file exists
+			if _, e := os.Stat(filepath.Join(modelPath, ociName)); errors.Is(e, os.ErrNotExist) {
+				modelDefinitionFilePath := filepath.Join(modelPath, ociName)
+				e := downloader.DownloadFile(url, modelDefinitionFilePath, "", 0, 0, func(fileName, current, total string, percent float64) {
+					utils.DisplayDownloadFunction(fileName, current, total, percent)
+				})
+				if e != nil {
+					log.Error().Err(e).Str("url", url).Str("filepath", modelDefinitionFilePath).Msg("error downloading model")
+					err = errors.Join(err, e)
+				}
+			}
+
+			log.Info().Msgf("[startup] installed model from OCI repository: %s", ociName)
 		case downloader.LooksLikeURL(url):
 			log.Debug().Msgf("[startup] resolved model to download: %s", url)