Commit e7a4aa05 by CalciumIon

feat: 支持设置流模式超时时间

parent 4b45e550
...@@ -83,6 +83,8 @@ ...@@ -83,6 +83,8 @@
``` ```
可以实现400错误转为500错误,从而重试 可以实现400错误转为500错误,从而重试
## 比原版One API多出的配置
- `STREAMING_TIMEOUT`:设置流式一次回复的超时时间,默认为 30 秒
## 部署 ## 部署
### 部署要求 ### 部署要求
......
...@@ -3,6 +3,7 @@ package common ...@@ -3,6 +3,7 @@ package common
import ( import (
"fmt" "fmt"
"runtime/debug" "runtime/debug"
"time"
) )
func SafeGoroutine(f func()) { func SafeGoroutine(f func()) {
...@@ -45,3 +46,21 @@ func SafeSendString(ch chan string, value string) (closed bool) { ...@@ -45,3 +46,21 @@ func SafeSendString(ch chan string, value string) (closed bool) {
// If the code reaches here, then the channel was not closed. // If the code reaches here, then the channel was not closed.
return false return false
} }
// SafeSendStringTimeout send, return true, else return false
func SafeSendStringTimeout(ch chan string, value string, timeout int) (closed bool) {
defer func() {
// Recover from panic if one occured. A panic would mean the channel was closed.
if recover() != nil {
closed = false
}
}()
// This will panic if the channel is closed.
select {
case ch <- value:
return true
case <-time.After(time.Duration(timeout) * time.Second):
return false
}
}
package constant package constant
import "one-api/common"
var ServerAddress = "http://localhost:3000" var ServerAddress = "http://localhost:3000"
var WorkerUrl = "" var WorkerUrl = ""
var WorkerValidKey = "" var WorkerValidKey = ""
var StreamingTimeout = common.GetOrDefault("STREAMING_TIMEOUT", 30)
func EnableWorker() bool { func EnableWorker() bool {
return WorkerUrl != "" return WorkerUrl != ""
} }
...@@ -8,6 +8,7 @@ import ( ...@@ -8,6 +8,7 @@ import (
"io" "io"
"net/http" "net/http"
"one-api/common" "one-api/common"
"one-api/constant"
"one-api/dto" "one-api/dto"
relaycommon "one-api/relay/common" relaycommon "one-api/relay/common"
relayconstant "one-api/relay/constant" relayconstant "one-api/relay/constant"
...@@ -51,7 +52,11 @@ func OpenaiStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon. ...@@ -51,7 +52,11 @@ func OpenaiStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.
if data[:6] != "data: " && data[:6] != "[DONE]" { if data[:6] != "data: " && data[:6] != "[DONE]" {
continue continue
} }
common.SafeSendString(dataChan, data) if !common.SafeSendStringTimeout(dataChan, data, constant.StreamingTimeout) {
// send data timeout, stop the stream
common.LogInfo(c, "send data timeout, stop the stream")
break
}
data = data[6:] data = data[6:]
if !strings.HasPrefix(data, "[DONE]") { if !strings.HasPrefix(data, "[DONE]") {
streamItems = append(streamItems, data) streamItems = append(streamItems, data)
......
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