Added validation for number of generated images

and defined image generation ratios

common/model-ratio.go

@@ -18,6 +18,12 @@ var DalleSizeRatios = map[string]map[string]float64{
 	},
 }
 
+// Array in values, of maxmium and minimum
+var DalleGenerationImageAmountRatios = map[string][2]int{
+	"dall-e-2": {1, 10},
+	"dall-e-3": {1, 1}, // OpenAI allows n=1 currently.
+}
+
 // ModelRatio
 // https://platform.openai.com/docs/models/model-endpoint-compatibility
 // https://cloud.baidu.com/doc/WENXINWORKSHOP/s/Blfmc9dlf

controller/relay-image.go

@@ -14,6 +14,17 @@ import (
 	"github.com/gin-gonic/gin"
 )
 
+func isWithinRange(element string, value int) bool {
+	if _, ok := common.DalleGenerationImageAmountRatios[element]; !ok {
+		return false
+	}
+
+	min := common.DalleGenerationImageAmountRatios[element][0]
+	max := common.DalleGenerationImageAmountRatios[element][1]
+
+	return value >= min && value <= max
+}
+
 func relayImageHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 	imageModel := "dall-e-2"
 	requestSize := "1024x1024"
@@ -49,6 +60,11 @@ func relayImageHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode
 		return errorWrapper(errors.New("prompt is required"), "required_field_missing", http.StatusBadRequest)
 	}
 
+	// Number of generated images validation
+	if isWithinRange(imageModel, imageRequest.N) == false {
+		return errorWrapper(errors.New("invalud value of n"), "number_of_generated_images_not_within_range", http.StatusBadRequest)
+	}
+
 	// map model name
 	modelMapping := c.GetString("model_mapping")
 	isModelMapped := false