Commit a9770319 by CalciumIon

feat: update stream_options

parent 5d4c5628
...@@ -16,6 +16,7 @@ import ( ...@@ -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 ...@@ -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 ...@@ -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 != "" {
if err != nil { err := service.StringData(c, lastStreamData)
common.LogError(c, "streaming error: "+err.Error()) if err != nil {
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 ...@@ -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 {
......
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