Commit de23f667 by CaIon

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

parent 7eccf425
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
......
......@@ -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)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment