mirror of
https://github.com/linux-do/new-api.git
synced 2025-09-20 01:16:36 +08:00
feat: update stream_options
This commit is contained in:
parent
436d08b48f
commit
20da8228df
@ -16,6 +16,7 @@ import (
|
||||
relayconstant "one-api/relay/constant"
|
||||
"one-api/service"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
@ -41,7 +42,10 @@ func OaiStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.Rel
|
||||
|
||||
stopChan := make(chan bool)
|
||||
defer close(stopChan)
|
||||
|
||||
var (
|
||||
lastStreamData string
|
||||
mu sync.Mutex
|
||||
)
|
||||
gopool.Go(func() {
|
||||
for scanner.Scan() {
|
||||
info.SetFirstResponseTime()
|
||||
@ -53,14 +57,19 @@ func OaiStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.Rel
|
||||
if data[:6] != "data: " && data[:6] != "[DONE]" {
|
||||
continue
|
||||
}
|
||||
mu.Lock()
|
||||
data = data[6:]
|
||||
if !strings.HasPrefix(data, "[DONE]") {
|
||||
err := service.StringData(c, data)
|
||||
if err != nil {
|
||||
common.LogError(c, "streaming error: "+err.Error())
|
||||
if lastStreamData != "" {
|
||||
err := service.StringData(c, lastStreamData)
|
||||
if err != nil {
|
||||
common.LogError(c, "streaming error: "+err.Error())
|
||||
}
|
||||
}
|
||||
lastStreamData = data
|
||||
streamItems = append(streamItems, data)
|
||||
}
|
||||
mu.Unlock()
|
||||
}
|
||||
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
|
||||
streamResp := "[" + strings.Join(streamItems, ",") + "]"
|
||||
switch info.RelayMode {
|
||||
|
Loading…
Reference in New Issue
Block a user