73ad85ea69
Refactor code to support starting of multiple back end llama.cpp servers. This functionality is exposed as `profiles` to create a simple configuration format. Changes: * refactor proxy tests to get ready for multi-process support * update proxy/ProxyManager to support multiple processes (#7) * Add support for Groups in configuration * improve handling of Model alias configs * implement multi-model swapping * improve code clarity for swapModel * improve docs, rename groups to profiles in config
93 lines
2.0 KiB
Go
93 lines
2.0 KiB
Go
package proxy
|
|
|
|
import (
|
|
"fmt"
|
|
"os"
|
|
"strings"
|
|
|
|
"gopkg.in/yaml.v3"
|
|
)
|
|
|
|
type ModelConfig struct {
|
|
Cmd string `yaml:"cmd"`
|
|
Proxy string `yaml:"proxy"`
|
|
Aliases []string `yaml:"aliases"`
|
|
Env []string `yaml:"env"`
|
|
CheckEndpoint string `yaml:"checkEndpoint"`
|
|
UnloadAfter int `yaml:"ttl"`
|
|
}
|
|
|
|
func (m *ModelConfig) SanitizedCommand() ([]string, error) {
|
|
return SanitizeCommand(m.Cmd)
|
|
}
|
|
|
|
type Config struct {
|
|
HealthCheckTimeout int `yaml:"healthCheckTimeout"`
|
|
Models map[string]ModelConfig `yaml:"models"`
|
|
Profiles map[string][]string `yaml:"profiles"`
|
|
|
|
// map aliases to actual model IDs
|
|
aliases map[string]string
|
|
}
|
|
|
|
func (c *Config) RealModelName(search string) (string, bool) {
|
|
if _, found := c.Models[search]; found {
|
|
return search, true
|
|
} else if name, found := c.aliases[search]; found {
|
|
return name, found
|
|
} else {
|
|
return "", false
|
|
}
|
|
}
|
|
|
|
func (c *Config) FindConfig(modelName string) (ModelConfig, string, bool) {
|
|
if realName, found := c.RealModelName(modelName); !found {
|
|
return ModelConfig{}, "", false
|
|
} else {
|
|
return c.Models[realName], realName, true
|
|
}
|
|
}
|
|
|
|
func LoadConfig(path string) (*Config, error) {
|
|
data, err := os.ReadFile(path)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
var config Config
|
|
err = yaml.Unmarshal(data, &config)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if config.HealthCheckTimeout < 15 {
|
|
config.HealthCheckTimeout = 15
|
|
}
|
|
|
|
// Populate the aliases map
|
|
config.aliases = make(map[string]string)
|
|
for modelName, modelConfig := range config.Models {
|
|
for _, alias := range modelConfig.Aliases {
|
|
config.aliases[alias] = modelName
|
|
}
|
|
}
|
|
|
|
return &config, nil
|
|
}
|
|
|
|
func SanitizeCommand(cmdStr string) ([]string, error) {
|
|
// Remove trailing backslashes
|
|
cmdStr = strings.ReplaceAll(cmdStr, "\\ \n", " ")
|
|
cmdStr = strings.ReplaceAll(cmdStr, "\\\n", " ")
|
|
|
|
// Split the command into arguments
|
|
args := strings.Fields(cmdStr)
|
|
|
|
// Ensure the command is not empty
|
|
if len(args) == 0 {
|
|
return nil, fmt.Errorf("empty command")
|
|
}
|
|
|
|
return args, nil
|
|
}
|