feat: upgrading azure client impl (k8sgpt-ai#526)

* feat: upgrading azure client impl Signed-off-by: Alex Jones <[email protected]> * feat: upgrading azure client impl Signed-off-by: Alex Jones <[email protected]> * fix: displaying correct yaml config location on app start (k8sgpt-ai#521) Signed-off-by: Ali Mohsin <[email protected]> Co-authored-by: Alex Jones <[email protected]> * fix(deps): update module github.com/aws/aws-sdk-go to v1.44.289 (k8sgpt-ai#524) Signed-off-by: renovate[bot] <29139614+renovate[bot]@users.noreply.github.com> Co-authored-by: renovate[bot] <29139614+renovate[bot]@users.noreply.github.com> * feat: upgrading azure client impl Signed-off-by: Alex Jones <[email protected]> --------- Signed-off-by: Alex Jones <[email protected]> Signed-off-by: Ali Mohsin <[email protected]> Signed-off-by: renovate[bot] <29139614+renovate[bot]@users.noreply.github.com> Co-authored-by: Ali Mohsin <[email protected]> Co-authored-by: renovate[bot] <29139614+renovate[bot]@users.noreply.github.com>
ruanxin · Jun 29, 2023 · 367fe8f · 367fe8f
1 parent 566f752
commit 367fe8f
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 3 deletions.
diff --git a/go.mod b/go.mod
@@ -7,7 +7,7 @@ require (
 	github.com/fatih/color v1.15.0
 	github.com/magiconair/properties v1.8.7
 	github.com/mittwald/go-helm-client v0.12.1
-	github.com/sashabaranov/go-openai v1.9.3
+	github.com/sashabaranov/go-openai v1.12.0
 	github.com/schollz/progressbar/v3 v3.13.1
 	github.com/spf13/cobra v1.7.0
 	github.com/spf13/viper v1.16.0

diff --git a/go.sum b/go.sum
@@ -1008,6 +1008,8 @@ github.com/samber/lo v1.38.1 h1:j2XEAqXKb09Am4ebOg31SpvzUTTs6EN3VfgeLUhPdXM=
 github.com/samber/lo v1.38.1/go.mod h1:+m/ZKRl6ClXCE2Lgf3MsQlWfh4bn1bz6CXEOxnEXnEA=
 github.com/sashabaranov/go-openai v1.9.3 h1:uNak3Rn5pPsKRs9bdT7RqRZEyej/zdZOEI2/8wvrFtM=
 github.com/sashabaranov/go-openai v1.9.3/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
+github.com/sashabaranov/go-openai v1.12.0 h1:aRNHH0gtVfrpIaEolD0sWrLLRnYQNK4cH/bIAHwL8Rk=
+github.com/sashabaranov/go-openai v1.12.0/go.mod h1:lj5b/K+zjTSFxVLijLSTDZuP7adOgerWeFyZLUhAKRg=
 github.com/schollz/progressbar/v3 v3.13.1 h1:o8rySDYiQ59Mwzy2FELeHY5ZARXZTVJC7iHD6PEFUiE=
 github.com/schollz/progressbar/v3 v3.13.1/go.mod h1:xvrbki8kfT1fzWzBT/UZd9L6GA+jdL7HAgq2RFnO6fQ=
 github.com/sean-/seed v0.0.0-20170313163322-e2103e2c3529/go.mod h1:DxrIzT+xaE7yg65j358z/aeFdxmN0P9QXhEzd20vsDc=

diff --git a/pkg/ai/azureopenai.go b/pkg/ai/azureopenai.go
@@ -25,7 +25,16 @@ func (c *AzureAIClient) Configure(config IAIConfig, lang string) error {
 	token := config.GetPassword()
 	baseURL := config.GetBaseURL()
 	engine := config.GetEngine()
-	defaultConfig := openai.DefaultAzureConfig(token, baseURL, engine)
+	defaultConfig := openai.DefaultAzureConfig(token, baseURL)
+
+	defaultConfig.AzureModelMapperFunc = func(model string) string {
+		// If you use a deployment name different from the model name, you can customize the AzureModelMapperFunc function
+		azureModelMapping := map[string]string{
+			model: engine,
+		}
+		return azureModelMapping[model]
+
+	}
 	client := openai.NewClientWithConfig(defaultConfig)
 	if client == nil {
 		return errors.New("error creating Azure OpenAI client")
@@ -42,7 +51,7 @@ func (c *AzureAIClient) GetCompletion(ctx context.Context, prompt string, prompt
 		Model: c.model,
 		Messages: []openai.ChatCompletionMessage{
 			{
-				Role:    "user",
+				Role:    openai.ChatMessageRoleUser,
 				Content: fmt.Sprintf(default_prompt, c.language, prompt),
 			},
 		},