diff --git a/api/core/model_runtime/entities/message_entities.py b/api/core/model_runtime/entities/message_entities.py index 0d104204bf..579e68c8a1 100644 --- a/api/core/model_runtime/entities/message_entities.py +++ b/api/core/model_runtime/entities/message_entities.py @@ -123,6 +123,7 @@ PromptMessageContentUnionTypes = Annotated[ Field(discriminator="type"), ] + class PromptMessage(BaseModel): """ Model class for prompt message. diff --git a/api/tests/unit_tests/core/prompt/test_prompt_message.py b/api/tests/unit_tests/core/prompt/test_prompt_message.py index 5de713b81a..e5da51d733 100644 --- a/api/tests/unit_tests/core/prompt/test_prompt_message.py +++ b/api/tests/unit_tests/core/prompt/test_prompt_message.py @@ -24,4 +24,4 @@ def test_dump_prompt_message(): ] ) data = prompt.model_dump() - assert data["content"][0].get("url") == example_url \ No newline at end of file + assert data["content"][0].get("url") == example_url