From cb36192c555b9ce2b85a3186d21df7ae6966d689 Mon Sep 17 00:00:00 2001
From: ckt1031 <65409152+ckt1031@users.noreply.github.com>
Date: Wed, 15 Nov 2023 20:16:50 +0800
Subject: [PATCH] Added validation for prompt length in relayImageHelper
 function

---
 common/model-ratio.go     | 6 +++++-
 controller/relay-image.go | 5 +++++
 2 files changed, 10 insertions(+), 1 deletion(-)

diff --git a/common/model-ratio.go b/common/model-ratio.go
index 204c1128..954ac3e4 100644
--- a/common/model-ratio.go
+++ b/common/model-ratio.go
@@ -18,12 +18,16 @@ var DalleSizeRatios = map[string]map[string]float64{
 	},
 }
 
-// Array in values, of maxmium and minimum
 var DalleGenerationImageAmounts = map[string][2]int{
 	"dall-e-2": {1, 10},
 	"dall-e-3": {1, 1}, // OpenAI allows n=1 currently.
 }
 
+var DalleImagePromptLengthLimitations = map[string]int{
+	"dall-e-2": 1000,
+	"dall-e-3": 4000,
+}
+
 // ModelRatio
 // https://platform.openai.com/docs/models/model-endpoint-compatibility
 // https://cloud.baidu.com/doc/WENXINWORKSHOP/s/Blfmc9dlf
diff --git a/controller/relay-image.go b/controller/relay-image.go
index 02ffcfb3..84893ebf 100644
--- a/controller/relay-image.go
+++ b/controller/relay-image.go
@@ -65,6 +65,11 @@ func relayImageHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode
 		return errorWrapper(errors.New("invalud value of n"), "number_of_generated_images_not_within_range", http.StatusBadRequest)
 	}
 
+	// Check prompt length
+	if len(imageRequest.Prompt) > common.DalleImagePromptLengthLimitations[imageModel] {
+		return errorWrapper(errors.New("prompt is too long"), "prompt_too_long", http.StatusBadRequest)
+	}
+
 	// map model name
 	modelMapping := c.GetString("model_mapping")
 	isModelMapped := false