|
@@ -9,6 +9,7 @@ import (
|
|
|
"net/http"
|
|
"net/http"
|
|
|
"one-api/common"
|
|
"one-api/common"
|
|
|
"one-api/dto"
|
|
"one-api/dto"
|
|
|
|
|
+ relaycommon "one-api/relay/common"
|
|
|
relayconstant "one-api/relay/constant"
|
|
relayconstant "one-api/relay/constant"
|
|
|
"one-api/service"
|
|
"one-api/service"
|
|
|
"strings"
|
|
"strings"
|
|
@@ -16,7 +17,7 @@ import (
|
|
|
"time"
|
|
"time"
|
|
|
)
|
|
)
|
|
|
|
|
|
|
|
-func OpenaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*dto.OpenAIErrorWithStatusCode, string, int) {
|
|
|
|
|
|
|
+func OpenaiStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.RelayInfo) (*dto.OpenAIErrorWithStatusCode, string, int) {
|
|
|
//checkSensitive := constant.ShouldCheckCompletionSensitive()
|
|
//checkSensitive := constant.ShouldCheckCompletionSensitive()
|
|
|
var responseTextBuilder strings.Builder
|
|
var responseTextBuilder strings.Builder
|
|
|
toolCount := 0
|
|
toolCount := 0
|
|
@@ -57,7 +58,7 @@ func OpenaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*d
|
|
|
}
|
|
}
|
|
|
}
|
|
}
|
|
|
streamResp := "[" + strings.Join(streamItems, ",") + "]"
|
|
streamResp := "[" + strings.Join(streamItems, ",") + "]"
|
|
|
- switch relayMode {
|
|
|
|
|
|
|
+ switch info.RelayMode {
|
|
|
case relayconstant.RelayModeChatCompletions:
|
|
case relayconstant.RelayModeChatCompletions:
|
|
|
var streamResponses []dto.ChatCompletionsStreamResponseSimple
|
|
var streamResponses []dto.ChatCompletionsStreamResponseSimple
|
|
|
err := json.Unmarshal(common.StringToByteSlice(streamResp), &streamResponses)
|
|
err := json.Unmarshal(common.StringToByteSlice(streamResp), &streamResponses)
|
|
@@ -126,9 +127,14 @@ func OpenaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*d
|
|
|
common.SafeSendBool(stopChan, true)
|
|
common.SafeSendBool(stopChan, true)
|
|
|
}()
|
|
}()
|
|
|
service.SetEventStreamHeaders(c)
|
|
service.SetEventStreamHeaders(c)
|
|
|
|
|
+ isFirst := true
|
|
|
c.Stream(func(w io.Writer) bool {
|
|
c.Stream(func(w io.Writer) bool {
|
|
|
select {
|
|
select {
|
|
|
case data := <-dataChan:
|
|
case data := <-dataChan:
|
|
|
|
|
+ if isFirst {
|
|
|
|
|
+ isFirst = false
|
|
|
|
|
+ info.FirstResponseTime = time.Now()
|
|
|
|
|
+ }
|
|
|
if strings.HasPrefix(data, "data: [DONE]") {
|
|
if strings.HasPrefix(data, "data: [DONE]") {
|
|
|
data = data[:12]
|
|
data = data[:12]
|
|
|
}
|
|
}
|