78 lines
2.3 KiB
Go
78 lines
2.3 KiB
Go
package lmcli
|
|
|
|
import (
|
|
"fmt"
|
|
"os"
|
|
|
|
"git.mlow.ca/mlow/lmcli/pkg/util"
|
|
"gopkg.in/yaml.v3"
|
|
)
|
|
|
|
type Config struct {
|
|
Defaults *struct {
|
|
Model *string `yaml:"model" default:"gpt-4"`
|
|
MaxTokens *int `yaml:"maxTokens" default:"256"`
|
|
Temperature *float32 `yaml:"temperature" default:"0.2"`
|
|
SystemPrompt string `yaml:"systemPrompt,omitempty"`
|
|
SystemPromptFile string `yaml:"systemPromptFile,omitempty"`
|
|
Agent string `yaml:"agent"`
|
|
} `yaml:"defaults"`
|
|
Conversations *struct {
|
|
TitleGenerationModel *string `yaml:"titleGenerationModel" default:"gpt-3.5-turbo"`
|
|
} `yaml:"conversations"`
|
|
Chroma *struct {
|
|
Style *string `yaml:"style" default:"onedark"`
|
|
Formatter *string `yaml:"formatter" default:"terminal16m"`
|
|
} `yaml:"chroma"`
|
|
Agents []*struct {
|
|
Name string `yaml:"name"`
|
|
SystemPrompt string `yaml:"systemPrompt"`
|
|
Tools []string `yaml:"tools"`
|
|
} `yaml:"agents"`
|
|
Providers []*struct {
|
|
Name string `yaml:"name,omitempty"`
|
|
Display string `yaml:"display,omitempty"`
|
|
Kind string `yaml:"kind"`
|
|
BaseURL string `yaml:"baseUrl,omitempty"`
|
|
APIKey string `yaml:"apiKey,omitempty"`
|
|
Models []string `yaml:"models"`
|
|
Headers map[string]string `yaml:"headers"`
|
|
} `yaml:"providers"`
|
|
}
|
|
|
|
func NewConfig(configFile string) (*Config, error) {
|
|
shouldWriteDefaults := false
|
|
c := &Config{}
|
|
|
|
configExists := true
|
|
configBytes, err := os.ReadFile(configFile)
|
|
if os.IsNotExist(err) {
|
|
configExists = false
|
|
} else if err != nil {
|
|
return nil, fmt.Errorf("Could not read config file: %v", err)
|
|
} else {
|
|
yaml.Unmarshal(configBytes, c)
|
|
}
|
|
|
|
shouldWriteDefaults = util.SetStructDefaults(c)
|
|
if !configExists || shouldWriteDefaults {
|
|
if configExists {
|
|
fmt.Printf("Saving new defaults to configuration, backing up existing configuration to %s\n", configFile+".bak")
|
|
os.Rename(configFile, configFile+".bak")
|
|
}
|
|
fmt.Printf("Writing configuration file to %s\n", configFile)
|
|
file, err := os.Create(configFile)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("Could not open config file for writing: %v", err)
|
|
}
|
|
encoder := yaml.NewEncoder(file)
|
|
encoder.SetIndent(2)
|
|
err = encoder.Encode(c)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("Could not save default configuration: %v", err)
|
|
}
|
|
}
|
|
|
|
return c, nil
|
|
}
|