|
|
@@ -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 {
|