Underground-Digital/Workflow-Engine

1import logging2 3from flask_restful import Resource, fields, marshal_with, reqparse4from flask_restful.inputs import int_range5from werkzeug.exceptions import BadRequest, InternalServerError, NotFound6 7import services8from controllers.service_api import api9from controllers.service_api.app.error import NotChatAppError10from controllers.service_api.wraps import FetchUserArg, WhereisUserArg, validate_app_token11from core.app.entities.app_invoke_entities import InvokeFrom12from fields.conversation_fields import message_file_fields13from libs.helper import TimestampField, uuid_value14from models.model import App, AppMode, EndUser15from services.errors.message import SuggestedQuestionsAfterAnswerDisabledError16from services.message_service import MessageService17 18 19class MessageListApi(Resource):20    feedback_fields = {"rating": fields.String}21    retriever_resource_fields = {22        "id": fields.String,23        "message_id": fields.String,24        "position": fields.Integer,25        "dataset_id": fields.String,26        "dataset_name": fields.String,27        "document_id": fields.String,28        "document_name": fields.String,29        "data_source_type": fields.String,30        "segment_id": fields.String,31        "score": fields.Float,32        "hit_count": fields.Integer,33        "word_count": fields.Integer,34        "segment_position": fields.Integer,35        "index_node_hash": fields.String,36        "content": fields.String,37        "created_at": TimestampField,38    }39 40    agent_thought_fields = {41        "id": fields.String,42        "chain_id": fields.String,43        "message_id": fields.String,44        "position": fields.Integer,45        "thought": fields.String,46        "tool": fields.String,47        "tool_labels": fields.Raw,48        "tool_input": fields.String,49        "created_at": TimestampField,50        "observation": fields.String,51        "message_files": fields.List(fields.Nested(message_file_fields)),52    }53 54    message_fields = {55        "id": fields.String,56        "conversation_id": fields.String,57        "parent_message_id": fields.String,58        "inputs": fields.Raw,59        "query": fields.String,60        "answer": fields.String(attribute="re_sign_file_url_answer"),61        "message_files": fields.List(fields.Nested(message_file_fields)),62        "feedback": fields.Nested(feedback_fields, attribute="user_feedback", allow_null=True),63        "retriever_resources": fields.List(fields.Nested(retriever_resource_fields)),64        "created_at": TimestampField,65        "agent_thoughts": fields.List(fields.Nested(agent_thought_fields)),66        "status": fields.String,67        "error": fields.String,68    }69 70    message_infinite_scroll_pagination_fields = {71        "limit": fields.Integer,72        "has_more": fields.Boolean,73        "data": fields.List(fields.Nested(message_fields)),74    }75 76    @validate_app_token(fetch_user_arg=FetchUserArg(fetch_from=WhereisUserArg.QUERY))77    @marshal_with(message_infinite_scroll_pagination_fields)78    def get(self, app_model: App, end_user: EndUser):79        app_mode = AppMode.value_of(app_model.mode)80        if app_mode not in {AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT}:81            raise NotChatAppError()82 83        parser = reqparse.RequestParser()84        parser.add_argument("conversation_id", required=True, type=uuid_value, location="args")85        parser.add_argument("first_id", type=uuid_value, location="args")86        parser.add_argument("limit", type=int_range(1, 100), required=False, default=20, location="args")87        args = parser.parse_args()88 89        try:90            return MessageService.pagination_by_first_id(91                app_model, end_user, args["conversation_id"], args["first_id"], args["limit"]92            )93        except services.errors.conversation.ConversationNotExistsError:94            raise NotFound("Conversation Not Exists.")95        except services.errors.message.FirstMessageNotExistsError:96            raise NotFound("First Message Not Exists.")97 98 99class MessageFeedbackApi(Resource):100    @validate_app_token(fetch_user_arg=FetchUserArg(fetch_from=WhereisUserArg.JSON, required=True))101    def post(self, app_model: App, end_user: EndUser, message_id):102        message_id = str(message_id)103 104        parser = reqparse.RequestParser()105        parser.add_argument("rating", type=str, choices=["like", "dislike", None], location="json")106        args = parser.parse_args()107 108        try:109            MessageService.create_feedback(app_model, message_id, end_user, args["rating"])110        except services.errors.message.MessageNotExistsError:111            raise NotFound("Message Not Exists.")112 113        return {"result": "success"}114 115 116class MessageSuggestedApi(Resource):117    @validate_app_token(fetch_user_arg=FetchUserArg(fetch_from=WhereisUserArg.QUERY, required=True))118    def get(self, app_model: App, end_user: EndUser, message_id):119        message_id = str(message_id)120        app_mode = AppMode.value_of(app_model.mode)121        if app_mode not in {AppMode.CHAT, AppMode.AGENT_CHAT, AppMode.ADVANCED_CHAT}:122            raise NotChatAppError()123 124        try:125            questions = MessageService.get_suggested_questions_after_answer(126                app_model=app_model, user=end_user, message_id=message_id, invoke_from=InvokeFrom.SERVICE_API127            )128        except services.errors.message.MessageNotExistsError:129            raise NotFound("Message Not Exists.")130        except SuggestedQuestionsAfterAnswerDisabledError:131            raise BadRequest("Suggested Questions Is Disabled.")132        except Exception:133            logging.exception("internal server error.")134            raise InternalServerError()135 136        return {"result": "success", "data": questions}137 138 139api.add_resource(MessageListApi, "/messages")140api.add_resource(MessageFeedbackApi, "/messages/<uuid:message_id>/feedbacks")141api.add_resource(MessageSuggestedApi, "/messages/<uuid:message_id>/suggested")142