Просмотр исходного кода

fix: handle case where no response is received from Gemini API

CaIon 7 месяцев назад
Родитель
Сommit
07a92293e4

+ 6 - 0
relay/channel/gemini/relay-gemini-native.go

@@ -1,6 +1,7 @@
 package gemini
 
 import (
+	"github.com/pkg/errors"
 	"io"
 	"net/http"
 	"one-api/common"
@@ -107,6 +108,7 @@ func GeminiTextGenerationStreamHandler(c *gin.Context, info *relaycommon.RelayIn
 
 		// 直接发送 GeminiChatResponse 响应
 		err = helper.StringData(c, data)
+		info.SendResponseCount++
 		if err != nil {
 			common.LogError(c, err.Error())
 		}
@@ -114,6 +116,10 @@ func GeminiTextGenerationStreamHandler(c *gin.Context, info *relaycommon.RelayIn
 		return true
 	})
 
+	if info.SendResponseCount == 0 {
+		return nil, types.NewOpenAIError(errors.New("no response received from Gemini API"), types.ErrorCodeEmptyResponse, http.StatusInternalServerError)
+	}
+
 	if imageCount != 0 {
 		if usage.CompletionTokens == 0 {
 			usage.CompletionTokens = imageCount * 258

+ 2 - 5
relay/channel/gemini/relay-gemini.go

@@ -827,8 +827,6 @@ func GeminiChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *
 	var usage = &dto.Usage{}
 	var imageCount int
 
-	respCount := 0
-
 	helper.StreamScannerHandler(c, resp, info, func(data string) bool {
 		var geminiResponse GeminiChatResponse
 		err := common.UnmarshalJsonStr(data, &geminiResponse)
@@ -858,7 +856,7 @@ func GeminiChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *
 			}
 		}
 
-		if respCount == 0 {
+		if info.SendResponseCount == 0 {
 			// send first response
 			err = handleStream(c, info, helper.GenerateStartEmptyResponse(id, createAt, info.UpstreamModelName, nil))
 			if err != nil {
@@ -873,11 +871,10 @@ func GeminiChatStreamHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *
 		if isStop {
 			_ = handleStream(c, info, helper.GenerateStopResponse(id, createAt, info.UpstreamModelName, constant.FinishReasonStop))
 		}
-		respCount++
 		return true
 	})
 
-	if respCount == 0 {
+	if info.SendResponseCount == 0 {
 		// 空补全,报错不计费
 		// empty response, throw an error
 		return nil, types.NewOpenAIError(errors.New("no response received from Gemini API"), types.ErrorCodeEmptyResponse, http.StatusInternalServerError)