feat(cli): allow to install backends from OCI tar files (#5816)

Signed-off-by: Ettore Di Giacinto <[email protected]>

Author: mudler

core/cli/backends.go

@@ -8,7 +8,6 @@ import (
 	"github.com/mudler/LocalAI/core/config"
 
 	"github.com/mudler/LocalAI/core/gallery"
-	"github.com/mudler/LocalAI/pkg/downloader"
 	"github.com/mudler/LocalAI/pkg/startup"
 	"github.com/rs/zerolog/log"
 	"github.com/schollz/progressbar/v3"
@@ -23,12 +22,6 @@ type BackendsList struct {
 	BackendsCMDFlags `embed:""`
 }
 
-type BackendsInstallSingle struct {
-	InstallArgs []string `arg:"" optional:"" name:"backend" help:"Backend images to install"`
-
-	BackendsCMDFlags `embed:""`
-}
-
 type BackendsInstall struct {
 	BackendArgs []string `arg:"" optional:"" name:"backends" help:"Backend configuration URLs to load"`
 
@@ -42,36 +35,9 @@ type BackendsUninstall struct {
 }
 
 type BackendsCMD struct {
-	List          BackendsList          `cmd:"" help:"List the backends available in your galleries" default:"withargs"`
-	Install       BackendsInstall       `cmd:"" help:"Install a backend from the gallery"`
-	InstallSingle BackendsInstallSingle `cmd:"" help:"Install a single backend from the gallery"`
-	Uninstall     BackendsUninstall     `cmd:"" help:"Uninstall a backend"`
-}
-
-func (bi *BackendsInstallSingle) Run(ctx *cliContext.Context) error {
-	for _, backend := range bi.InstallArgs {
-		progressBar := progressbar.NewOptions(
-			1000,
-			progressbar.OptionSetDescription(fmt.Sprintf("downloading backend %s", backend)),
-			progressbar.OptionShowBytes(false),
-			progressbar.OptionClearOnFinish(),
-		)
-		progressCallback := func(fileName string, current string, total string, percentage float64) {
-			v := int(percentage * 10)
-			err := progressBar.Set(v)
-			if err != nil {
-				log.Error().Err(err).Str("filename", fileName).Int("value", v).Msg("error while updating progress bar")
-			}
-		}
-
-		if err := gallery.InstallBackend(bi.BackendsPath, &gallery.GalleryBackend{
-			URI: backend,
-		}, progressCallback); err != nil {
-			return err
-		}
-	}
-
-	return nil
+	List      BackendsList      `cmd:"" help:"List the backends available in your galleries" default:"withargs"`
+	Install   BackendsInstall   `cmd:"" help:"Install a backend from the gallery"`
+	Uninstall BackendsUninstall `cmd:"" help:"Uninstall a backend"`
 }
 
 func (bl *BackendsList) Run(ctx *cliContext.Context) error {
@@ -116,23 +82,6 @@ func (bi *BackendsInstall) Run(ctx *cliContext.Context) error {
 			}
 		}
 
-		backendURI := downloader.URI(backendName)
-
-		if !backendURI.LooksLikeOCI() {
-			backends, err := gallery.AvailableBackends(galleries, bi.BackendsPath)
-			if err != nil {
-				return err
-			}
-
-			backend := gallery.FindGalleryElement(backends, backendName, bi.BackendsPath)
-			if backend == nil {
-				log.Error().Str("backend", backendName).Msg("backend not found")
-				return fmt.Errorf("backend not found: %s", backendName)
-			}
-
-			log.Info().Str("backend", backendName).Str("license", backend.License).Msg("installing backend")
-		}
-
 		err := startup.InstallExternalBackends(galleries, bi.BackendsPath, progressCallback, backendName)
 		if err != nil {
 			return err

core/gallery/backends.go

@@ -9,8 +9,8 @@ import (
 
 	"github.com/mudler/LocalAI/core/config"
 	"github.com/mudler/LocalAI/core/system"
+	"github.com/mudler/LocalAI/pkg/downloader"
 	"github.com/mudler/LocalAI/pkg/model"
-	"github.com/mudler/LocalAI/pkg/oci"
 	"github.com/rs/zerolog/log"
 )
 
@@ -151,19 +151,15 @@ func InstallBackend(basePath string, config *GalleryBackend, downloadStatus func
 	}
 
 	name := config.Name
-
-	img, err := oci.GetImage(config.URI, "", nil, nil)
-	if err != nil {
-		return fmt.Errorf("failed to get image %q: %v", config.URI, err)
-	}
-
 	backendPath := filepath.Join(basePath, name)
-	if err := os.MkdirAll(backendPath, 0750); err != nil {
-		return fmt.Errorf("failed to create backend path %q: %v", backendPath, err)
+	err = os.MkdirAll(backendPath, 0750)
+	if err != nil {
+		return fmt.Errorf("failed to create base path: %v", err)
 	}
 
-	if err := oci.ExtractOCIImage(img, config.URI, backendPath, downloadStatus); err != nil {
-		return fmt.Errorf("failed to extract image %q: %v", config.URI, err)
+	uri := downloader.URI(config.URI)
+	if err := uri.DownloadFile(backendPath, "", 1, 1, downloadStatus); err != nil {
+		return fmt.Errorf("failed to download backend %q: %v", config.URI, err)
 	}
 
 	// Create metadata for the backend

pkg/downloader/uri.go

@@ -13,6 +13,7 @@ import (
 	"strconv"
 	"strings"
 
+	"github.com/google/go-containerregistry/pkg/v1/tarball"
 	ocispec "github.com/opencontainers/image-spec/specs-go/v1"
 
 	"github.com/mudler/LocalAI/pkg/oci"
@@ -25,6 +26,7 @@ const (
 	HuggingFacePrefix1 = "hf://"
 	HuggingFacePrefix2 = "hf.co/"
 	OCIPrefix          = "oci://"
+	OCIFilePrefix      = "ocifile://"
 	OllamaPrefix       = "ollama://"
 	HTTPPrefix         = "http://"
 	HTTPSPrefix        = "https://"
@@ -137,8 +139,18 @@ func (u URI) LooksLikeURL() bool {
 		strings.HasPrefix(string(u), GithubURI2)
 }
 
+func (u URI) LooksLikeHTTPURL() bool {
+	return strings.HasPrefix(string(u), HTTPPrefix) ||
+		strings.HasPrefix(string(u), HTTPSPrefix)
+}
+
 func (s URI) LooksLikeOCI() bool {
-	return strings.HasPrefix(string(s), OCIPrefix) || strings.HasPrefix(string(s), OllamaPrefix)
+	return strings.HasPrefix(string(s), "quay.io") ||
+		strings.HasPrefix(string(s), OCIPrefix) ||
+		strings.HasPrefix(string(s), OllamaPrefix) ||
+		strings.HasPrefix(string(s), OCIFilePrefix) ||
+		strings.HasPrefix(string(s), "ghcr.io") ||
+		strings.HasPrefix(string(s), "docker.io")
 }
 
 func (s URI) ResolveURL() string {
@@ -234,6 +246,13 @@ func (uri URI) checkSeverSupportsRangeHeader() (bool, error) {
 func (uri URI) DownloadFile(filePath, sha string, fileN, total int, downloadStatus func(string, string, string, float64)) error {
 	url := uri.ResolveURL()
 	if uri.LooksLikeOCI() {
+
+		// Only Ollama wants to download to the file, for the rest, we want to download to the directory
+		// so we check if filepath has any extension, otherwise we assume it's a directory
+		if filepath.Ext(filePath) != "" && !strings.HasPrefix(url, OllamaPrefix) {
+			filePath = filepath.Dir(filePath)
+		}
+
 		progressStatus := func(desc ocispec.Descriptor) io.Writer {
 			return &progressWriter{
 				fileName:       filePath,
@@ -245,18 +264,32 @@ func (uri URI) DownloadFile(filePath, sha string, fileN, total int, downloadStat
 			}
 		}
 
-		if strings.HasPrefix(url, OllamaPrefix) {
-			url = strings.TrimPrefix(url, OllamaPrefix)
+		if url, ok := strings.CutPrefix(url, OllamaPrefix); ok {
 			return oci.OllamaFetchModel(url, filePath, progressStatus)
 		}
 
+		if url, ok := strings.CutPrefix(url, OCIFilePrefix); ok {
+			// Open the tarball
+			img, err := tarball.ImageFromPath(url, nil)
+			if err != nil {
+				return fmt.Errorf("failed to open tarball: %s", err.Error())
+			}
+
+			return oci.ExtractOCIImage(img, url, filePath, downloadStatus)
+		}
+
 		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, url, filepath.Dir(filePath), downloadStatus)
+		return oci.ExtractOCIImage(img, url, filePath, downloadStatus)
+	}
+
+	// We need to check if url looks like an URL or bail out
+	if !URI(url).LooksLikeHTTPURL() {
+		return fmt.Errorf("url %q does not look like an HTTP URL", url)
 	}
 
 	// Check if the file already exists

pkg/startup/backend_preload.go

@@ -3,11 +3,14 @@ package startup
 import (
 	"errors"
 	"fmt"
+	"path/filepath"
 	"strings"
 
 	"github.com/mudler/LocalAI/core/config"
 	"github.com/mudler/LocalAI/core/gallery"
 	"github.com/mudler/LocalAI/core/system"
+	"github.com/mudler/LocalAI/pkg/downloader"
+	"github.com/rs/zerolog/log"
 )
 
 func InstallExternalBackends(galleries []config.Gallery, backendPath string, downloadStatus func(string, string, string, float64), backends ...string) error {
@@ -17,11 +20,21 @@ func InstallExternalBackends(galleries []config.Gallery, backendPath string, dow
 		return fmt.Errorf("failed to get system state: %w", err)
 	}
 	for _, backend := range backends {
+		uri := downloader.URI(backend)
 		switch {
-		case strings.HasPrefix(backend, "oci://"):
-			backend = strings.TrimPrefix(backend, "oci://")
+		case uri.LooksLikeOCI():
+			name, err := uri.FilenameFromUrl()
+			if err != nil {
+				return fmt.Errorf("failed to get filename from URL: %w", err)
+			}
+			// strip extension if any
+			name = strings.TrimSuffix(name, filepath.Ext(name))
 
+			log.Info().Str("backend", backend).Str("name", name).Msg("Installing backend from OCI image")
 			if err := gallery.InstallBackend(backendPath, &gallery.GalleryBackend{
+				Metadata: gallery.Metadata{
+					Name: name,
+				},
 				URI: backend,
 			}, downloadStatus); err != nil {
 				errs = errors.Join(err, fmt.Errorf("error installing backend %s", backend))