feat: add azure openai provider (#309)

* feat: add azure openai provider Signed-off-by: Aris Boutselis <aris.boutselis@senseon.io> * feat: validate backend name Signed-off-by: Aris Boutselis <aris.boutselis@senseon.io> * fix: remove BaseURL from the mandatory env variables Signed-off-by: Aris Boutselis <arisboutselis08@gmail.com> * fix: conflicts Signed-off-by: Aris Boutselis <aris.boutselis@senseon.io> * chore: updated logo (#365) Signed-off-by: Alex Jones <alexsimonjones@gmail.com> * chore: added changing banners (#367) Signed-off-by: Alex Jones <alexsimonjones@gmail.com> * feat: add additionalLabels to Service Monitor (#366) * feat: add additionalLabels to Service Monitor Signed-off-by: Brad McCoy <bradmccoydev@gmail.com> * feat: update additionalLabels Signed-off-by: Brad McCoy <bradmccoydev@gmail.com> --------- Signed-off-by: Brad McCoy <bradmccoydev@gmail.com> * fix: update README file's ai provider section. Signed-off-by: Aris Boutselis <aris.boutselis@senseon.io> --------- Signed-off-by: Aris Boutselis <aris.boutselis@senseon.io> Signed-off-by: Aris Boutselis <arisboutselis08@gmail.com> Signed-off-by: Alex Jones <alexsimonjones@gmail.com> Signed-off-by: Brad McCoy <bradmccoydev@gmail.com> Co-authored-by: Aris Boutselis <arisboutselis08@gmail.com> Co-authored-by: Alex Jones <alexsimonjones@gmail.com> Co-authored-by: Brad McCoy <bradmccoydev@gmail.com>
k8sgpt-ai · May 2, 2023 · d8357ce · d8357ce
1 parent a89a5cf
commit d8357ce
Show file tree

Hide file tree

Showing 5 changed files with 195 additions and 45 deletions.
diff --git a/README.md b/README.md
@@ -262,47 +262,37 @@ _Analysis with serve mode_
 ```
 curl -X GET "http://localhost:8080/analyze?namespace=k8sgpt&explain=false"
 ```
-
 </details>
 
-## How does anonymization work?
+## Additional AI providers  
 
-With this option, the data is anonymized before being sent to the AI Backend. During the analysis execution, `k8sgpt` retrieves sensitive data (Kubernetes object names, labels, etc.). This data is masked when sent to the AI backend and replaced by a key that can be used to de-anonymize the data when the solution is returned to the user.
+### Azure OpenAI  
+<em>Prerequisites:</em> an Azure OpenAI deployment is needed, please visit MS official [documentation](https://learn.microsoft.com/en-us/azure/cognitive-services/openai/how-to/create-resource?pivots=web-portal#create-a-resource) to create your own.
 
+To authenticate with k8sgpt, you will need the Azure OpenAI endpoint of your tenant `"https://your Azure OpenAI Endpoint"`, the api key to access your deployment, the deployment name of your model and the model name itself.  
 <details>
 
-1. Error reported during analysis:
-```bash
-Error: HorizontalPodAutoscaler uses StatefulSet/fake-deployment as ScaleTargetRef which does not exist.
+### Run k8sgpt  
+To run k8sgpt, run `k8sgpt auth` with the `azureopenai` backend:  
 ```
-
-2. Payload sent to the AI backend:
-```bash
-Error: HorizontalPodAutoscaler uses StatefulSet/tGLcCRcHa1Ce5Rs as ScaleTargetRef which does not exist.
+k8sgpt auth --backend azureopenai --baseurl https://<your Azure OpenAI endpoint> --engine <deployment_name> --model <model_name>
 ```
+Lastly, enter your Azure API key, after the prompt.
 
-3. Payload returned by the AI:
-```bash
-The Kubernetes system is trying to scale a StatefulSet named tGLcCRcHa1Ce5Rs using the HorizontalPodAutoscaler, but it cannot find the StatefulSet. The solution is to verify that the StatefulSet name is spelled correctly and exists in the same namespace as the HorizontalPodAutoscaler.
+Now you are ready to analyze with the azure openai backend:  
 ```
-
-4. Payload returned to the user:
-```bash
-The Kubernetes system is trying to scale a StatefulSet named fake-deployment using the HorizontalPodAutoscaler, but it cannot find the StatefulSet. The solution is to verify that the StatefulSet name is spelled correctly and exists in the same namespace as the HorizontalPodAutoscaler.
+k8sgpt analyze --explain --backend azureopenai
 ```
 
-**Anonymization does not currently apply to events.**
-
 </details>
 
-
-## Running local models
+### Running local models
 
 To run local models, it is possible to use OpenAI compatible APIs, for instance [LocalAI](https://github.com/go-skynet/LocalAI) which uses [llama.cpp](https://github.com/ggerganov/llama.cpp) and [ggml](https://github.com/ggerganov/ggml) to run inference on consumer-grade hardware. Models supported by LocalAI for instance are Vicuna, Alpaca, LLaMA, Cerebras, GPT4ALL, GPT4ALL-J and koala. 
 
 <details>
 
-To run local inference, you need to download the models first, for instance you can find `ggml` compatible models in [huggingface.co](https://huggingface.co/models?search=ggml) (for example vicuna, alpaca and koala).
+To run local inference, you need to download the models first, for instance you can find `ggml` compatible models in [huggingface.com](https://huggingface.co/models?search=ggml) (for example vicuna, alpaca and koala).
 
 ### Start the API server
 
@@ -316,14 +306,42 @@ To run k8sgpt, run `k8sgpt auth` with the `localai` backend:
 k8sgpt auth --backend localai --model <model_name> --baseurl http://localhost:8080/v1
 ```
 
-When being asked for an API key, just press enter.
-
 Now you can analyze with the `localai` backend:
 
 ```
 k8sgpt analyze --explain --backend localai
 ```
 
+</details>  
+
+## How does anonymization work?
+
+With this option, the data is anonymized before being sent to the AI Backend. During the analysis execution, `k8sgpt` retrieves sensitive data (Kubernetes object names, labels, etc.). This data is masked when sent to the AI backend and replaced by a key that can be used to de-anonymize the data when the solution is returned to the user.
+
+<details>
+
+1. Error reported during analysis:
+```bash
+Error: HorizontalPodAutoscaler uses StatefulSet/fake-deployment as ScaleTargetRef which does not exist.
+```
+
+2. Payload sent to the AI backend:
+```bash
+Error: HorizontalPodAutoscaler uses StatefulSet/tGLcCRcHa1Ce5Rs as ScaleTargetRef which does not exist.
+```
+
+3. Payload returned by the AI:
+```bash
+The Kubernetes system is trying to scale a StatefulSet named tGLcCRcHa1Ce5Rs using the HorizontalPodAutoscaler, but it cannot find the StatefulSet. The solution is to verify that the StatefulSet name is spelled correctly and exists in the same namespace as the HorizontalPodAutoscaler.
+```
+
+4. Payload returned to the user:
+```bash
+The Kubernetes system is trying to scale a StatefulSet named fake-deployment using the HorizontalPodAutoscaler, but it cannot find the StatefulSet. The solution is to verify that the StatefulSet name is spelled correctly and exists in the same namespace as the HorizontalPodAutoscaler.
+```
+
+**Anonymization does not currently apply to events.**
+
 </details>
 
 ## Configuration

diff --git a/cmd/auth/auth.go b/cmd/auth/auth.go
@@ -31,13 +31,21 @@ var (
 	password string
 	baseURL  string
 	model    string
+	engine   string
 )
 
 // authCmd represents the auth command
 var AuthCmd = &cobra.Command{
 	Use:   "auth",
 	Short: "Authenticate with your chosen backend",
 	Long:  `Provide the necessary credentials to authenticate with your chosen backend.`,
+	PreRun: func(cmd *cobra.Command, args []string) {
+		backend, _ := cmd.Flags().GetString("backend")
+		if strings.ToLower(backend) == "azureopenai" {
+			cmd.MarkFlagRequired("engine")
+			cmd.MarkFlagRequired("baseurl")
+		}
+	},
 	Run: func(cmd *cobra.Command, args []string) {
 
 		// get ai configuration
@@ -57,9 +65,18 @@ var AuthCmd = &cobra.Command{
 			}
 		}
 
-		// check if backend is not empty
-		if backend == "" {
-			color.Red("Error: Backend AI cannot be empty.")
+		validBackend := func(validBackends []string, backend string) bool {
+			for _, b := range validBackends {
+				if b == backend {
+					return true
+				}
+			}
+			return false
+		}
+
+		// check if backend is not empty and a valid value
+		if backend == "" || !validBackend(ai.Backends, backend) {
+			color.Red("Error: Backend AI cannot be empty and accepted values are '%v'", strings.Join(ai.Backends, ", "))
 			os.Exit(1)
 		}
 
@@ -88,6 +105,7 @@ var AuthCmd = &cobra.Command{
 			Model:    model,
 			Password: password,
 			BaseURL:  baseURL,
+			Engine:   engine,
 		}
 
 		if providerIndex == -1 {
@@ -117,4 +135,6 @@ func init() {
 	AuthCmd.Flags().StringVarP(&password, "password", "p", "", "Backend AI password")
 	// add flag for url
 	AuthCmd.Flags().StringVarP(&baseURL, "baseurl", "u", "", "URL AI provider, (e.g `http://localhost:8080/v1`)")
+	// add flag for azure open ai engine/deployment name
+	AuthCmd.Flags().StringVarP(&engine, "engine", "e", "", "Azure AI deployment name")
 }
diff --git a/cmd/serve/serve.go b/cmd/serve/serve.go
@@ -47,17 +47,17 @@ var ServeCmd = &cobra.Command{
 			password := os.Getenv("K8SGPT_PASSWORD")
 			model := os.Getenv("K8SGPT_MODEL")
 			baseURL := os.Getenv("K8SGPT_BASEURL")
-
+			engine := os.Getenv("K8SGPT_ENGINE")
 			// If the envs are set, allocate in place to the aiProvider
 			// else exit with error
-			envIsSet := backend != "" || password != "" || model != "" || baseURL != ""
-
+			envIsSet := backend != "" || password != "" || model != ""
 			if envIsSet {
 				aiProvider = &ai.AIProvider{
 					Name:     backend,
 					Password: password,
 					Model:    model,
 					BaseURL:  baseURL,
+					Engine:   engine,
 				}
 
 				configAI.Providers = append(configAI.Providers, *aiProvider)
@@ -75,10 +75,10 @@ var ServeCmd = &cobra.Command{
 		if aiProvider == nil {
 			for _, provider := range configAI.Providers {
 				if backend == provider.Name {
-          // he pointer to the range variable is not really an issue here, as there
-          // is a break right after, but to prevent potential future issues, a temp
-          // variable is assigned
-          p := provider
+					// the pointer to the range variable is not really an issue here, as there
+					// is a break right after, but to prevent potential future issues, a temp
+					// variable is assigned
+					p := provider
 					aiProvider = &p
 					break
 				}

diff --git a/pkg/ai/azureopenai.go b/pkg/ai/azureopenai.go
@@ -0,0 +1,94 @@
+package ai
+
+import (
+	"context"
+	"encoding/base64"
+	"errors"
+	"fmt"
+	"strings"
+
+	"github.com/k8sgpt-ai/k8sgpt/pkg/cache"
+	"github.com/k8sgpt-ai/k8sgpt/pkg/util"
+
+	"github.com/fatih/color"
+
+	"github.com/sashabaranov/go-openai"
+)
+
+type AzureAIClient struct {
+	client   *openai.Client
+	language string
+	model    string
+}
+
+func (c *AzureAIClient) Configure(config IAIConfig, lang string) error {
+	token := config.GetPassword()
+	baseURL := config.GetBaseURL()
+	engine := config.GetEngine()
+	defaultConfig := openai.DefaultAzureConfig(token, baseURL, engine)
+	client := openai.NewClientWithConfig(defaultConfig)
+	if client == nil {
+		return errors.New("error creating Azure OpenAI client")
+	}
+	c.language = lang
+	c.client = client
+	c.model = config.GetModel()
+	return nil
+}
+
+func (c *AzureAIClient) GetCompletion(ctx context.Context, prompt string) (string, error) {
+	// Create a completion request
+	resp, err := c.client.CreateChatCompletion(ctx, openai.ChatCompletionRequest{
+		Model: c.model,
+		Messages: []openai.ChatCompletionMessage{
+			{
+				Role:    "user",
+				Content: fmt.Sprintf(default_prompt, c.language, prompt),
+			},
+		},
+	})
+	if err != nil {
+		return "", err
+	}
+	return resp.Choices[0].Message.Content, nil
+}
+
+func (a *AzureAIClient) Parse(ctx context.Context, prompt []string, cache cache.ICache) (string, error) {
+	inputKey := strings.Join(prompt, " ")
+	// Check for cached data
+	cacheKey := util.GetCacheKey(a.GetName(), a.language, inputKey)
+
+	if !cache.IsCacheDisabled() && cache.Exists(cacheKey) {
+		response, err := cache.Load(cacheKey)
+		if err != nil {
+			return "", err
+		}
+
+		if response != "" {
+			output, err := base64.StdEncoding.DecodeString(response)
+			if err != nil {
+				color.Red("error decoding cached data: %v", err)
+				return "", nil
+			}
+			return string(output), nil
+		}
+	}
+
+	response, err := a.GetCompletion(ctx, inputKey)
+	if err != nil {
+		return "", err
+	}
+
+	err = cache.Store(cacheKey, base64.StdEncoding.EncodeToString([]byte(response)))
+
+	if err != nil {
+		color.Red("error storing value to cache: %v", err)
+		return "", nil
+	}
+
+	return response, nil
+}
+
+func (a *AzureAIClient) GetName() string {
+	return "azureopenai"
+}
diff --git a/pkg/ai/iai.go b/pkg/ai/iai.go
@@ -19,6 +19,21 @@ import (
 	"github.com/k8sgpt-ai/k8sgpt/pkg/cache"
 )
 
+var (
+	clients = []IAI{
+		&OpenAIClient{},
+		&AzureAIClient{},
+		&LocalAIClient{},
+		&NoOpAIClient{},
+	}
+	Backends = []string{
+		"openai",
+		"localai",
+		"azureopenai",
+		"noopai",
+	}
+)
+
 type IAI interface {
 	Configure(config IAIConfig, language string) error
 	GetCompletion(ctx context.Context, prompt string) (string, error)
@@ -30,19 +45,17 @@ type IAIConfig interface {
 	GetPassword() string
 	GetModel() string
 	GetBaseURL() string
+	GetEngine() string
 }
 
 func NewClient(provider string) IAI {
-	switch provider {
-	case "openai":
-		return &OpenAIClient{}
-	case "localai":
-		return &LocalAIClient{}
-	case "noopai":
-		return &NoOpAIClient{}
-	default:
-		return &OpenAIClient{}
+	for _, c := range clients {
+		if provider == c.GetName() {
+			return c
+		}
 	}
+	// default client
+	return &OpenAIClient{}
 }
 
 type AIConfiguration struct {
@@ -52,8 +65,9 @@ type AIConfiguration struct {
 type AIProvider struct {
 	Name     string `mapstructure:"name"`
 	Model    string `mapstructure:"model"`
-	Password string `mapstructure:"password"`
-	BaseURL  string `mapstructure:"baseurl"`
+	Password string `mapstructure:"password" yaml:"password,omitempty"`
+	BaseURL  string `mapstructure:"baseurl" yaml:"baseurl,omitempty"`
+	Engine   string `mapstructure:"engine" yaml:"engine,omitempty"`
 }
 
 func (p *AIProvider) GetBaseURL() string {
@@ -68,6 +82,10 @@ func (p *AIProvider) GetModel() string {
 	return p.Model
 }
 
+func (p *AIProvider) GetEngine() string {
+	return p.Engine
+}
+
 func NeedPassword(backend string) bool {
 	return backend != "localai"
 }