core/gallery/models.go

package gallery

import (
	"context"
	"errors"
	"fmt"
	"os"
	"path/filepath"
	"slices"
	"strings"

	"dario.cat/mergo"
	lconfig "github.com/mudler/LocalAI/core/config"
	"github.com/mudler/LocalAI/pkg/downloader"
	"github.com/mudler/LocalAI/pkg/model"
	"github.com/mudler/LocalAI/pkg/system"
	"github.com/mudler/LocalAI/pkg/utils"

	"github.com/mudler/xlog"
	"gopkg.in/yaml.v3"
)

/*

description: |
    foo
license: ""

urls:
-
-

name: "bar"

config_file: |
    # Note, name will be injected. or generated by the alias wanted by the user
    threads: 14

files:
    - filename: ""
      sha: ""
      uri: ""

prompt_templates:
    - name: ""
      content: ""

*/
// ModelConfig is the model configuration which contains all the model details
// This configuration is read from the gallery endpoint and is used to download and install the model
// It is the internal structure, separated from the request
type ModelConfig struct {
	Description     string           `yaml:"description"`
	Icon            string           `yaml:"icon"`
	License         string           `yaml:"license"`
	URLs            []string         `yaml:"urls"`
	Name            string           `yaml:"name"`
	ConfigFile      string           `yaml:"config_file"`
	Files           []File           `yaml:"files"`
	PromptTemplates []PromptTemplate `yaml:"prompt_templates"`
}

type File struct {
	Filename string `yaml:"filename" json:"filename"`
	SHA256   string `yaml:"sha256" json:"sha256"`
	URI      string `yaml:"uri" json:"uri"`
}

type PromptTemplate struct {
	Name    string `yaml:"name"`
	Content string `yaml:"content"`
}

// Installs a model from the gallery
func InstallModelFromGallery(
	ctx context.Context,
	modelGalleries, backendGalleries []lconfig.Gallery,
	systemState *system.SystemState,
	modelLoader *model.ModelLoader,
	name string, req GalleryModel, downloadStatus func(string, string, string, float64), enforceScan, automaticallyInstallBackend bool) error {

	applyModel := func(model *GalleryModel) error {
		name = strings.ReplaceAll(name, string(os.PathSeparator), "__")

		var config ModelConfig

		if len(model.URL) > 0 {
			var err error
			config, err = GetGalleryConfigFromURLWithContext[ModelConfig](ctx, model.URL, systemState.Model.ModelsPath)
			if err != nil {
				return err
			}
			config.Description = model.Description
			config.License = model.License
		} else if len(model.ConfigFile) > 0 {
			// TODO: is this worse than using the override method with a blank cfg yaml?
			reYamlConfig, err := yaml.Marshal(model.ConfigFile)
			if err != nil {
				return err
			}
			config = ModelConfig{
				ConfigFile:  string(reYamlConfig),
				Description: model.Description,
				License:     model.License,
				URLs:        model.URLs,
				Name:        model.Name,
				Files:       make([]File, 0), // Real values get added below, must be blank
				// Prompt Template Skipped for now - I expect in this mode that they will be delivered as files.
			}
		} else {
			return fmt.Errorf("invalid gallery model %+v", model)
		}

		installName := model.Name
		if req.Name != "" {
			installName = req.Name
		}

		// Copy the model configuration from the request schema
		config.URLs = append(config.URLs, model.URLs...)
		config.Icon = model.Icon
		config.Files = append(config.Files, req.AdditionalFiles...)
		config.Files = append(config.Files, model.AdditionalFiles...)

		// TODO model.Overrides could be merged with user overrides (not defined yet)
		if req.Overrides != nil {
			if err := mergo.Merge(&model.Overrides, req.Overrides, mergo.WithOverride); err != nil {
				return err
			}
		}

		installedModel, err := InstallModel(ctx, systemState, installName, &config, model.Overrides, downloadStatus, enforceScan)
		if err != nil {
			return err
		}
		xlog.Debug("Installed model", "model", installedModel.Name)
		if automaticallyInstallBackend && installedModel.Backend != "" {
			xlog.Debug("Installing backend", "backend", installedModel.Backend)

			if err := InstallBackendFromGallery(ctx, backendGalleries, systemState, modelLoader, installedModel.Backend, downloadStatus, false); err != nil {
				return err
			}
		}

		return nil
	}

	models, err := AvailableGalleryModels(modelGalleries, systemState)
	if err != nil {
		return err
	}

	model := FindGalleryElement(models, name)
	if model == nil {
		return fmt.Errorf("no model found with name %q", name)
	}

	return applyModel(model)
}

