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

README.md

@@ -83,6 +83,8 @@
 ```
 可以实现400错误转为500错误，从而重试
 
+## 比原版One API多出的配置
+- `STREAMING_TIMEOUT`：设置流式一次回复的超时时间，默认为 30 秒
 
 ## 部署
 ### 部署要求

common/go-channel.go

@@ -3,6 +3,7 @@ package common
 import (
 	"fmt"
 	"runtime/debug"
+	"time"
 )
 
 func SafeGoroutine(f func()) {
@@ -45,3 +46,21 @@ func SafeSendString(ch chan string, value string) (closed bool) {
 	// If the code reaches here, then the channel was not closed.
 	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
+	}
+}

constant/system-setting.go

@@ -1,9 +1,13 @@
 package constant
 
+import "one-api/common"
+
 var ServerAddress = "http://localhost:3000"
 var WorkerUrl = ""
 var WorkerValidKey = ""
 
+var StreamingTimeout = common.GetOrDefault("STREAMING_TIMEOUT", 30)
+
 func EnableWorker() bool {
 	return WorkerUrl != ""
 }

relay/channel/openai/relay-openai.go

@@ -8,6 +8,7 @@ import (
 	"io"
 	"net/http"
 	"one-api/common"
+	"one-api/constant"
 	"one-api/dto"
 	relaycommon "one-api/relay/common"
 	relayconstant "one-api/relay/constant"
@@ -51,7 +52,11 @@ func OpenaiStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.
 			if data[:6] != "data: " && data[:6] != "[DONE]" {
 				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:]
 			if !strings.HasPrefix(data, "[DONE]") {
 				streamItems = append(streamItems, data)