From 9b243cdcaab1742fca2516bc1ae5710505e0eb65 Mon Sep 17 00:00:00 2001 From: Matthis Holleville Date: Wed, 5 Apr 2023 09:39:12 +0200 Subject: [PATCH] feat: add support for new configuration format This commit adds support for a new configuration format that is not backwards compatible with the previous format. This is a breaking change and requires users to update their configuration files to use the new format. BREAKING CHANGE: The format of the configuration file has changed. Users must update their configuration files to use the new format. Signed-off-by: Matthis Holleville --- cmd/analyze/analyze.go | 43 +++++++++++++++++----------- cmd/auth/auth.go | 64 +++++++++++++++++++++++++++++++++--------- pkg/ai/ai.go | 19 +++++++++++-- pkg/ai/iai.go | 2 +- 4 files changed, 93 insertions(+), 35 deletions(-) diff --git a/cmd/analyze/analyze.go b/cmd/analyze/analyze.go index 2a3f176c1a..05c16c95a3 100644 --- a/cmd/analyze/analyze.go +++ b/cmd/analyze/analyze.go @@ -4,11 +4,12 @@ import ( "context" "encoding/json" "fmt" - "github.com/k8sgpt-ai/k8sgpt/pkg/analysis" - "github.com/k8sgpt-ai/k8sgpt/pkg/analyzer" "os" "strings" + "github.com/k8sgpt-ai/k8sgpt/pkg/analysis" + "github.com/k8sgpt-ai/k8sgpt/pkg/analyzer" + "github.com/fatih/color" "github.com/k8sgpt-ai/k8sgpt/pkg/ai" "github.com/k8sgpt-ai/k8sgpt/pkg/kubernetes" @@ -36,26 +37,34 @@ var AnalyzeCmd = &cobra.Command{ provide you with a list of issues that need to be resolved`, Run: func(cmd *cobra.Command, args []string) { - // get backend from file - backendType := viper.GetString("backend_type") - if backendType == "" { - color.Red("No backend set. Please run k8sgpt auth") + // get ai configuration + var configAI ai.AIConfiguration + err := viper.UnmarshalKey("ai", &configAI) + if err != nil { + color.Red("Error: %v", err) os.Exit(1) } - // override the default backend if a flag is provided - if backend != "" { - backendType = backend + + if len(configAI.Providers) == 0 { + color.Red("Error: AI provider not specified in configuration. Please run k8sgpt auth") + os.Exit(1) + } + + var aiProvider ai.AIProvider + for _, provider := range configAI.Providers { + if backend == provider.Name { + aiProvider = provider + break + } } - // get the token with viper - token := viper.GetString(fmt.Sprintf("%s_key", backendType)) - // check if nil - if token == "" { - color.Red("No %s key set. Please run k8sgpt auth", backendType) + + if aiProvider.Name == "" { + color.Red("Error: AI provider %s not specified in configuration. Please run k8sgpt auth", backend) os.Exit(1) } - aiClient := ai.NewClient(backendType) - if err := aiClient.Configure(token, language); err != nil { + aiClient := ai.NewClient(aiProvider.Name) + if err := aiClient.Configure(aiProvider.Password, aiProvider.Model, language); err != nil { color.Red("Error: %v", err) os.Exit(1) } @@ -73,7 +82,7 @@ var AnalyzeCmd = &cobra.Command{ Context: ctx, } - err := config.RunAnalysis() + err = config.RunAnalysis() if err != nil { color.Red("Error: %v", err) os.Exit(1) diff --git a/cmd/auth/auth.go b/cmd/auth/auth.go index 3017f0a9bc..d0c910f367 100644 --- a/cmd/auth/auth.go +++ b/cmd/auth/auth.go @@ -7,6 +7,7 @@ import ( "syscall" "github.com/fatih/color" + "github.com/k8sgpt-ai/k8sgpt/pkg/ai" "github.com/spf13/cobra" "github.com/spf13/viper" "golang.org/x/term" @@ -15,6 +16,7 @@ import ( var ( backend string password string + model string ) // authCmd represents the auth command @@ -24,33 +26,65 @@ var AuthCmd = &cobra.Command{ Long: `Provide the necessary credentials to authenticate with your chosen backend.`, Run: func(cmd *cobra.Command, args []string) { - backendType := viper.GetString("backend_type") - if backendType == "" { - // Set the default backend - viper.Set("backend_type", "openai") - if err := viper.WriteConfig(); err != nil { - color.Red("Error writing config file: %s", err.Error()) - os.Exit(1) + // get ai configuration + var configAI ai.AIConfiguration + err := viper.UnmarshalKey("ai", &configAI) + if err != nil { + color.Red("Error: %v", err) + os.Exit(1) + } + + // search for provider with same name + providerIndex := -1 + for i, provider := range configAI.Providers { + if backend == provider.Name { + providerIndex = i + break } } - // override the default backend if a flag is provided - if backend != "" { - backendType = backend - color.Green("Using %s as backend AI provider", backendType) + + // check if backend is not empty + if backend == "" { + color.Red("Error: Backend AI cannot be empty.") + os.Exit(1) + } + + color.Green("Using %s as backend AI provider", backend) + + // check if model is not empty + if model == "" { + color.Red("Error: Model cannot be empty.") + os.Exit(1) } if password == "" { - fmt.Printf("Enter %s Key: ", backendType) + fmt.Printf("Enter %s Key: ", backend) bytePassword, err := term.ReadPassword(int(syscall.Stdin)) if err != nil { - color.Red("Error reading %s Key from stdin: %s", backendType, + color.Red("Error reading %s Key from stdin: %s", backend, err.Error()) os.Exit(1) } password = strings.TrimSpace(string(bytePassword)) } - viper.Set(fmt.Sprintf("%s_key", backendType), password) + // create new provider object + newProvider := ai.AIProvider{ + Name: backend, + Model: model, + Password: password, + } + + if providerIndex == -1 { + // provider with same name does not exist, add new provider to list + configAI.Providers = append(configAI.Providers, newProvider) + color.Green("New provider added") + } else { + // provider with same name exists, update provider info + configAI.Providers[providerIndex] = newProvider + color.Green("Provider updated") + } + viper.Set("ai", configAI) if err := viper.WriteConfig(); err != nil { color.Red("Error writing config file: %s", err.Error()) os.Exit(1) @@ -62,6 +96,8 @@ var AuthCmd = &cobra.Command{ func init() { // add flag for backend AuthCmd.Flags().StringVarP(&backend, "backend", "b", "openai", "Backend AI provider") + // add flag for model + AuthCmd.Flags().StringVarP(&model, "model", "m", "gpt-3.5-turbo", "Backend AI model") // add flag for password AuthCmd.Flags().StringVarP(&password, "password", "p", "", "Backend AI password") } diff --git a/pkg/ai/ai.go b/pkg/ai/ai.go index 9bf0997be5..a7fa99c14f 100644 --- a/pkg/ai/ai.go +++ b/pkg/ai/ai.go @@ -5,9 +5,10 @@ import ( "encoding/base64" "errors" "fmt" + "strings" + "github.com/fatih/color" "github.com/spf13/viper" - "strings" "github.com/sashabaranov/go-openai" ) @@ -19,25 +20,37 @@ const ( prompt_c = "Reading the following %s error message and it's accompanying log message %s, how would you simplify this message?" ) +type AIConfiguration struct { + Providers []AIProvider `mapstructure:"providers"` +} + +type AIProvider struct { + Name string `mapstructure:"name"` + Model string `mapstructure:"model"` + Password string `mapstructure:"password"` +} + type OpenAIClient struct { client *openai.Client language string + model string } -func (c *OpenAIClient) Configure(token string, language string) error { +func (c *OpenAIClient) Configure(token string, model string, language string) error { client := openai.NewClient(token) if client == nil { return errors.New("error creating OpenAI client") } c.language = language c.client = client + c.model = model return nil } func (c *OpenAIClient) GetCompletion(ctx context.Context, prompt string) (string, error) { // Create a completion request resp, err := c.client.CreateChatCompletion(ctx, openai.ChatCompletionRequest{ - Model: openai.GPT3Dot5Turbo, + Model: c.model, Messages: []openai.ChatCompletionMessage{ { Role: "user", diff --git a/pkg/ai/iai.go b/pkg/ai/iai.go index cb12efa7ca..29c0e73803 100644 --- a/pkg/ai/iai.go +++ b/pkg/ai/iai.go @@ -5,7 +5,7 @@ import ( ) type IAI interface { - Configure(token string, language string) error + Configure(token string, model string, language string) error GetCompletion(ctx context.Context, prompt string) (string, error) Parse(ctx context.Context, prompt []string, nocache bool) (string, error) }