diff --git a/api/core/app/task_pipeline/based_generate_task_pipeline.py b/api/core/app/task_pipeline/based_generate_task_pipeline.py index 5331c0cc94..ed9710469b 100644 --- a/api/core/app/task_pipeline/based_generate_task_pipeline.py +++ b/api/core/app/task_pipeline/based_generate_task_pipeline.py @@ -19,7 +19,7 @@ from core.app.entities.task_entities import ( from core.errors.error import QuotaExceededError from core.model_runtime.errors.invoke import InvokeAuthorizationError, InvokeError from core.moderation.output_moderation import ModerationRule, OutputModeration -from models.model import Message +from models.model import Message, MessageStatus logger = logging.getLogger(__name__) @@ -62,7 +62,7 @@ class BasedGenerateTaskPipeline: return err err_desc = self._error_to_desc(err) - message.status = "error" + message.status = MessageStatus.ERROR message.error = err_desc return err diff --git a/api/models/model.py b/api/models/model.py index 93737043d5..a1959d8424 100644 --- a/api/models/model.py +++ b/api/models/model.py @@ -892,6 +892,15 @@ class Conversation(Base): } +class MessageStatus(StrEnum): + """ + Message Status Enum + """ + + NORMAL = "normal" + ERROR = "error" + + class Message(Base): __tablename__ = "messages" __table_args__ = ( diff --git a/api/services/audio_service.py b/api/services/audio_service.py index a259f5a4c4..ee3297e31f 100644 --- a/api/services/audio_service.py +++ b/api/services/audio_service.py @@ -8,7 +8,7 @@ from werkzeug.datastructures import FileStorage from constants import AUDIO_EXTENSIONS from core.model_manager import ModelManager from core.model_runtime.entities.model_entities import ModelType -from models.model import App, AppMode, AppModelConfig, Message +from models.model import App, AppMode, AppModelConfig, Message, MessageStatus from services.errors.audio import ( AudioTooLargeServiceError, NoAudioUploadedServiceError, @@ -132,7 +132,7 @@ class AudioService: message = db.session.query(Message).filter(Message.id == message_id).first() if message is None: return None - if message.answer == "" and message.status == "normal": + if message.answer == "" and message.status == MessageStatus.NORMAL: return None else: