LocalAI/core/gallery/gallery.go

package gallery

import (
	"context"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"
	"time"

	"github.com/lithammer/fuzzysearch/fuzzy"
	"github.com/mudler/LocalAI/core/config"
	"github.com/mudler/LocalAI/pkg/downloader"
	"github.com/mudler/LocalAI/pkg/system"
	"github.com/mudler/LocalAI/pkg/xsync"
	"github.com/mudler/xlog"

	"gopkg.in/yaml.v2"
)

func GetGalleryConfigFromURL[T any](url string, basePath string) (T, error) {
	var config T
	uri := downloader.URI(url)
	err := uri.ReadWithCallback(basePath, func(url string, d []byte) error {
		return yaml.Unmarshal(d, &config)
	})
	if err != nil {
		xlog.Error("failed to get gallery config for url", "error", err, "url", url)
		return config, err
	}
	return config, nil
}

func GetGalleryConfigFromURLWithContext[T any](ctx context.Context, url string, basePath string) (T, error) {
	var config T
	uri := downloader.URI(url)
	err := uri.ReadWithAuthorizationAndCallback(ctx, basePath, "", func(url string, d []byte) error {
		return yaml.Unmarshal(d, &config)
	})
	if err != nil {
		xlog.Error("failed to get gallery config for url", "error", err, "url", url)
		return config, err
	}
	return config, nil
}

func ReadConfigFile[T any](filePath string) (*T, error) {
	// Read the YAML file
	yamlFile, err := os.ReadFile(filePath)
	if err != nil {
		return nil, fmt.Errorf("failed to read YAML file: %v", err)
	}

	// Unmarshal YAML data into a Config struct
	var config T
	err = yaml.Unmarshal(yamlFile, &config)
	if err != nil {
		return nil, fmt.Errorf("failed to unmarshal YAML: %v", err)
	}

	return &config, nil
}

type GalleryElement interface {
	SetGallery(gallery config.Gallery)
	SetInstalled(installed bool)
	GetName() string
	GetDescription() string
	GetTags() []string
	GetInstalled() bool
	GetLicense() string
	GetGallery() config.Gallery
}

type GalleryElements[T GalleryElement] []T

func (gm GalleryElements[T]) Search(term string) GalleryElements[T] {
	var filteredModels GalleryElements[T]
	term = strings.ToLower(term)
	for _, m := range gm {
		if fuzzy.Match(term, strings.ToLower(m.GetName())) ||
			fuzzy.Match(term, strings.ToLower(m.GetGallery().Name)) ||
			strings.Contains(strings.ToLower(m.GetName()), term) ||
			strings.Contains(strings.ToLower(m.GetDescription()), term) ||
			strings.Contains(strings.ToLower(m.GetGallery().Name), term) ||
			strings.Contains(strings.ToLower(strings.Join(m.GetTags(), ",")), term) {
			filteredModels = append(filteredModels, m)
		}
	}

	return filteredModels
}

func (gm GalleryElements[T]) SortByName(sortOrder string) GalleryElements[T] {
	sort.Slice(gm, func(i, j int) bool {
		if sortOrder == "asc" {
			return strings.ToLower(gm[i].GetName()) < strings.ToLower(gm[j].GetName())
		} else {
			return strings.ToLower(gm[i].GetName()) > strings.ToLower(gm[j].GetName())
		}
	})
	return gm
}

func (gm GalleryElements[T]) SortByRepository(sortOrder string) GalleryElements[T] {
	sort.Slice(gm, func(i, j int) bool {
		if sortOrder == "asc" {
			return strings.ToLower(gm[i].GetGallery().Name) < strings.ToLower(gm[j].GetGallery().Name)
		} else {
			return strings.ToLower(gm[i].GetGallery().Name) > strings.ToLower(gm[j].GetGallery().Name)
		}
	})
	return gm
}

func (gm GalleryElements[T]) SortByLicense(sortOrder string) GalleryElements[T] {
	sort.Slice(gm, func(i, j int) bool {
		licenseI := gm[i].GetLicense()
		licenseJ := gm[j].GetLicense()
		var result bool
		if licenseI == "" && licenseJ != "" {
			return sortOrder == "desc"
		} else if licenseI != "" && licenseJ == "" {
			return sortOrder == "asc"
		} else if licenseI == "" && licenseJ == "" {
			return false
		} else {
			result = strings.ToLower(licenseI) < strings.ToLower(licenseJ)
		}
		if sortOrder == "desc" {
			return !result
		} else {
			return result
		}
	})
	return gm
}

func (gm GalleryElements[T]) SortByInstalled(sortOrder string) GalleryElements[T] {
	sort.Slice(gm, func(i, j int) bool {
		var result bool
		// Sort by installed status: installed items first (true > false)
		if gm[i].GetInstalled() != gm[j].GetInstalled() {
			result = gm[i].GetInstalled()
		} else {
			result = strings.ToLower(gm[i].GetName()) < strings.ToLower(gm[j].GetName())
		}
		if sortOrder == "desc" {
			return !result
		} else {
			return result
		}
	})
	return gm
}

func (gm GalleryElements[T]) FindByName(name string) T {
	for _, m := range gm {
		if strings.EqualFold(m.GetName(), name) {
			return m
		}
	}
	var zero T
	return zero
}

func (gm GalleryElements[T]) Paginate(pageNum int, itemsNum int) GalleryElements[T] {
	start := (pageNum - 1) * itemsNum
	end := start + itemsNum
	if start > len(gm) {
		start = len(gm)
	}
	if end > len(gm) {
		end = len(gm)
	}
	return gm[start:end]
}

