refactor: more ns.model to BaseModel (#30445)

api/controllers/console/explore/conversation.py

@@ -1,8 +1,7 @@
 from typing import Any
 
 from flask import request
-from flask_restx import marshal_with
-from pydantic import BaseModel, Field, model_validator
+from pydantic import BaseModel, Field, TypeAdapter, model_validator
 from sqlalchemy.orm import Session
 from werkzeug.exceptions import NotFound
 
@@ -11,7 +10,11 @@ from controllers.console.explore.error import NotChatAppError
 from controllers.console.explore.wraps import InstalledAppResource
 from core.app.entities.app_invoke_entities import InvokeFrom
 from extensions.ext_database import db
-from fields.conversation_fields import conversation_infinite_scroll_pagination_fields, simple_conversation_fields
+from fields.conversation_fields import (
+    ConversationInfiniteScrollPagination,
+    ResultResponse,
+    SimpleConversation,
+)
 from libs.helper import UUIDStrOrEmpty
 from libs.login import current_user
 from models import Account
@@ -49,7 +52,6 @@ register_schema_models(console_ns, ConversationListQuery, ConversationRenamePayl
     endpoint="installed_app_conversations",
 )
 class ConversationListApi(InstalledAppResource):
-    @marshal_with(conversation_infinite_scroll_pagination_fields)
     @console_ns.expect(console_ns.models[ConversationListQuery.__name__])
     def get(self, installed_app):
         app_model = installed_app.app
@@ -73,7 +75,7 @@ class ConversationListApi(InstalledAppResource):
             if not isinstance(current_user, Account):
                 raise ValueError("current_user must be an Account instance")
             with Session(db.engine) as session:
-                return WebConversationService.pagination_by_last_id(
+                pagination = WebConversationService.pagination_by_last_id(
                     session=session,
                     app_model=app_model,
                     user=current_user,
@@ -82,6 +84,13 @@ class ConversationListApi(InstalledAppResource):
                     invoke_from=InvokeFrom.EXPLORE,
                     pinned=args.pinned,
                 )
+                adapter = TypeAdapter(SimpleConversation)
+                conversations = [adapter.validate_python(item, from_attributes=True) for item in pagination.data]
+                return ConversationInfiniteScrollPagination(
+                    limit=pagination.limit,
+                    has_more=pagination.has_more,
+                    data=conversations,
+                ).model_dump(mode="json")
         except LastConversationNotExistsError:
             raise NotFound("Last Conversation Not Exists.")
 
@@ -105,7 +114,7 @@ class ConversationApi(InstalledAppResource):
         except ConversationNotExistsError:
             raise NotFound("Conversation Not Exists.")
 
-        return {"result": "success"}, 204
+        return ResultResponse(result="success").model_dump(mode="json"), 204
 
 
 @console_ns.route(
@@ -113,7 +122,6 @@ class ConversationApi(InstalledAppResource):
     endpoint="installed_app_conversation_rename",
 )
 class ConversationRenameApi(InstalledAppResource):
-    @marshal_with(simple_conversation_fields)
     @console_ns.expect(console_ns.models[ConversationRenamePayload.__name__])
     def post(self, installed_app, c_id):
         app_model = installed_app.app
@@ -128,9 +136,14 @@ class ConversationRenameApi(InstalledAppResource):
         try:
             if not isinstance(current_user, Account):
                 raise ValueError("current_user must be an Account instance")
-            return ConversationService.rename(
+            conversation = ConversationService.rename(
                 app_model, conversation_id, current_user, payload.name, payload.auto_generate
             )
+            return (
+                TypeAdapter(SimpleConversation)
+                .validate_python(conversation, from_attributes=True)
+                .model_dump(mode="json")
+            )
         except ConversationNotExistsError:
             raise NotFound("Conversation Not Exists.")
 
@@ -155,7 +168,7 @@ class ConversationPinApi(InstalledAppResource):
         except ConversationNotExistsError:
             raise NotFound("Conversation Not Exists.")
 
-        return {"result": "success"}
+        return ResultResponse(result="success").model_dump(mode="json")
 
 
 @console_ns.route(
@@ -174,4 +187,4 @@ class ConversationUnPinApi(InstalledAppResource):
             raise ValueError("current_user must be an Account instance")
         WebConversationService.unpin(app_model, conversation_id, current_user)
 
-        return {"result": "success"}
+        return ResultResponse(result="success").model_dump(mode="json")

api/controllers/console/explore/message.py

@@ -2,8 +2,7 @@ import logging
 from typing import Literal
 
 from flask import request
-from flask_restx import marshal_with
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, TypeAdapter
 from werkzeug.exceptions import InternalServerError, NotFound
 
 from controllers.common.schema import register_schema_models
@@ -23,7 +22,8 @@ from controllers.console.explore.wraps import InstalledAppResource
 from core.app.entities.app_invoke_entities import InvokeFrom
 from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError
 from core.model_runtime.errors.invoke import InvokeError
-from fields.message_fields import message_infinite_scroll_pagination_fields
+from fields.conversation_fields import ResultResponse
+from fields.message_fields import MessageInfiniteScrollPagination, MessageListItem, SuggestedQuestionsResponse
 from libs import helper
 from libs.helper import UUIDStrOrEmpty
 from libs.login import current_account_with_tenant
@@ -66,7 +66,6 @@ register_schema_models(console_ns, MessageListQuery, MessageFeedbackPayload, Mor
     endpoint="installed_app_messages",
 )
 class MessageListApi(InstalledAppResource):
-    @marshal_with(message_infinite_scroll_pagination_fields)
     @console_ns.expect(console_ns.models[MessageListQuery.__name__])
     def get(self, installed_app):
         current_user, _ = current_account_with_tenant()
@@ -78,13 +77,20 @@ class MessageListApi(InstalledAppResource):
         args = MessageListQuery.model_validate(request.args.to_dict())
 
         try:
-            return MessageService.pagination_by_first_id(
+            pagination = MessageService.pagination_by_first_id(
                 app_model,
                 current_user,
                 str(args.conversation_id),
                 str(args.first_id) if args.first_id else None,
                 args.limit,
             )
+            adapter = TypeAdapter(MessageListItem)
+            items = [adapter.validate_python(message, from_attributes=True) for message in pagination.data]
+            return MessageInfiniteScrollPagination(
+                limit=pagination.limit,
+                has_more=pagination.has_more,
+                data=items,
+            ).model_dump(mode="json")
         except ConversationNotExistsError:
             raise NotFound("Conversation Not Exists.")
         except FirstMessageNotExistsError:
@@ -116,7 +122,7 @@ class MessageFeedbackApi(InstalledAppResource):
         except MessageNotExistsError:
             raise NotFound("Message Not Exists.")
 
-        return {"result": "success"}
+        return ResultResponse(result="success").model_dump(mode="json")
 
 
 @console_ns.route(
@@ -201,4 +207,4 @@ class MessageSuggestedQuestionApi(InstalledAppResource):
             logger.exception("internal server error.")
             raise InternalServerError()
 
-        return {"data": questions}
+        return SuggestedQuestionsResponse(data=questions).model_dump(mode="json")

api/controllers/console/explore/saved_message.py

@@ -1,14 +1,14 @@
 from flask import request
-from flask_restx import fields, marshal_with
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, TypeAdapter
 from werkzeug.exceptions import NotFound
 
 from controllers.common.schema import register_schema_models
 from controllers.console import console_ns
 from controllers.console.explore.error import NotCompletionAppError
 from controllers.console.explore.wraps import InstalledAppResource
-from fields.conversation_fields import message_file_fields
-from libs.helper import TimestampField, UUIDStrOrEmpty
+from fields.conversation_fields import ResultResponse
+from fields.message_fields import SavedMessageInfiniteScrollPagination, SavedMessageItem
+from libs.helper import UUIDStrOrEmpty
 from libs.login import current_account_with_tenant
 from services.errors.message import MessageNotExistsError
 from services.saved_message_service import SavedMessageService
@@ -26,28 +26,8 @@ class SavedMessageCreatePayload(BaseModel):
 register_schema_models(console_ns, SavedMessageListQuery, SavedMessageCreatePayload)
 
 
-feedback_fields = {"rating": fields.String}
-
-message_fields = {
-    "id": fields.String,
-    "inputs": fields.Raw,
-    "query": fields.String,
-    "answer": fields.String,
-    "message_files": fields.List(fields.Nested(message_file_fields)),
-    "feedback": fields.Nested(feedback_fields, attribute="user_feedback", allow_null=True),
-    "created_at": TimestampField,
-}
-
-
 @console_ns.route("/installed-apps/<uuid:installed_app_id>/saved-messages", endpoint="installed_app_saved_messages")
 class SavedMessageListApi(InstalledAppResource):
-    saved_message_infinite_scroll_pagination_fields = {
-        "limit": fields.Integer,
-        "has_more": fields.Boolean,
-        "data": fields.List(fields.Nested(message_fields)),
-    }
-
-    @marshal_with(saved_message_infinite_scroll_pagination_fields)
     @console_ns.expect(console_ns.models[SavedMessageListQuery.__name__])
     def get(self, installed_app):
         current_user, _ = current_account_with_tenant()
@@ -57,12 +37,19 @@ class SavedMessageListApi(InstalledAppResource):
 
         args = SavedMessageListQuery.model_validate(request.args.to_dict())
 
-        return SavedMessageService.pagination_by_last_id(
+        pagination = SavedMessageService.pagination_by_last_id(
             app_model,
             current_user,
             str(args.last_id) if args.last_id else None,
             args.limit,
         )
+        adapter = TypeAdapter(SavedMessageItem)
+        items = [adapter.validate_python(message, from_attributes=True) for message in pagination.data]
+        return SavedMessageInfiniteScrollPagination(
+            limit=pagination.limit,
+            has_more=pagination.has_more,
+            data=items,
+        ).model_dump(mode="json")
 
     @console_ns.expect(console_ns.models[SavedMessageCreatePayload.__name__])
     def post(self, installed_app):
@@ -78,7 +65,7 @@ class SavedMessageListApi(InstalledAppResource):
         except MessageNotExistsError:
             raise NotFound("Message Not Exists.")
 
-        return {"result": "success"}
+        return ResultResponse(result="success").model_dump(mode="json")
 
 
 @console_ns.route(
@@ -96,4 +83,4 @@ class SavedMessageApi(InstalledAppResource):
 
         SavedMessageService.delete(app_model, current_user, message_id)
 
-        return {"result": "success"}, 204
+        return ResultResponse(result="success").model_dump(mode="json"), 204