groundwork: ListModels Filtering Upgrade (#2773)

* seperate the filtering from the middleware changes

---------

Signed-off-by: Dave Lee <dave@gray101.com>
This commit is contained in:
Dave 2024-10-01 14:55:46 -04:00 committed by GitHub
parent f84b55d1ef
commit 307a835199
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 387 additions and 81 deletions

View file

@ -1,55 +1,47 @@
package services
import (
"regexp"
"github.com/mudler/LocalAI/core/config"
"github.com/mudler/LocalAI/pkg/model"
)
func ListModels(bcl *config.BackendConfigLoader, ml *model.ModelLoader, filter string, excludeConfigured bool) ([]string, error) {
type LooseFilePolicy int
models, err := ml.ListFilesInModelPath()
if err != nil {
return nil, err
}
const (
SKIP_IF_CONFIGURED LooseFilePolicy = iota
SKIP_ALWAYS
ALWAYS_INCLUDE
LOOSE_ONLY
)
var mm map[string]interface{} = map[string]interface{}{}
func ListModels(bcl *config.BackendConfigLoader, ml *model.ModelLoader, filter config.BackendConfigFilterFn, looseFilePolicy LooseFilePolicy) ([]string, error) {
var skipMap map[string]interface{} = map[string]interface{}{}
dataModels := []string{}
var filterFn func(name string) bool
// If filter is not specified, do not filter the list by model name
if filter == "" {
filterFn = func(_ string) bool { return true }
} else {
// If filter _IS_ specified, we compile it to a regex which is used to create the filterFn
rxp, err := regexp.Compile(filter)
if err != nil {
return nil, err
}
filterFn = func(name string) bool {
return rxp.MatchString(name)
}
}
// Start with the known configurations
for _, c := range bcl.GetAllBackendConfigs() {
if excludeConfigured {
mm[c.Model] = nil
}
if filterFn(c.Name) {
// Start with known configurations
if looseFilePolicy != LOOSE_ONLY {
for _, c := range bcl.GetBackendConfigsByFilter(filter) {
if looseFilePolicy == SKIP_IF_CONFIGURED {
skipMap[c.Model] = nil
}
dataModels = append(dataModels, c.Name)
}
}
// Then iterate through the loose files:
for _, m := range models {
// And only adds them if they shouldn't be skipped.
if _, exists := mm[m]; !exists && filterFn(m) {
dataModels = append(dataModels, m)
// Then iterate through the loose files if requested.
if looseFilePolicy != SKIP_ALWAYS {
models, err := ml.ListFilesInModelPath()
if err != nil {
return nil, err
}
for _, m := range models {
// And only adds them if they shouldn't be skipped.
if _, exists := skipMap[m]; !exists && filter(m, nil) {
dataModels = append(dataModels, m)
}
}
}