feat: dall-e-3 (#47)

* Changes * 支持dall-e-3模型，对应的image.yml增加一个model参数 * 增加根据 model 检查 size
cssxsh · Nov 20, 2023 · 563d541 · 563d541
1 parent faa54d3
commit 563d541
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 1 deletion.
diff --git a/src/main/kotlin/xyz/cssxsh/mirai/openai/config/ImageConfig.kt b/src/main/kotlin/xyz/cssxsh/mirai/openai/config/ImageConfig.kt
@@ -23,9 +23,14 @@ internal object ImageConfig : ReadOnlyPluginConfig(saveName = "image") {
     @ValueDescription("url / b64_json")
     val format: ImageResponseFormat by value(ImageResponseFormat.URL)
 
+    @ValueName("model")
+    @ValueDescription("dall-e-2 | dall-e-3")
+    val model: String by value("dall-e-2")
+
     fun push(builder: ImageRequest.Builder) {
         builder.number = number
         builder.size = size
         builder.format = format
+        builder.model = model
     }
 }
diff --git a/src/main/kotlin/xyz/cssxsh/openai/image/ImageInfo.kt b/src/main/kotlin/xyz/cssxsh/openai/image/ImageInfo.kt
@@ -14,6 +14,8 @@ public data class ImageInfo(
         @SerialName("url")
         val url: String = "",
         @SerialName("b64_json")
-        val base64: String = ""
+        val base64: String = "",
+        @SerialName("revised_prompt")
+        var revisedPrompt:String = ""
     )
 }
diff --git a/src/main/kotlin/xyz/cssxsh/openai/image/ImageRequest.kt b/src/main/kotlin/xyz/cssxsh/openai/image/ImageRequest.kt
@@ -70,15 +70,41 @@ public data class ImageRequest(
             user = value
         }
 
+        @OpenAiDsl
+        public var model: String = "dall-e-2"
+
+        @OpenAiDsl
+        public fun model(value: String): Builder = apply {
+            model = value
+        }
+
         public fun build(): ImageRequest {
             require(number in 1..10) { "Must be between 1 and 10" }
+            // 这里添加对模型和尺寸的检查
+            validateSizeForModel(model, size)
             return ImageRequest(
                 prompt = prompt,
                 number = number,
                 size = size,
                 format = format,
                 user = user,
+                model = model,
             )
         }
+
+        private fun validateSizeForModel(model: String, size: ImageSize) {
+            when (model) {
+                "dall-e-2" -> require(size in listOf(ImageSize.SMALL, ImageSize.MIDDLE, ImageSize.LARGE)) {
+                    "dall-e-2 model requires one of the sizes: SMALL (256x256), MIDDLE (512x512), or LARGE (1024x1024)"
+                }
+
+                "dall-e-3" -> require(size in listOf(ImageSize.LARGE, ImageSize.LARGE_WIDTH, ImageSize.LARGE_HEIGHT)) {
+                    "dall-e-3 model requires one of the sizes: LARGE (1024x1024), LARGE_WIDTH (1792x1024), or LARGE_HEIGHT (1024x1792)"
+                }
+
+                else -> throw IllegalArgumentException("Unsupported model: $model")
+            }
+        }
+
     }
 }