perf: flush response after response handled (close songquanpeng#364)

common/constants.go

@@ -55,6 +55,8 @@ var EmailDomainWhitelist = []string{
 	"foxmail.com",
 }
 
+var DebugEnabled = os.Getenv("DEBUG") == "true"
+
 var LogConsumeEnabled = true
 
 var SMTPServer = ""

controller/relay-text.go

@@ -5,13 +5,12 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"github.com/gin-gonic/gin"
 	"io"
 	"net/http"
 	"one-api/common"
 	"one-api/model"
 	"strings"
-
-	"github.com/gin-gonic/gin"
 )
 
 const (
@@ -308,6 +307,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 	var textResponse TextResponse
 
 	defer func() {
+		c.Writer.Flush()
 		if consumeQuota {
 			quota := 0
 			completionRatio := 1.0

main.go

@@ -26,6 +26,9 @@ func main() {
 	if os.Getenv("GIN_MODE") != "debug" {
 		gin.SetMode(gin.ReleaseMode)
 	}
+	if common.DebugEnabled {
+		common.SysLog("running in debug mode")
+	}
 	// Initialize SQL Database
 	err := model.InitDB()
 	if err != nil {

model/main.go

@@ -57,6 +57,9 @@ func InitDB() (err error) {
 	}
 	common.SysLog("database connected")
 	if err == nil {
+		if common.DebugEnabled {
+			db = db.Debug()
+		}
 		DB = db
 		sqlDB, err := DB.DB()
 		if err != nil {