修正AIProxy访问gpt-4o的max_tokens参数必填问题

songquanpeng · Jun 12, 2024 · b01f0af · b01f0af
1 parent e54abc0
commit b01f0af
Showing 1 changed file with 16 additions and 1 deletion.
diff --git a/relay/controller/text.go b/relay/controller/text.go
@@ -15,9 +15,14 @@ import (
 	"github.com/songquanpeng/one-api/relay/meta"
 	"github.com/songquanpeng/one-api/relay/model"
 	"io"
+	"math"
 	"net/http"
+	"os"
+	"strings"
 )
 
+var fixAIProxyGpt4oMaxTokens = os.Getenv("FIX_AI_PROXY_GPT4O_MAX_TOKENS") == "1"
+
 func RelayTextHelper(c *gin.Context) *model.ErrorWithStatusCode {
 	ctx := c.Request.Context()
 	meta := meta.GetByContext(c)
@@ -57,8 +62,18 @@ func RelayTextHelper(c *gin.Context) *model.ErrorWithStatusCode {
 	var requestBody io.Reader
 	if meta.APIType == apitype.OpenAI {
 		// no need to convert request for openai
-		shouldResetRequestBody := isModelMapped || meta.ChannelType == channeltype.Baichuan // frequency_penalty 0 is not acceptable for baichuan
+		shouldResetRequestBody := isModelMapped ||
+			meta.ChannelType == channeltype.Baichuan /*frequency_penalty 0 is not acceptable for baichuan*/ ||
+			(meta.ChannelType == channeltype.AIProxy && fixAIProxyGpt4oMaxTokens && strings.HasPrefix(textRequest.Model, "gpt-4o"))
 		if shouldResetRequestBody {
+			if meta.ChannelType == channeltype.AIProxy {
+				maxTokens := textRequest.MaxTokens
+				maxTokens = int(math.Min(float64(maxTokens), 4096))
+				if maxTokens == 0 {
+					maxTokens = 4096
+				}
+				textRequest.MaxTokens = maxTokens
+			}
 			jsonStr, err := json.Marshal(textRequest)
 			if err != nil {
 				return openai.ErrorWrapper(err, "json_marshal_failed", http.StatusInternalServerError)