Browse Source

fix: streaming timeout

CalciumIon 1 year ago
parent
commit
1e536ee7d9
1 changed files with 6 additions and 0 deletions
  1. 6 0
      relay/channel/openai/relay-openai.go

+ 6 - 0
relay/channel/openai/relay-openai.go

@@ -133,13 +133,19 @@ func OpenaiStreamHandler(c *gin.Context, resp *http.Response, info *relaycommon.
 	}()
 	}()
 	service.SetEventStreamHeaders(c)
 	service.SetEventStreamHeaders(c)
 	isFirst := true
 	isFirst := true
+	ticker := time.NewTicker(time.Duration(constant.StreamingTimeout))
+	defer ticker.Stop()
 	c.Stream(func(w io.Writer) bool {
 	c.Stream(func(w io.Writer) bool {
 		select {
 		select {
+		case <-ticker.C:
+			common.LogError(c, "reading data from upstream timeout")
+			return false
 		case data := <-dataChan:
 		case data := <-dataChan:
 			if isFirst {
 			if isFirst {
 				isFirst = false
 				isFirst = false
 				info.FirstResponseTime = time.Now()
 				info.FirstResponseTime = time.Now()
 			}
 			}
+			ticker.Reset(time.Duration(constant.StreamingTimeout))
 			if strings.HasPrefix(data, "data: [DONE]") {
 			if strings.HasPrefix(data, "data: [DONE]") {
 				data = data[:12]
 				data = data[:12]
 			}
 			}