fix: Update prompt message content types to use Literal and add union type for content (#17136)

Co-authored-by: 朱庆超 <zhuqingchao@xiaomi.com>
Co-authored-by: crazywoola <427733928@qq.com>
This commit is contained in:
ZalterCitty
2025-04-22 16:17:55 +08:00
committed by GitHub
parent 404f8a790c
commit a1158cc946
10 changed files with 73 additions and 39 deletions

View File

@@ -7,9 +7,9 @@ from core.model_runtime.entities import (
AudioPromptMessageContent,
DocumentPromptMessageContent,
ImagePromptMessageContent,
MultiModalPromptMessageContent,
VideoPromptMessageContent,
)
from core.model_runtime.entities.message_entities import PromptMessageContentUnionTypes
from extensions.ext_storage import storage
from . import helpers
@@ -43,7 +43,7 @@ def to_prompt_message_content(
/,
*,
image_detail_config: ImagePromptMessageContent.DETAIL | None = None,
) -> MultiModalPromptMessageContent:
) -> PromptMessageContentUnionTypes:
if f.extension is None:
raise ValueError("Missing file extension")
if f.mime_type is None:
@@ -58,7 +58,7 @@ def to_prompt_message_content(
if f.type == FileType.IMAGE:
params["detail"] = image_detail_config or ImagePromptMessageContent.DETAIL.LOW
prompt_class_map: Mapping[FileType, type[MultiModalPromptMessageContent]] = {
prompt_class_map: Mapping[FileType, type[PromptMessageContentUnionTypes]] = {
FileType.IMAGE: ImagePromptMessageContent,
FileType.AUDIO: AudioPromptMessageContent,
FileType.VIDEO: VideoPromptMessageContent,