feat: update stream_options

This commit is contained in:
CalciumIon 2024-07-19 14:46:25 +08:00
parent 436d08b48f
commit 20da8228df

View File

@ -16,6 +16,7 @@ import (
relayconstant "one-api/relay/constant" relayconstant "one-api/relay/constant"
"one-api/service" "one-api/service"
"strings" "strings"
"sync"
"time" "time"
) )
@ -41,7 +42,10 @@ func OaiStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.Rel
stopChan := make(chan bool) stopChan := make(chan bool)
defer close(stopChan) defer close(stopChan)
var (
lastStreamData string
mu sync.Mutex
)
gopool.Go(func() { gopool.Go(func() {
for scanner.Scan() { for scanner.Scan() {
info.SetFirstResponseTime() info.SetFirstResponseTime()
@ -53,14 +57,19 @@ func OaiStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.Rel
if data[:6] != "data: " && data[:6] != "[DONE]" { if data[:6] != "data: " && data[:6] != "[DONE]" {
continue continue
} }
mu.Lock()
data = data[6:] data = data[6:]
if !strings.HasPrefix(data, "[DONE]") { if !strings.HasPrefix(data, "[DONE]") {
err := service.StringData(c, data) if lastStreamData != "" {
err := service.StringData(c, lastStreamData)
if err != nil { if err != nil {
common.LogError(c, "streaming error: "+err.Error()) common.LogError(c, "streaming error: "+err.Error())
} }
}
lastStreamData = data
streamItems = append(streamItems, data) streamItems = append(streamItems, data)
} }
mu.Unlock()
} }
common.SafeSendBool(stopChan, true) common.SafeSendBool(stopChan, true)
}) })
@ -73,6 +82,22 @@ func OaiStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.Rel
// 正常结束 // 正常结束
} }
shouldSendLastResp := true
var lastStreamResponse dto.ChatCompletionsStreamResponse
err := json.Unmarshal(common.StringToByteSlice(lastStreamData), &lastStreamResponse)
if err == nil {
if lastStreamResponse.Usage != nil && service.ValidUsage(lastStreamResponse.Usage) {
if info.ShouldIncludeUsage {
containStreamUsage = true
} else {
shouldSendLastResp = false
}
}
}
if shouldSendLastResp {
service.StringData(c, lastStreamData)
}
// 计算token // 计算token
streamResp := "[" + strings.Join(streamItems, ",") + "]" streamResp := "[" + strings.Join(streamItems, ",") + "]"
switch info.RelayMode { switch info.RelayMode {