CaIon 1 год назад
Родитель
Сommit
21f32605c8
2 измененных файлов с 18 добавлено и 3 удалено
  1. 16 1
      common/go-channel.go
  2. 2 2
      relay/channel/openai/relay-openai.go

+ 16 - 1
common/go-channel.go

@@ -16,7 +16,22 @@ func SafeGoroutine(f func()) {
 	}()
 	}()
 }
 }
 
 
-func SafeSend(ch chan bool, value bool) (closed bool) {
+func SafeSendBool(ch chan bool, value bool) (closed bool) {
+	defer func() {
+		// Recover from panic if one occured. A panic would mean the channel was closed.
+		if recover() != nil {
+			closed = true
+		}
+	}()
+
+	// This will panic if the channel is closed.
+	ch <- value
+
+	// If the code reaches here, then the channel was not closed.
+	return false
+}
+
+func SafeSendString(ch chan string, value string) (closed bool) {
 	defer func() {
 	defer func() {
 		// Recover from panic if one occured. A panic would mean the channel was closed.
 		// Recover from panic if one occured. A panic would mean the channel was closed.
 		if recover() != nil {
 		if recover() != nil {

+ 2 - 2
relay/channel/openai/relay-openai.go

@@ -50,7 +50,7 @@ func OpenaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*d
 			if data[:6] != "data: " && data[:6] != "[DONE]" {
 			if data[:6] != "data: " && data[:6] != "[DONE]" {
 				continue
 				continue
 			}
 			}
-			dataChan <- data
+			common.SafeSendString(dataChan, data)
 			data = data[6:]
 			data = data[6:]
 			if !strings.HasPrefix(data, "[DONE]") {
 			if !strings.HasPrefix(data, "[DONE]") {
 				streamItems = append(streamItems, data)
 				streamItems = append(streamItems, data)
@@ -123,7 +123,7 @@ func OpenaiStreamHandler(c *gin.Context, resp *http.Response, relayMode int) (*d
 			// wait data out
 			// wait data out
 			time.Sleep(2 * time.Second)
 			time.Sleep(2 * time.Second)
 		}
 		}
-		common.SafeSend(stopChan, true)
+		common.SafeSendBool(stopChan, true)
 	}()
 	}()
 	service.SetEventStreamHeaders(c)
 	service.SetEventStreamHeaders(c)
 	c.Stream(func(w io.Writer) bool {
 	c.Stream(func(w io.Writer) bool {