Underground-Digital/Workflow-Engine

1import logging2 3from flask_login import current_user4from flask_restful import Resource, fields, marshal_with, reqparse5from flask_restful.inputs import int_range6from werkzeug.exceptions import Forbidden, InternalServerError, NotFound7 8from controllers.console import api9from controllers.console.app.error import (10    CompletionRequestError,11    ProviderModelCurrentlyNotSupportError,12    ProviderNotInitializeError,13    ProviderQuotaExceededError,14)15from controllers.console.app.wraps import get_app_model16from controllers.console.explore.error import AppSuggestedQuestionsAfterAnswerDisabledError17from controllers.console.wraps import (18    account_initialization_required,19    cloud_edition_billing_resource_check,20    setup_required,21)22from core.app.entities.app_invoke_entities import InvokeFrom23from core.errors.error import ModelCurrentlyNotSupportError, ProviderTokenNotInitError, QuotaExceededError24from core.model_runtime.errors.invoke import InvokeError25from extensions.ext_database import db26from fields.conversation_fields import annotation_fields, message_detail_fields27from libs.helper import uuid_value28from libs.infinite_scroll_pagination import InfiniteScrollPagination29from libs.login import login_required30from models.model import AppMode, Conversation, Message, MessageAnnotation, MessageFeedback31from services.annotation_service import AppAnnotationService32from services.errors.conversation import ConversationNotExistsError33from services.errors.message import MessageNotExistsError, SuggestedQuestionsAfterAnswerDisabledError34from services.message_service import MessageService35 36 37class ChatMessageListApi(Resource):38    message_infinite_scroll_pagination_fields = {39        "limit": fields.Integer,40        "has_more": fields.Boolean,41        "data": fields.List(fields.Nested(message_detail_fields)),42    }43 44    @setup_required45    @login_required46    @get_app_model(mode=[AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT])47    @account_initialization_required48    @marshal_with(message_infinite_scroll_pagination_fields)49    def get(self, app_model):50        parser = reqparse.RequestParser()51        parser.add_argument("conversation_id", required=True, type=uuid_value, location="args")52        parser.add_argument("first_id", type=uuid_value, location="args")53        parser.add_argument("limit", type=int_range(1, 100), required=False, default=20, location="args")54        args = parser.parse_args()55 56        conversation = (57            db.session.query(Conversation)58            .filter(Conversation.id == args["conversation_id"], Conversation.app_id == app_model.id)59            .first()60        )61 62        if not conversation:63            raise NotFound("Conversation Not Exists.")64 65        if args["first_id"]:66            first_message = (67                db.session.query(Message)68                .filter(Message.conversation_id == conversation.id, Message.id == args["first_id"])69                .first()70            )71 72            if not first_message:73                raise NotFound("First message not found")74 75            history_messages = (76                db.session.query(Message)77                .filter(78                    Message.conversation_id == conversation.id,79                    Message.created_at < first_message.created_at,80                    Message.id != first_message.id,81                )82                .order_by(Message.created_at.desc())83                .limit(args["limit"])84                .all()85            )86        else:87            history_messages = (88                db.session.query(Message)89                .filter(Message.conversation_id == conversation.id)90                .order_by(Message.created_at.desc())91                .limit(args["limit"])92                .all()93            )94 95        has_more = False96        if len(history_messages) == args["limit"]:97            current_page_first_message = history_messages[-1]98            rest_count = (99                db.session.query(Message)100                .filter(101                    Message.conversation_id == conversation.id,102                    Message.created_at < current_page_first_message.created_at,103                    Message.id != current_page_first_message.id,104                )105                .count()106            )107 108            if rest_count > 0:109                has_more = True110 111        history_messages = list(reversed(history_messages))112 113        return InfiniteScrollPagination(data=history_messages, limit=args["limit"], has_more=has_more)114 115 116class MessageFeedbackApi(Resource):117    @setup_required118    @login_required119    @account_initialization_required120    @get_app_model121    def post(self, app_model):122        parser = reqparse.RequestParser()123        parser.add_argument("message_id", required=True, type=uuid_value, location="json")124        parser.add_argument("rating", type=str, choices=["like", "dislike", None], location="json")125        args = parser.parse_args()126 127        message_id = str(args["message_id"])128 129        message = db.session.query(Message).filter(Message.id == message_id, Message.app_id == app_model.id).first()130 131        if not message:132            raise NotFound("Message Not Exists.")133 134        feedback = message.admin_feedback135 136        if not args["rating"] and feedback:137            db.session.delete(feedback)138        elif args["rating"] and feedback:139            feedback.rating = args["rating"]140        elif not args["rating"] and not feedback:141            raise ValueError("rating cannot be None when feedback not exists")142        else:143            feedback = MessageFeedback(144                app_id=app_model.id,145                conversation_id=message.conversation_id,146                message_id=message.id,147                rating=args["rating"],148                from_source="admin",149                from_account_id=current_user.id,150            )151            db.session.add(feedback)152 153        db.session.commit()154 155        return {"result": "success"}156 157 158class MessageAnnotationApi(Resource):159    @setup_required160    @login_required161    @account_initialization_required162    @cloud_edition_billing_resource_check("annotation")163    @get_app_model164    @marshal_with(annotation_fields)165    def post(self, app_model):166        if not current_user.is_editor:167            raise Forbidden()168 169        parser = reqparse.RequestParser()170        parser.add_argument("message_id", required=False, type=uuid_value, location="json")171        parser.add_argument("question", required=True, type=str, location="json")172        parser.add_argument("answer", required=True, type=str, location="json")173        parser.add_argument("annotation_reply", required=False, type=dict, location="json")174        args = parser.parse_args()175        annotation = AppAnnotationService.up_insert_app_annotation_from_message(args, app_model.id)176 177        return annotation178 179 180class MessageAnnotationCountApi(Resource):181    @setup_required182    @login_required183    @account_initialization_required184    @get_app_model185    def get(self, app_model):186        count = db.session.query(MessageAnnotation).filter(MessageAnnotation.app_id == app_model.id).count()187 188        return {"count": count}189 190 191class MessageSuggestedQuestionApi(Resource):192    @setup_required193    @login_required194    @account_initialization_required195    @get_app_model(mode=[AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT])196    def get(self, app_model, message_id):197        message_id = str(message_id)198 199        try:200            questions = MessageService.get_suggested_questions_after_answer(201                app_model=app_model, message_id=message_id, user=current_user, invoke_from=InvokeFrom.DEBUGGER202            )203        except MessageNotExistsError:204            raise NotFound("Message not found")205        except ConversationNotExistsError:206            raise NotFound("Conversation not found")207        except ProviderTokenNotInitError as ex:208            raise ProviderNotInitializeError(ex.description)209        except QuotaExceededError:210            raise ProviderQuotaExceededError()211        except ModelCurrentlyNotSupportError:212            raise ProviderModelCurrentlyNotSupportError()213        except InvokeError as e:214            raise CompletionRequestError(e.description)215        except SuggestedQuestionsAfterAnswerDisabledError:216            raise AppSuggestedQuestionsAfterAnswerDisabledError()217        except Exception:218            logging.exception("internal server error.")219            raise InternalServerError()220 221        return {"data": questions}222 223 224class MessageApi(Resource):225    @setup_required226    @login_required227    @account_initialization_required228    @get_app_model229    @marshal_with(message_detail_fields)230    def get(self, app_model, message_id):231        message_id = str(message_id)232 233        message = db.session.query(Message).filter(Message.id == message_id, Message.app_id == app_model.id).first()234 235        if not message:236            raise NotFound("Message Not Exists.")237 238        return message239 240 241api.add_resource(MessageSuggestedQuestionApi, "/apps/<uuid:app_id>/chat-messages/<uuid:message_id>/suggested-questions")242api.add_resource(ChatMessageListApi, "/apps/<uuid:app_id>/chat-messages", endpoint="console_chat_messages")243api.add_resource(MessageFeedbackApi, "/apps/<uuid:app_id>/feedbacks")244api.add_resource(MessageAnnotationApi, "/apps/<uuid:app_id>/annotations")245api.add_resource(MessageAnnotationCountApi, "/apps/<uuid:app_id>/annotations/count")246api.add_resource(MessageApi, "/apps/<uuid:app_id>/messages/<uuid:message_id>", endpoint="console_message")247