Skip to content

Commit d6c9ba5

Browse files
liyanwei93yanweili
authored andcommitted
fix: add default maxToken value of watsonxai backend (k8sgpt-ai#1209)
Signed-off-by: yanweili <[email protected]> Co-authored-by: yanweili <[email protected]> Signed-off-by: Ronald Petty <[email protected]>
1 parent 89ca9a7 commit d6c9ba5

File tree

1 file changed

+16
-12
lines changed

1 file changed

+16
-12
lines changed

pkg/ai/watsonxai.go

+16-12
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,10 @@
11
package ai
22

33
import (
4-
"os"
5-
"fmt"
64
"context"
75
"errors"
6+
"fmt"
7+
"os"
88

99
wx "github.com/IBM/watsonx-go/pkg/models"
1010
)
@@ -14,28 +14,33 @@ const watsonxAIClientName = "watsonxai"
1414
type WatsonxAIClient struct {
1515
nopCloser
1616

17-
client *wx.Client
18-
model string
19-
temperature float32
20-
topP float32
21-
topK int32
22-
maxNewTokens int
17+
client *wx.Client
18+
model string
19+
temperature float32
20+
topP float32
21+
topK int32
22+
maxNewTokens int
2323
}
2424

2525
const (
2626
modelMetallama = "ibm/granite-13b-chat-v2"
27+
maxTokens = 2048
2728
)
2829

2930
func (c *WatsonxAIClient) Configure(config IAIConfig) error {
30-
if(config.GetModel() == "") {
31+
if config.GetModel() == "" {
32+
c.model = modelMetallama
33+
} else {
3134
c.model = config.GetModel()
35+
}
36+
if config.GetMaxTokens() == 0 {
37+
c.maxNewTokens = maxTokens
3238
} else {
33-
c.model = modelMetallama
39+
c.maxNewTokens = config.GetMaxTokens()
3440
}
3541
c.temperature = config.GetTemperature()
3642
c.topP = config.GetTopP()
3743
c.topK = config.GetTopK()
38-
c.maxNewTokens = config.GetMaxTokens()
3944

4045
// WatsonxAPIKeyEnvVarName = "WATSONX_API_KEY"
4146
// WatsonxProjectIDEnvVarName = "WATSONX_PROJECT_ID"
@@ -75,7 +80,6 @@ func (c *WatsonxAIClient) GetCompletion(ctx context.Context, prompt string) (str
7580
if result.Text == "" {
7681
return "", errors.New("Expected a result, but got an empty string")
7782
}
78-
7983
return result.Text, nil
8084
}
8185

0 commit comments

Comments
 (0)