From e2f5c1eb8cc188bf0d689f1ec40dc2a5e1e6e47f Mon Sep 17 00:00:00 2001
From: ckt1031 <65409152+ckt1031@users.noreply.github.com>
Date: Thu, 13 Jul 2023 22:07:07 +0800
Subject: [PATCH] fix: channel testing for reverse proxy

---
 controller/channel-test.go | 24 +++++++++++++++++++++++-
 controller/relay-text.go   | 19 +++++++++++++++++++
 controller/relay.go        |  2 +-
 3 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/controller/channel-test.go b/controller/channel-test.go
index 93d09d38..39aec6ba 100644
--- a/controller/channel-test.go
+++ b/controller/channel-test.go
@@ -55,7 +55,10 @@ func testChannel(channel *model.Channel, request ChatRequest) error {
 	}
 
 	if resp.StatusCode != http.StatusOK {
-		return errors.New("invalid status code: " + strconv.Itoa(resp.StatusCode))
+		// Prinnt the body in string
+		buf := new(bytes.Buffer)
+		buf.ReadFrom(resp.Body)
+		return errors.New("error response: " + strconv.Itoa(resp.StatusCode) + " " + buf.String())
 	}
 
 	var done = false
@@ -83,6 +86,25 @@ func testChannel(channel *model.Channel, request ChatRequest) error {
 			common.SysError("invalid stream response: " + data)
 			continue
 		}
+		// If data has event: event content inside, remove it, it can be prefix or inside the data
+		if strings.HasPrefix(data, "event:") || strings.Contains(data, "event:") {
+			// Remove event: event in the front or back
+			data = strings.TrimPrefix(data, "event: event")
+			data = strings.TrimSuffix(data, "event: event")
+			// Remove everything, only keep `data: {...}` <--- this is the json
+			// Find the start and end indices of `data: {...}` substring
+			startIndex := strings.Index(data, "data:")
+			endIndex := strings.LastIndex(data, "}")
+
+			// If both indices are found and end index is greater than start index
+			if startIndex != -1 && endIndex != -1 && endIndex > startIndex {
+				// Extract the `data: {...}` substring
+				data = data[startIndex : endIndex+1]
+			}
+
+			// Trim whitespace and newlines from the modified data string
+			data = strings.TrimSpace(data)
+		}
 		if !strings.HasPrefix(data, "data:") {
 			continue
 		}
diff --git a/controller/relay-text.go b/controller/relay-text.go
index de7f1768..fdbd6660 100644
--- a/controller/relay-text.go
+++ b/controller/relay-text.go
@@ -278,6 +278,25 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 					common.SysError("invalid stream response: " + data)
 					continue
 				}
+				// If data has event: event content inside, remove it, it can be prefix or inside the data
+				if strings.HasPrefix(data, "event:") || strings.Contains(data, "event:") {
+					// Remove event: event in the front or back
+					data = strings.TrimPrefix(data, "event: event")
+					data = strings.TrimSuffix(data, "event: event")
+					// Remove everything, only keep `data: {...}` <--- this is the json
+					// Find the start and end indices of `data: {...}` substring
+					startIndex := strings.Index(data, "data:")
+					endIndex := strings.LastIndex(data, "}")
+
+					// If both indices are found and end index is greater than start index
+					if startIndex != -1 && endIndex != -1 && endIndex > startIndex {
+						// Extract the `data: {...}` substring
+						data = data[startIndex : endIndex+1]
+					}
+
+					// Trim whitespace and newlines from the modified data string
+					data = strings.TrimSpace(data)
+				}
 				if !strings.HasPrefix(data, "data:") {
 					continue
 				}
diff --git a/controller/relay.go b/controller/relay.go
index 3b37910e..bef667ec 100644
--- a/controller/relay.go
+++ b/controller/relay.go
@@ -44,7 +44,7 @@ type GeneralOpenAIRequest struct {
 type ChatRequest struct {
 	Model     string    `json:"model"`
 	Messages  []Message `json:"messages"`
-	MaxTokens int       `json:"max_tokens"`
+	MaxTokens *int      `json:"max_tokens,omitempty"`
 	Stream    bool      `json:"stream"`
 }