func InstallModel(ctx context.Context, systemState *system.SystemState, nameOverride string, config *ModelConfig, configOverrides map[string]interface{}, downloadStatus func(string, string, string, float64), enforceScan bool) (*lconfig.ModelConfig, error) {
	basePath := systemState.Model.ModelsPath
	// Create base path if it doesn't exist
	err := os.MkdirAll(basePath, 0750)
	if err != nil {
		return nil, fmt.Errorf("failed to create base path: %v", err)
	}

	if len(configOverrides) > 0 {
		xlog.Debug("Config overrides", "overrides", configOverrides)
	}

	// Download files and verify their SHA
	for i, file := range config.Files {
		// Check for cancellation before each file
		select {
		case <-ctx.Done():
			return nil, ctx.Err()
		default:
		}

		xlog.Debug("Checking file exists and matches SHA", "filename", file.Filename)

		if err := utils.VerifyPath(file.Filename, basePath); err != nil {
			return nil, err
		}

		// Create file path
		filePath := filepath.Join(basePath, file.Filename)

		if enforceScan {
			scanResults, err := downloader.HuggingFaceScan(downloader.URI(file.URI))
			if err != nil && errors.Is(err, downloader.ErrUnsafeFilesFound) {
				xlog.Error("Contains unsafe file(s)!", "model", config.Name, "clamAV", scanResults.ClamAVInfectedFiles, "pickles", scanResults.DangerousPickles)
				return nil, err
			}
		}
		uri := downloader.URI(file.URI)
		if err := uri.DownloadFileWithContext(ctx, filePath, file.SHA256, i, len(config.Files), downloadStatus); err != nil {
			return nil, err
		}
	}

	// Write prompt template contents to separate files
	for _, template := range config.PromptTemplates {
		if err := utils.VerifyPath(template.Name+".tmpl", basePath); err != nil {
			return nil, err
		}
		// Create file path
		filePath := filepath.Join(basePath, template.Name+".tmpl")

		// Create parent directory
		err := os.MkdirAll(filepath.Dir(filePath), 0750)
		if err != nil {
			return nil, fmt.Errorf("failed to create parent directory for prompt template %q: %v", template.Name, err)
		}
		// Create and write file content
		err = os.WriteFile(filePath, []byte(template.Content), 0600)
		if err != nil {
			return nil, fmt.Errorf("failed to write prompt template %q: %v", template.Name, err)
		}

		xlog.Debug("Prompt template written", "template", template.Name)
	}

	name := config.Name
	if nameOverride != "" {
		name = nameOverride
	}

	if err := utils.VerifyPath(name+".yaml", basePath); err != nil {
		return nil, err
	}

	modelConfig := lconfig.ModelConfig{}

	// write config file
	if len(configOverrides) != 0 || len(config.ConfigFile) != 0 {
		configFilePath := filepath.Join(basePath, name+".yaml")

		// Read and update config file as map[string]interface{}
		configMap := make(map[string]interface{})
		err = yaml.Unmarshal([]byte(config.ConfigFile), &configMap)
		if err != nil {
			return nil, fmt.Errorf("failed to unmarshal config YAML: %v", err)
		}

		configMap["name"] = name

		if configOverrides != nil {
			if err := mergo.Merge(&configMap, configOverrides, mergo.WithOverride); err != nil {
				return nil, err
			}
		}

		// Write updated config file
		updatedConfigYAML, err := yaml.Marshal(configMap)
		if err != nil {
			return nil, fmt.Errorf("failed to marshal updated config YAML: %v", err)
		}

		err = yaml.Unmarshal(updatedConfigYAML, &modelConfig)
		if err != nil {
			return nil, fmt.Errorf("failed to unmarshal updated config YAML: %v", err)
		}

		if valid, err := modelConfig.Validate(); !valid {
			return nil, fmt.Errorf("failed to validate updated config YAML: %v", err)
		}

		err = os.WriteFile(configFilePath, updatedConfigYAML, 0600)
		if err != nil {
			return nil, fmt.Errorf("failed to write updated config file: %v", err)
		}

		xlog.Debug("Written config file", "file", configFilePath)
	}

	// Save the model gallery file for further reference
	modelFile := filepath.Join(basePath, galleryFileName(name))
	data, err := yaml.Marshal(config)
	if err != nil {
		return nil, err
	}

	xlog.Debug("Written gallery file", "file", modelFile)

	return &modelConfig, os.WriteFile(modelFile, data, 0600)
}

func galleryFileName(name string) string {
	return "._gallery_" + name + ".yaml"
}

func GetLocalModelConfiguration(basePath string, name string) (*ModelConfig, error) {
	name = strings.ReplaceAll(name, string(os.PathSeparator), "__")
	galleryFile := filepath.Join(basePath, galleryFileName(name))
	return ReadConfigFile[ModelConfig](galleryFile)
}

