From bdb2e739d4ad7c09103e64fde8923149ee5711a0 Mon Sep 17 00:00:00 2001 From: Alex Jones Date: Fri, 24 Mar 2023 12:01:36 +0000 Subject: [PATCH] interfaced out ai clients Signed-off-by: Alex Jones --- README.md | 7 +++---- cmd/auth/auth.go | 36 ++++++++++++++++++++++++------- cmd/auth/key.go | 42 ------------------------------------- cmd/find/problems.go | 31 ++++++++++++++++++++++----- cmd/init/init.go | 42 +++++++++++++++++++++++++++++++++++++ cmd/root.go | 4 +++- pkg/ai/ai.go | 29 ++++++++----------------- pkg/ai/iai.go | 8 +++++++ pkg/analyzer/analyzer.go | 3 ++- pkg/analyzer/podAnalyzer.go | 7 ++++--- pkg/analyzer/rsAnalyzer.go | 7 ++++--- 11 files changed, 129 insertions(+), 87 deletions(-) delete mode 100644 cmd/auth/key.go create mode 100644 cmd/init/init.go create mode 100644 pkg/ai/iai.go diff --git a/README.md b/README.md index f3e9745..bd6e230 100644 --- a/README.md +++ b/README.md @@ -3,7 +3,7 @@ Text changing depending on mode. Light: 'So light!' Dark: 'So dark!' -_Try it out now_ +_Install it now_ ``` brew tap k8sgpt-ai/k8sgpt @@ -20,12 +20,11 @@ It has SRE experience codified into it's analyzers and helps to pull out the mos ``` # Ensure KUBECONFIG env is set to an active Kubernetes cluster -k8sgpt auth key - +k8sgpt init +k8sgpt auth k8sgpt find problems # for more detail k8s find problems --explain - ``` ### What about kubectl-ai? diff --git a/cmd/auth/auth.go b/cmd/auth/auth.go index de5de0d..77fb0fd 100644 --- a/cmd/auth/auth.go +++ b/cmd/auth/auth.go @@ -1,21 +1,41 @@ package auth import ( + "fmt" + "os" + "strings" + "syscall" + + "github.com/fatih/color" "github.com/spf13/cobra" + "github.com/spf13/viper" + "golang.org/x/term" ) // authCmd represents the auth command var AuthCmd = &cobra.Command{ Use: "auth", - Short: "A brief description of your command", - Long: `A longer description that spans multiple lines and likely contains examples -and usage of using your command. For example: - -Cobra is a CLI library for Go that empowers applications. -This application is a tool to generate the needed files -to quickly create a Cobra application.`, + Short: "Authenticate with your chosen backend", + Long: `Provide the necessary credentials to authenticate with your chosen backend.`, Run: func(cmd *cobra.Command, args []string) { - cmd.Help() + + backendType := viper.GetString("backend_type") + + fmt.Printf("Enter %s Key: ", backendType) + bytePassword, err := term.ReadPassword(int(syscall.Stdin)) + if err != nil { + color.Red("Error reading %s Key from stdin: %s", backendType, + err.Error()) + os.Exit(1) + } + password := strings.TrimSpace(string(bytePassword)) + + viper.Set(fmt.Sprintf("%s_key", backendType), password) + if err := viper.WriteConfig(); err != nil { + color.Red("Error writing config file: %s", err.Error()) + os.Exit(1) + } + color.Green("key added") }, } diff --git a/cmd/auth/key.go b/cmd/auth/key.go deleted file mode 100644 index 91ad5be..0000000 --- a/cmd/auth/key.go +++ /dev/null @@ -1,42 +0,0 @@ -package auth - -import ( - "fmt" - "os" - "strings" - "syscall" - - "github.com/fatih/color" - "github.com/spf13/cobra" - "github.com/spf13/viper" - "golang.org/x/term" -) - -// keyCmd represents the key command -var keyCmd = &cobra.Command{ - Use: "key", - Short: "Add a key to OpenAI", - Long: `This command will add a key from OpenAI to enable you to interact with the API`, - Run: func(cmd *cobra.Command, args []string) { - - fmt.Print("Enter OpenAI API Key: ") - bytePassword, err := term.ReadPassword(int(syscall.Stdin)) - if err != nil { - color.Red("Error reading OpenAI API Key from stdin: %s", err.Error()) - os.Exit(1) - } - password := strings.TrimSpace(string(bytePassword)) - - viper.Set("openai_api_key", password) - if err := viper.WriteConfig(); err != nil { - color.Red("Error writing config file: %s", err.Error()) - os.Exit(1) - } - color.Green("key added") - }, -} - -func init() { - AuthCmd.AddCommand(keyCmd) - -} diff --git a/cmd/find/problems.go b/cmd/find/problems.go index 015b19e..22aac13 100644 --- a/cmd/find/problems.go +++ b/cmd/find/problems.go @@ -2,6 +2,7 @@ package find import ( "context" + "fmt" "os" "github.com/fatih/color" @@ -22,10 +23,30 @@ var problemsCmd = &cobra.Command{ provide you with a list of issues that need to be resolved`, Run: func(cmd *cobra.Command, args []string) { - // Initialise the openAI client - openAIClient, err := ai.NewClient() - if err != nil { - color.Red("Error: %v", err) + // get backend from file + backendType := viper.GetString("backend_type") + if backendType == "" { + color.Red("No backend set. Please run k8sgpt init") + os.Exit(1) + } + // 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) + os.Exit(1) + } + + var aiClient ai.IAI + switch backendType { + case "openai": + aiClient = &ai.OpenAIClient{} + if err := aiClient.Configure(token); err != nil { + color.Red("Error: %v", err) + os.Exit(1) + } + default: + color.Red("Backend not supported") os.Exit(1) } @@ -33,7 +54,7 @@ var problemsCmd = &cobra.Command{ // Get kubernetes client from viper client := viper.Get("kubernetesClient").(*kubernetes.Client) - if err := analyzer.RunAnalysis(ctx, client, openAIClient, explain); err != nil { + if err := analyzer.RunAnalysis(ctx, client, aiClient, explain); err != nil { color.Red("Error: %v", err) os.Exit(1) } diff --git a/cmd/init/init.go b/cmd/init/init.go new file mode 100644 index 0000000..b3d2557 --- /dev/null +++ b/cmd/init/init.go @@ -0,0 +1,42 @@ +package init + +import ( + "os" + + "github.com/fatih/color" + "github.com/spf13/cobra" + "github.com/spf13/viper" +) + +var ( + backend string +) + +// authCmd represents the auth command +var InitCmd = &cobra.Command{ + Use: "init", + Short: "Initialise k8sgpt with a backend AI provider", + Long: `Currently only OpenAI is supported.`, + Run: func(cmd *cobra.Command, args []string) { + + /* + This is largely a placeholder for now. In the future we will support + multiple backends and this will allow us to set the backend we want to use. + */ + if backend != "openai" { + color.Yellow("Only OpenAI is supported at the moment") + } + viper.Set("backend_type", backend) + if err := viper.WriteConfig(); err != nil { + color.Red("Error writing config file: %s", err.Error()) + os.Exit(1) + } + + color.Green("Backend set to %s", backend) + }, +} + +func init() { + // add flag for backend + InitCmd.Flags().StringVarP(&backend, "backend", "b", "openai", "Backend AI provider") +} diff --git a/cmd/root.go b/cmd/root.go index 75dd862..7d7d4fa 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -6,6 +6,7 @@ import ( "github.com/fatih/color" "github.com/k8sgpt-ai/k8sgpt/cmd/auth" "github.com/k8sgpt-ai/k8sgpt/cmd/find" + i "github.com/k8sgpt-ai/k8sgpt/cmd/init" "github.com/k8sgpt-ai/k8sgpt/pkg/kubernetes" "github.com/spf13/cobra" "github.com/spf13/viper" @@ -44,7 +45,7 @@ func init() { // will be global for your application. rootCmd.AddCommand(auth.AuthCmd) rootCmd.AddCommand(find.FindCmd) - + rootCmd.AddCommand(i.InitCmd) rootCmd.PersistentFlags().StringVar(&cfgFile, "config", "", "config file (default is $HOME/.k8sgpt.git.yaml)") rootCmd.PersistentFlags().StringVar(&masterURL, "master", "", "The address of the Kubernetes API server. Overrides any value in kubeconfig. Only required if out-of-cluster.") rootCmd.PersistentFlags().StringVar(&kubeconfig, "kubeconfig", "", "Path to a kubeconfig. Only required if out-of-cluster.") @@ -59,6 +60,7 @@ func init() { } viper.Set("kubernetesClient", kubernetesClient) + } // initConfig reads in config file and ENV variables if set. diff --git a/pkg/ai/ai.go b/pkg/ai/ai.go index 2e44e70..1c42b37 100644 --- a/pkg/ai/ai.go +++ b/pkg/ai/ai.go @@ -2,36 +2,25 @@ package ai import ( "context" - "fmt" + "errors" "github.com/sashabaranov/go-openai" - "github.com/spf13/viper" ) -type Client struct { +type OpenAIClient struct { client *openai.Client } -func (c *Client) GetClient() *openai.Client { - return c.client -} - -func NewClient() (*Client, error) { - - // get the token with viper - token := viper.GetString("openai_api_key") - // check if nil - if token == "" { - return nil, fmt.Errorf("no OpenAI API Key found") - } - +func (c *OpenAIClient) Configure(token string) error { client := openai.NewClient(token) - return &Client{ - client: client, - }, nil + + if client == nil { + return errors.New("error creating OpenAI client") + } + return nil } -func (c *Client) GetCompletion(ctx context.Context, prompt string) (string, error) { +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, diff --git a/pkg/ai/iai.go b/pkg/ai/iai.go new file mode 100644 index 0000000..8a3a91e --- /dev/null +++ b/pkg/ai/iai.go @@ -0,0 +1,8 @@ +package ai + +import "context" + +type IAI interface { + Configure(token string) error + GetCompletion(ctx context.Context, prompt string) (string, error) +} diff --git a/pkg/analyzer/analyzer.go b/pkg/analyzer/analyzer.go index 37b957c..4e9b4cd 100644 --- a/pkg/analyzer/analyzer.go +++ b/pkg/analyzer/analyzer.go @@ -2,11 +2,12 @@ package analyzer import ( "context" + "github.com/k8sgpt-ai/k8sgpt/pkg/ai" "github.com/k8sgpt-ai/k8sgpt/pkg/kubernetes" ) -func RunAnalysis(ctx context.Context, client *kubernetes.Client, aiClient *ai.Client, explain bool) error { +func RunAnalysis(ctx context.Context, client *kubernetes.Client, aiClient ai.IAI, explain bool) error { err := AnalyzePod(ctx, client, aiClient, explain) if err != nil { return err diff --git a/pkg/analyzer/podAnalyzer.go b/pkg/analyzer/podAnalyzer.go index f562c03..c8e1215 100644 --- a/pkg/analyzer/podAnalyzer.go +++ b/pkg/analyzer/podAnalyzer.go @@ -4,17 +4,18 @@ import ( "context" "encoding/base64" "fmt" + "strings" + "time" + "github.com/briandowns/spinner" "github.com/fatih/color" "github.com/k8sgpt-ai/k8sgpt/pkg/ai" "github.com/k8sgpt-ai/k8sgpt/pkg/kubernetes" "github.com/spf13/viper" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" - "strings" - "time" ) -func AnalyzePod(ctx context.Context, client *kubernetes.Client, aiClient *ai.Client, explain bool) error { +func AnalyzePod(ctx context.Context, client *kubernetes.Client, aiClient ai.IAI, explain bool) error { // search all namespaces for pods that are not running list, err := client.GetClient().CoreV1().Pods("").List(ctx, metav1.ListOptions{}) diff --git a/pkg/analyzer/rsAnalyzer.go b/pkg/analyzer/rsAnalyzer.go index 2841da2..a08a6d0 100644 --- a/pkg/analyzer/rsAnalyzer.go +++ b/pkg/analyzer/rsAnalyzer.go @@ -4,17 +4,18 @@ import ( "context" "encoding/base64" "fmt" + "strings" + "time" + "github.com/briandowns/spinner" "github.com/fatih/color" "github.com/k8sgpt-ai/k8sgpt/pkg/ai" "github.com/k8sgpt-ai/k8sgpt/pkg/kubernetes" "github.com/spf13/viper" metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" - "strings" - "time" ) -func AnalyzeReplicaSet(ctx context.Context, client *kubernetes.Client, aiClient *ai.Client, explain bool) error { +func AnalyzeReplicaSet(ctx context.Context, client *kubernetes.Client, aiClient ai.IAI, explain bool) error { // search all namespaces for pods that are not running list, err := client.GetClient().AppsV1().ReplicaSets("").List(ctx, metav1.ListOptions{})