diff --git a/core/cli/backends.go b/core/cli/backends.go index ca9ed7741..480d68901 100644 --- a/core/cli/backends.go +++ b/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 diff --git a/core/gallery/backends.go b/core/gallery/backends.go index 42812c2fc..f46100057 100644 --- a/core/gallery/backends.go +++ b/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 diff --git a/pkg/downloader/uri.go b/pkg/downloader/uri.go index a4da4f574..b485676ba 100644 --- a/pkg/downloader/uri.go +++ b/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 diff --git a/pkg/startup/backend_preload.go b/pkg/startup/backend_preload.go index b68a811dd..37ca911d8 100644 --- a/pkg/startup/backend_preload.go +++ b/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))