func listModelFiles(systemState *system.SystemState, name string) ([]string, error) {

	configFile := filepath.Join(systemState.Model.ModelsPath, fmt.Sprintf("%s.yaml", name))
	if err := utils.VerifyPath(configFile, systemState.Model.ModelsPath); err != nil {
		return nil, fmt.Errorf("failed to verify path %s: %w", configFile, err)
	}

	// os.PathSeparator is not allowed in model names. Replace them with "__" to avoid conflicts with file paths.
	name = strings.ReplaceAll(name, string(os.PathSeparator), "__")

	galleryFile := filepath.Join(systemState.Model.ModelsPath, galleryFileName(name))
	if err := utils.VerifyPath(galleryFile, systemState.Model.ModelsPath); err != nil {
		return nil, fmt.Errorf("failed to verify path %s: %w", galleryFile, err)
	}

	additionalFiles := []string{}
	allFiles := []string{}

	// Galleryname is the name of the model in this case
	dat, err := os.ReadFile(configFile)
	if err == nil {
		modelConfig := &lconfig.ModelConfig{}

		err = yaml.Unmarshal(dat, &modelConfig)
		if err != nil {
			return nil, err
		}
		if modelConfig.Model != "" {
			additionalFiles = append(additionalFiles, modelConfig.ModelFileName())
		}

		if modelConfig.MMProj != "" {
			additionalFiles = append(additionalFiles, modelConfig.MMProjFileName())
		}
	}

	// read the model config
	galleryconfig, err := ReadConfigFile[ModelConfig](galleryFile)
	if err == nil && galleryconfig != nil {
		for _, f := range galleryconfig.Files {
			fullPath := filepath.Join(systemState.Model.ModelsPath, f.Filename)
			if err := utils.VerifyPath(fullPath, systemState.Model.ModelsPath); err != nil {
				return allFiles, fmt.Errorf("failed to verify path %s: %w", fullPath, err)
			}
			allFiles = append(allFiles, fullPath)
		}
	} else {
		xlog.Error("failed to read gallery file", "error", err, "file", configFile)
	}

	for _, f := range additionalFiles {
		fullPath := filepath.Join(filepath.Join(systemState.Model.ModelsPath, f))
		if err := utils.VerifyPath(fullPath, systemState.Model.ModelsPath); err != nil {
			return allFiles, fmt.Errorf("failed to verify path %s: %w", fullPath, err)
		}
		allFiles = append(allFiles, fullPath)
	}

	allFiles = append(allFiles, galleryFile)

	// skip duplicates
	allFiles = utils.Unique(allFiles)

	return allFiles, nil
}

func DeleteModelFromSystem(systemState *system.SystemState, name string) error {
	configFile := filepath.Join(systemState.Model.ModelsPath, fmt.Sprintf("%s.yaml", name))

	filesToRemove, err := listModelFiles(systemState, name)
	if err != nil {
		return err
	}

	allOtherFiles := []string{}
	// Get all files of all other models
	fi, err := os.ReadDir(systemState.Model.ModelsPath)
	if err != nil {
		return err
	}
	for _, f := range fi {
		if f.IsDir() {
			continue
		}
		if strings.HasPrefix(f.Name(), "._gallery_") {
			continue
		}
		if !strings.HasSuffix(f.Name(), ".yaml") && !strings.HasSuffix(f.Name(), ".yml") {
			continue
		}
		if f.Name() == fmt.Sprintf("%s.yaml", name) || f.Name() == fmt.Sprintf("%s.yml", name) {
			continue
		}

		name := strings.TrimSuffix(f.Name(), ".yaml")
		name = strings.TrimSuffix(name, ".yml")

		xlog.Debug("Checking file", "file", f.Name())
		files, err := listModelFiles(systemState, name)
		if err != nil {
			xlog.Debug("failed to list files for model", "error", err, "model", f.Name())
			continue
		}
		allOtherFiles = append(allOtherFiles, files...)
	}

	xlog.Debug("Files to remove", "files", filesToRemove)
	xlog.Debug("All other files", "files", allOtherFiles)

	// Removing files
	for _, f := range filesToRemove {
		if slices.Contains(allOtherFiles, f) {
			xlog.Debug("Skipping file because it is part of another model", "file", f)
			continue
		}
		if e := os.Remove(f); e != nil {
			xlog.Error("failed to remove file", "error", e, "file", f)
		}
	}

	return os.Remove(configFile)
}

// This is ***NEVER*** going to be perfect or finished.
// This is a BEST EFFORT function to surface known-vulnerable models to users.
func SafetyScanGalleryModels(galleries []lconfig.Gallery, systemState *system.SystemState) error {
	galleryModels, err := AvailableGalleryModels(galleries, systemState)
	if err != nil {
		return err
	}
	for _, gM := range galleryModels {
		if gM.Installed {
			err = errors.Join(err, SafetyScanGalleryModel(gM))
		}
	}
	return err
}

func SafetyScanGalleryModel(galleryModel *GalleryModel) error {
	for _, file := range galleryModel.AdditionalFiles {
		scanResults, err := downloader.HuggingFaceScan(downloader.URI(file.URI))
		if err != nil && errors.Is(err, downloader.ErrUnsafeFilesFound) {
			xlog.Error("Contains unsafe file(s)!", "model", galleryModel.Name, "clamAV", scanResults.ClamAVInfectedFiles, "pickles", scanResults.DangerousPickles)
			return err
		}
	}
	return nil
}