From 7c7eb6b7ece0865a94fb563de7376fe7d3e11940 Mon Sep 17 00:00:00 2001
From: JustSong <songquanpeng@foxmail.com>
Date: Mon, 12 Jun 2023 16:11:57 +0800
Subject: [PATCH] fix: now the input field can be array type now (close #149)

---
 controller/relay-utils.go | 14 ++++++++++++++
 controller/relay.go       |  4 ++--
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/controller/relay-utils.go b/controller/relay-utils.go
index bb25fa3b..a2dc2685 100644
--- a/controller/relay-utils.go
+++ b/controller/relay-utils.go
@@ -58,6 +58,20 @@ func countTokenMessages(messages []Message, model string) int {
 	return tokenNum
 }
 
+func countTokenInput(input any, model string) int {
+	switch input.(type) {
+	case string:
+		return countTokenText(input.(string), model)
+	case []string:
+		text := ""
+		for _, s := range input.([]string) {
+			text += s
+		}
+		return countTokenText(text, model)
+	}
+	return 0
+}
+
 func countTokenText(text string, model string) int {
 	tokenEncoder := getTokenEncoder(model)
 	token := tokenEncoder.Encode(text, nil, nil)
diff --git a/controller/relay.go b/controller/relay.go
index 7882a09d..35897909 100644
--- a/controller/relay.go
+++ b/controller/relay.go
@@ -38,7 +38,7 @@ type GeneralOpenAIRequest struct {
 	Temperature float64   `json:"temperature"`
 	TopP        float64   `json:"top_p"`
 	N           int       `json:"n"`
-	Input       string    `json:"input"`
+	Input       any       `json:"input"`
 }
 
 type ChatRequest struct {
@@ -189,7 +189,7 @@ func relayHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 	case RelayModeCompletions:
 		promptTokens = countTokenText(textRequest.Prompt, textRequest.Model)
 	case RelayModeModeration:
-		promptTokens = countTokenText(textRequest.Input, textRequest.Model)
+		promptTokens = countTokenInput(textRequest.Input, textRequest.Model)
 	}
 	preConsumedTokens := common.PreConsumedQuota
 	if textRequest.MaxTokens != 0 {