lmcli/pkg/cli/config.go

package cli

import (
	"fmt"
	"os"
	"path/filepath"

	"github.com/go-yaml/yaml"
)

type Config struct {
	ModelDefaults *struct {
		SystemPrompt *string `yaml:"systemPrompt" default:"You are a helpful assistant."`
	} `yaml:"modelDefaults"`
	OpenAI *struct {
		APIKey           *string  `yaml:"apiKey" default:"your_key_here"`
		DefaultModel     *string  `yaml:"defaultModel" default:"gpt-4"`
		DefaultMaxLength *int     `yaml:"defaultMaxLength" default:"256"`
		EnabledTools     []string `yaml:"enabledTools"`
	} `yaml:"openai"`
	Chroma *struct {
		Style     *string `yaml:"style" default:"onedark"`
		Formatter *string `yaml:"formatter" default:"terminal16m"`
	} `yaml:"chroma"`
}

func configDir() string {
	var configDir string

	xdgConfigHome := os.Getenv("XDG_CONFIG_HOME")
	if xdgConfigHome != "" {
		configDir = filepath.Join(xdgConfigHome, "lmcli")
	} else {
		userHomeDir, _ := os.UserHomeDir()
		configDir = filepath.Join(userHomeDir, ".config/lmcli")
	}

	os.MkdirAll(configDir, 0755)
	return configDir
}

func NewConfig() (*Config, error) {
	configFile := filepath.Join(configDir(), "config.yaml")
	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 = 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
}