From 9a0e2e8f4271d9f4e72e9ead74804e1f2382d9e8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9C=B1=E5=BA=86=E8=B6=85?= <zhuqingchao@xiaomi.com>
Date: Sun, 30 Mar 2025 19:18:25 +0800
Subject: [PATCH] ix: Update prompt message content types to use Literal and
 add union type for content

---
 .../entities/message_entities.py              | 32 +++++++++++--------
 .../core/prompt/test_prompt_message.py        | 23 +++++++++++++
 2 files changed, 42 insertions(+), 13 deletions(-)
 create mode 100644 api/tests/unit_tests/core/prompt/test_prompt_message.py

diff --git a/api/core/model_runtime/entities/message_entities.py b/api/core/model_runtime/entities/message_entities.py
index 3bed2460dd..c307fc5515 100644
--- a/api/core/model_runtime/entities/message_entities.py
+++ b/api/core/model_runtime/entities/message_entities.py
@@ -1,6 +1,6 @@
 from collections.abc import Sequence
 from enum import Enum, StrEnum
-from typing import Any, Optional, Union
+from typing import Any, Annotated, Literal, Optional, Union, Union
 
 from pydantic import BaseModel, Field, field_serializer, field_validator
 
@@ -61,11 +61,7 @@ class PromptMessageContentType(StrEnum):
 
 
 class PromptMessageContent(BaseModel):
-    """
-    Model class for prompt message content.
-    """
-
-    type: PromptMessageContentType
+    pass
 
 
 class TextPromptMessageContent(PromptMessageContent):
@@ -73,7 +69,7 @@ class TextPromptMessageContent(PromptMessageContent):
     Model class for text prompt message content.
     """
 
-    type: PromptMessageContentType = PromptMessageContentType.TEXT
+    type: Literal[PromptMessageContentType.TEXT] = PromptMessageContentType.TEXT
     data: str
 
 
@@ -82,7 +78,6 @@ class MultiModalPromptMessageContent(PromptMessageContent):
     Model class for multi-modal prompt message content.
     """
 
-    type: PromptMessageContentType
     format: str = Field(default=..., description="the format of multi-modal file")
     base64_data: str = Field(default="", description="the base64 data of multi-modal file")
     url: str = Field(default="", description="the url of multi-modal file")
@@ -94,11 +89,11 @@ class MultiModalPromptMessageContent(PromptMessageContent):
 
 
 class VideoPromptMessageContent(MultiModalPromptMessageContent):
-    type: PromptMessageContentType = PromptMessageContentType.VIDEO
+    type: Literal[PromptMessageContentType.VIDEO] = PromptMessageContentType.VIDEO
 
 
 class AudioPromptMessageContent(MultiModalPromptMessageContent):
-    type: PromptMessageContentType = PromptMessageContentType.AUDIO
+    type: Literal[PromptMessageContentType.AUDIO] = PromptMessageContentType.AUDIO
 
 
 class ImagePromptMessageContent(MultiModalPromptMessageContent):
@@ -110,13 +105,24 @@ class ImagePromptMessageContent(MultiModalPromptMessageContent):
         LOW = "low"
         HIGH = "high"
 
-    type: PromptMessageContentType = PromptMessageContentType.IMAGE
+    type: Literal[PromptMessageContentType.IMAGE] = PromptMessageContentType.IMAGE
     detail: DETAIL = DETAIL.LOW
 
 
 class DocumentPromptMessageContent(MultiModalPromptMessageContent):
-    type: PromptMessageContentType = PromptMessageContentType.DOCUMENT
+    type: Literal[PromptMessageContentType.DOCUMENT] = PromptMessageContentType.DOCUMENT
+
 
+PromptMessageContentUnionTypes = Annotated[
+    Union[
+        TextPromptMessageContent,
+        ImagePromptMessageContent,
+        DocumentPromptMessageContent,
+        AudioPromptMessageContent,
+        VideoPromptMessageContent,
+    ],
+    Field(discriminator="type"),
+]
 
 class PromptMessage(BaseModel):
     """
@@ -124,7 +130,7 @@ class PromptMessage(BaseModel):
     """
 
     role: PromptMessageRole
-    content: Optional[str | Sequence[PromptMessageContent]] = None
+    content: Optional[str | list[PromptMessageContentUnionTypes]] = None
     name: Optional[str] = None
 
     def is_empty(self) -> bool:
diff --git a/api/tests/unit_tests/core/prompt/test_prompt_message.py b/api/tests/unit_tests/core/prompt/test_prompt_message.py
new file mode 100644
index 0000000000..4e93ad7957
--- /dev/null
+++ b/api/tests/unit_tests/core/prompt/test_prompt_message.py
@@ -0,0 +1,23 @@
+from core.model_runtime.entities.message_entities import ImagePromptMessageContent, TextPromptMessageContent, UserPromptMessage
+
+
+def test_build_prompt_message_with_prompt_message_contents():
+    prompt = UserPromptMessage(content=[TextPromptMessageContent(data="Hello, World!")])
+    assert isinstance(prompt.content, list)
+    assert isinstance(prompt.content[0], TextPromptMessageContent)
+    assert prompt.content[0].data == "Hello, World!"
+
+
+def test_dump_prompt_message():
+    example_url = "https://example.com/image.jpg"
+    prompt = UserPromptMessage(
+        content=[
+            ImagePromptMessageContent(
+                url=example_url,
+                format="jpeg",
+                mime_type="image/jpeg",
+            )
+        ]
+    )
+    data = prompt.model_dump()
+    assert data["content"][0].get("url") == example_url
\ No newline at end of file