func FindGalleryElement[T GalleryElement](models []T, name string) T {
	var model T
	name = strings.ReplaceAll(name, string(os.PathSeparator), "__")

	if !strings.Contains(name, "@") {
		for _, m := range models {
			if strings.EqualFold(strings.ToLower(m.GetName()), strings.ToLower(name)) {
				model = m
				break
			}
		}

	} else {
		for _, m := range models {
			if strings.EqualFold(strings.ToLower(name), strings.ToLower(fmt.Sprintf("%s@%s", m.GetGallery().Name, m.GetName()))) {
				model = m
				break
			}
		}
	}

	return model
}

// List available models
// Models galleries are a list of yaml files that are hosted on a remote server (for example github).
// Each yaml file contains a list of models that can be downloaded and optionally overrides to define a new model setting.
func AvailableGalleryModels(galleries []config.Gallery, systemState *system.SystemState) (GalleryElements[*GalleryModel], error) {
	var models []*GalleryModel

	// Get models from galleries
	for _, gallery := range galleries {
		galleryModels, err := getGalleryElements(gallery, systemState.Model.ModelsPath, func(model *GalleryModel) bool {
			if _, err := os.Stat(filepath.Join(systemState.Model.ModelsPath, fmt.Sprintf("%s.yaml", model.GetName()))); err == nil {
				return true
			}
			return false
		})
		if err != nil {
			return nil, err
		}
		models = append(models, galleryModels...)
	}

	return models, nil
}

// List available backends
func AvailableBackends(galleries []config.Gallery, systemState *system.SystemState) (GalleryElements[*GalleryBackend], error) {
	return availableBackendsWithFilter(galleries, systemState, true)
}

// AvailableBackendsUnfiltered returns all available backends without filtering by system capability.
func AvailableBackendsUnfiltered(galleries []config.Gallery, systemState *system.SystemState) (GalleryElements[*GalleryBackend], error) {
	return availableBackendsWithFilter(galleries, systemState, false)
}

// availableBackendsWithFilter is a helper function that lists available backends with optional filtering.
func availableBackendsWithFilter(galleries []config.Gallery, systemState *system.SystemState, filterByCapability bool) (GalleryElements[*GalleryBackend], error) {
	var backends []*GalleryBackend

	systemBackends, err := ListSystemBackends(systemState)
	if err != nil {
		return nil, err
	}

	// Get backends from galleries
	for _, gallery := range galleries {
		galleryBackends, err := getGalleryElements(gallery, systemState.Backend.BackendsPath, func(backend *GalleryBackend) bool {
			return systemBackends.Exists(backend.GetName())
		})
		if err != nil {
			return nil, err
		}

		// Filter backends by system capability if requested
		if filterByCapability {
			for _, backend := range galleryBackends {
				if backend.IsCompatibleWith(systemState) {
					backends = append(backends, backend)
				}
			}
		} else {
			backends = append(backends, galleryBackends...)
		}
	}

	return backends, nil
}

func findGalleryURLFromReferenceURL(url string, basePath string) (string, error) {
	var refFile string
	uri := downloader.URI(url)
	err := uri.ReadWithCallback(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)
		}
		cutPoint := strings.LastIndex(url, "/")
		refFile = url[:cutPoint+1] + refFile
		return nil
	})
	return refFile, err
}

type galleryCacheEntry struct {
	yamlEntry   []byte
	lastUpdated time.Time
}

func (entry galleryCacheEntry) hasExpired() bool {
	return entry.lastUpdated.Before(time.Now().Add(-1 * time.Hour))
}

var galleryCache = xsync.NewSyncedMap[string, galleryCacheEntry]()

func getGalleryElements[T GalleryElement](gallery config.Gallery, basePath string, isInstalledCallback func(T) bool) ([]T, error) {
	var models []T = []T{}

	if strings.HasSuffix(gallery.URL, ".ref") {
		var err error
		gallery.URL, err = findGalleryURLFromReferenceURL(gallery.URL, basePath)
		if err != nil {
			return models, err
		}
	}

	cacheKey := fmt.Sprintf("%s-%s", gallery.Name, gallery.URL)
	if galleryCache.Exists(cacheKey) {
		entry := galleryCache.Get(cacheKey)
		// refresh if last updated is more than 1 hour ago
		if !entry.hasExpired() {
			err := yaml.Unmarshal(entry.yamlEntry, &models)
			if err != nil {
				return models, err
			}
		} else {
			galleryCache.Delete(cacheKey)
		}
	}

	uri := downloader.URI(gallery.URL)

	if len(models) == 0 {
		err := uri.ReadWithCallback(basePath, func(url string, d []byte) error {
			galleryCache.Set(cacheKey, galleryCacheEntry{
				yamlEntry:   d,
				lastUpdated: time.Now(),
			})
			return yaml.Unmarshal(d, &models)
		})
		if err != nil {
			if yamlErr, ok := err.(*yaml.TypeError); ok {
				xlog.Debug("YAML errors", "errors", strings.Join(yamlErr.Errors, "\n"), "models", models)
			}
			return models, fmt.Errorf("failed to read gallery elements: %w", err)
		}
	}

	// Add gallery to models
	for _, model := range models {
		model.SetGallery(gallery)
		model.SetInstalled(isInstalledCallback(model))
	}
	return models, nil
}