xinyin025/new-api
1
0
Fork 0
mirror of https://github.com/linux-do/new-api.git synced 2025-09-17 16:06:38 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

feat: 支持设置流模式超时时间(gemini, claude)

Browse Source
This commit is contained in:
CalciumIon 2024-06-27 17:24:48 +08:00
parent 55c28b2f98
commit 402a415c79
3 changed files with 17 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
relay/channel/claude/relay-claude.go
View File

@ -8,6 +8,7 @@ import (
"io" "io"
"net/http" "net/http"
"one-api/common" "one-api/common"
"one-api/constant"
"one-api/dto" "one-api/dto"
relaycommon "one-api/relay/common" relaycommon "one-api/relay/common"
"one-api/service" "one-api/service"
@ -267,8 +268,8 @@ func claudeStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.
} }
return 0, nil, nil return 0, nil, nil
}) })
dataChan := make(chan string) dataChan := make(chan string, 5)
stopChan := make(chan bool) stopChan := make(chan bool, 2)
go func() { go func() {
for scanner.Scan() { for scanner.Scan() {
data := scanner.Text() data := scanner.Text()
@ -276,7 +277,11 @@ func claudeStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.
continue continue
} }
data = strings.TrimPrefix(data, "data: ") data = strings.TrimPrefix(data, "data: ")
dataChan <- data if !common.SafeSendStringTimeout(dataChan, data, constant.StreamingTimeout) {
// send data timeout, stop the stream
common.LogError(c, "send data timeout, stop the stream")
break
}
} }
stopChan <- true stopChan <- true
}() }()

11
relay/channel/gemini/relay-gemini.go
View File

@ -7,6 +7,7 @@ import (
"io" "io"
"net/http" "net/http"
"one-api/common" "one-api/common"
"one-api/constant"
"one-api/dto" "one-api/dto"
relaycommon "one-api/relay/common" relaycommon "one-api/relay/common"
"one-api/service" "one-api/service"
@ -163,8 +164,8 @@ func streamResponseGeminiChat2OpenAI(geminiResponse *GeminiChatResponse) *dto.Ch
func geminiChatStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*dto.OpenAIErrorWithStatusCode, string) { func geminiChatStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*dto.OpenAIErrorWithStatusCode, string) {
responseText := "" responseText := ""
dataChan := make(chan string) dataChan := make(chan string, 5)
stopChan := make(chan bool) stopChan := make(chan bool, 2)
scanner := bufio.NewScanner(resp.Body) scanner := bufio.NewScanner(resp.Body)
scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) { scanner.Split(func(data []byte, atEOF bool) (advance int, token []byte, err error) {
if atEOF && len(data) == 0 { if atEOF && len(data) == 0 {
@ -187,7 +188,11 @@ func geminiChatStreamHandler(c *gin.Context, resp *http.Response, info *relaycom
} }
data = strings.TrimPrefix(data, "\"text\": \"") data = strings.TrimPrefix(data, "\"text\": \"")
data = strings.TrimSuffix(data, "\"") data = strings.TrimSuffix(data, "\"")
dataChan <- data if !common.SafeSendStringTimeout(dataChan, data, constant.StreamingTimeout) {
// send data timeout, stop the stream
common.LogError(c, "send data timeout, stop the stream")
break
}
} }
stopChan <- true stopChan <- true
}() }()

2
relay/channel/openai/relay-openai.go
View File

@ -54,7 +54,7 @@ func OpenaiStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.
} }
if !common.SafeSendStringTimeout(dataChan, data, constant.StreamingTimeout) { if !common.SafeSendStringTimeout(dataChan, data, constant.StreamingTimeout) {
// send data timeout, stop the stream // send data timeout, stop the stream
common.LogInfo(c, "send data timeout, stop the stream") common.LogError(c, "send data timeout, stop the stream")
break break
} }
data = data[6:] data = data[6:]