fix: enabled auth add support watsonx backend

Signed-off-by: Guangya Liu <[email protected]>
k8sgpt-ai · Jul 22, 2024 · 14ba86a · 14ba86a
1 parent 34b6de3
commit 14ba86a
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 18 deletions.
diff --git a/cmd/auth/add.go b/cmd/auth/add.go
@@ -48,6 +48,9 @@ var addCmd = &cobra.Command{
 		if strings.ToLower(backend) == "amazonbedrock" {
 			_ = cmd.MarkFlagRequired("providerRegion")
 		}
+		if strings.ToLower(backend) == "watsonxai" {
+			_ = cmd.MarkFlagRequired("providerId")
+		}
 	},
 	Run: func(cmd *cobra.Command, args []string) {
 
@@ -173,8 +176,8 @@ func init() {
 	addCmd.Flags().StringVarP(&engine, "engine", "e", "", "Azure AI deployment name (only for azureopenai backend)")
 	//add flag for amazonbedrock region name
 	addCmd.Flags().StringVarP(&providerRegion, "providerRegion", "r", "", "Provider Region name (only for amazonbedrock, googlevertexai backend)")
-	//add flag for vertexAI Project ID
-	addCmd.Flags().StringVarP(&providerId, "providerId", "i", "", "Provider specific ID for e.g. project (only for googlevertexai backend)")
+	//add flag for vertexAI/WatsonxAI Project ID
+	addCmd.Flags().StringVarP(&providerId, "providerId", "i", "", "Provider specific ID for e.g. project (only for googlevertexai/watsonxai backend)")
 	//add flag for OCI Compartment ID
 	addCmd.Flags().StringVarP(&compartmentId, "compartmentId", "k", "", "Compartment ID for generative AI model (only for oci backend)")
 	// add flag for openai organization

diff --git a/pkg/ai/iai.go b/pkg/ai/iai.go
@@ -181,7 +181,7 @@ func (p *AIProvider) GetCustomHeaders() []http.Header {
 	return p.CustomHeaders
 }
 
-var passwordlessProviders = []string{"localai", "ollama", "amazonsagemaker", "amazonbedrock", "googlevertexai", "oci", "watsonxai"}
+var passwordlessProviders = []string{"localai", "ollama", "amazonsagemaker", "amazonbedrock", "googlevertexai", "oci"}
 
 func NeedPassword(backend string) bool {
 	for _, b := range passwordlessProviders {

diff --git a/pkg/ai/watsonxai.go b/pkg/ai/watsonxai.go
@@ -1,10 +1,9 @@
 package ai
 
 import (
-	"os"
-	"fmt"
 	"context"
 	"errors"
+	"fmt"
 
 	wx "github.com/IBM/watsonx-go/pkg/models"
 )
@@ -14,20 +13,20 @@ const watsonxAIClientName = "watsonxai"
 type WatsonxAIClient struct {
 	nopCloser
 
-	client         *wx.Client
-	model          string
-	temperature    float32
-	topP           float32
-	topK           int32
-	maxNewTokens   int
+	client       *wx.Client
+	model        string
+	temperature  float32
+	topP         float32
+	topK         int32
+	maxNewTokens int
 }
 
 const (
 	modelMetallama = "ibm/granite-13b-chat-v2"
 )
 
 func (c *WatsonxAIClient) Configure(config IAIConfig) error {
-	if(config.GetModel() == "") {
+	if config.GetModel() == "" {
 		c.model = config.GetModel()
 	} else {
 		c.model = modelMetallama
@@ -37,20 +36,19 @@ func (c *WatsonxAIClient) Configure(config IAIConfig) error {
 	c.topK = config.GetTopK()
 	c.maxNewTokens = config.GetMaxTokens()
 
-	// WatsonxAPIKeyEnvVarName    = "WATSONX_API_KEY"
-	// WatsonxProjectIDEnvVarName = "WATSONX_PROJECT_ID"
-	apiKey, projectID := os.Getenv(wx.WatsonxAPIKeyEnvVarName), os.Getenv(wx.WatsonxProjectIDEnvVarName)
-
+	apiKey := config.GetPassword()
 	if apiKey == "" {
 		return errors.New("No watsonx API key provided")
 	}
-	if projectID == "" {
+
+	projectId := config.GetProviderId()
+	if projectId == "" {
 		return errors.New("No watsonx project ID provided")
 	}
 
 	client, err := wx.NewClient(
 		wx.WithWatsonxAPIKey(apiKey),
-		wx.WithWatsonxProjectID(projectID),
+		wx.WithWatsonxProjectID(projectId),
 	)
 	if err != nil {
 		return fmt.Errorf("Failed to create client for testing. Error: %v", err)