72 lines
2.0 KiB
Go
72 lines
2.0 KiB
Go
package lmcli
|
|
|
|
import (
|
|
"fmt"
|
|
"os"
|
|
|
|
"git.mlow.ca/mlow/lmcli/pkg/util"
|
|
"github.com/go-yaml/yaml"
|
|
)
|
|
|
|
type Config struct {
|
|
Defaults *struct {
|
|
SystemPrompt *string `yaml:"systemPrompt" default:"You are a helpful assistant."`
|
|
MaxTokens *int `yaml:"maxTokens" default:"256"`
|
|
Temperature *float32 `yaml:"temperature" default:"0.7"`
|
|
Model *string `yaml:"model" default:"gpt-4"`
|
|
} `yaml:"defaults"`
|
|
Conversations *struct {
|
|
TitleGenerationModel *string `yaml:"titleGenerationModel" default:"gpt-3.5-turbo"`
|
|
} `yaml:"conversations"`
|
|
Tools *struct {
|
|
EnabledTools *[]string `yaml:"enabledTools"`
|
|
} `yaml:"tools"`
|
|
OpenAI *struct {
|
|
APIKey *string `yaml:"apiKey" default:"your_key_here"`
|
|
Models *[]string `yaml:"models"`
|
|
} `yaml:"openai"`
|
|
Anthropic *struct {
|
|
APIKey *string `yaml:"apiKey" default:"your_key_here"`
|
|
Models *[]string `yaml:"models"`
|
|
} `yaml:"anthropic"`
|
|
Chroma *struct {
|
|
Style *string `yaml:"style" default:"onedark"`
|
|
Formatter *string `yaml:"formatter" default:"terminal16m"`
|
|
} `yaml:"chroma"`
|
|
}
|
|
|
|
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)
|
|
}
|
|
bytes, _ := yaml.Marshal(c)
|
|
_, err = file.Write(bytes)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("Could not save default configuration: %v", err)
|
|
}
|
|
}
|
|
|
|
return c, nil
|
|
}
|