Underground-Digital/Workflow-Engine

1import json2from typing import Optional, Union3 4from core.app.apps.advanced_chat.app_config_manager import AdvancedChatAppConfigManager5from core.app.entities.app_invoke_entities import InvokeFrom6from core.llm_generator.llm_generator import LLMGenerator7from core.memory.token_buffer_memory import TokenBufferMemory8from core.model_manager import ModelManager9from core.model_runtime.entities.model_entities import ModelType10from core.ops.entities.trace_entity import TraceTaskName11from core.ops.ops_trace_manager import TraceQueueManager, TraceTask12from core.ops.utils import measure_time13from extensions.ext_database import db14from libs.infinite_scroll_pagination import InfiniteScrollPagination15from models.account import Account16from models.model import App, AppMode, AppModelConfig, EndUser, Message, MessageFeedback17from services.conversation_service import ConversationService18from services.errors.conversation import ConversationCompletedError, ConversationNotExistsError19from services.errors.message import (20    FirstMessageNotExistsError,21    LastMessageNotExistsError,22    MessageNotExistsError,23    SuggestedQuestionsAfterAnswerDisabledError,24)25from services.workflow_service import WorkflowService26 27 28class MessageService:29    @classmethod30    def pagination_by_first_id(31        cls,32        app_model: App,33        user: Optional[Union[Account, EndUser]],34        conversation_id: str,35        first_id: Optional[str],36        limit: int,37        order: str = "asc",38    ) -> InfiniteScrollPagination:39        if not user:40            return InfiniteScrollPagination(data=[], limit=limit, has_more=False)41 42        if not conversation_id:43            return InfiniteScrollPagination(data=[], limit=limit, has_more=False)44 45        conversation = ConversationService.get_conversation(46            app_model=app_model, user=user, conversation_id=conversation_id47        )48 49        if first_id:50            first_message = (51                db.session.query(Message)52                .filter(Message.conversation_id == conversation.id, Message.id == first_id)53                .first()54            )55 56            if not first_message:57                raise FirstMessageNotExistsError()58 59            history_messages = (60                db.session.query(Message)61                .filter(62                    Message.conversation_id == conversation.id,63                    Message.created_at < first_message.created_at,64                    Message.id != first_message.id,65                )66                .order_by(Message.created_at.desc())67                .limit(limit)68                .all()69            )70        else:71            history_messages = (72                db.session.query(Message)73                .filter(Message.conversation_id == conversation.id)74                .order_by(Message.created_at.desc())75                .limit(limit)76                .all()77            )78 79        has_more = False80        if len(history_messages) == limit:81            current_page_first_message = history_messages[-1]82            rest_count = (83                db.session.query(Message)84                .filter(85                    Message.conversation_id == conversation.id,86                    Message.created_at < current_page_first_message.created_at,87                    Message.id != current_page_first_message.id,88                )89                .count()90            )91 92            if rest_count > 0:93                has_more = True94 95        if order == "asc":96            history_messages = list(reversed(history_messages))97 98        return InfiniteScrollPagination(data=history_messages, limit=limit, has_more=has_more)99 100    @classmethod101    def pagination_by_last_id(102        cls,103        app_model: App,104        user: Optional[Union[Account, EndUser]],105        last_id: Optional[str],106        limit: int,107        conversation_id: Optional[str] = None,108        include_ids: Optional[list] = None,109    ) -> InfiniteScrollPagination:110        if not user:111            return InfiniteScrollPagination(data=[], limit=limit, has_more=False)112 113        base_query = db.session.query(Message)114 115        if conversation_id is not None:116            conversation = ConversationService.get_conversation(117                app_model=app_model, user=user, conversation_id=conversation_id118            )119 120            base_query = base_query.filter(Message.conversation_id == conversation.id)121 122        if include_ids is not None:123            base_query = base_query.filter(Message.id.in_(include_ids))124 125        if last_id:126            last_message = base_query.filter(Message.id == last_id).first()127 128            if not last_message:129                raise LastMessageNotExistsError()130 131            history_messages = (132                base_query.filter(Message.created_at < last_message.created_at, Message.id != last_message.id)133                .order_by(Message.created_at.desc())134                .limit(limit)135                .all()136            )137        else:138            history_messages = base_query.order_by(Message.created_at.desc()).limit(limit).all()139 140        has_more = False141        if len(history_messages) == limit:142            current_page_first_message = history_messages[-1]143            rest_count = base_query.filter(144                Message.created_at < current_page_first_message.created_at, Message.id != current_page_first_message.id145            ).count()146 147            if rest_count > 0:148                has_more = True149 150        return InfiniteScrollPagination(data=history_messages, limit=limit, has_more=has_more)151 152    @classmethod153    def create_feedback(154        cls, app_model: App, message_id: str, user: Optional[Union[Account, EndUser]], rating: Optional[str]155    ) -> MessageFeedback:156        if not user:157            raise ValueError("user cannot be None")158 159        message = cls.get_message(app_model=app_model, user=user, message_id=message_id)160 161        feedback = message.user_feedback if isinstance(user, EndUser) else message.admin_feedback162 163        if not rating and feedback:164            db.session.delete(feedback)165        elif rating and feedback:166            feedback.rating = rating167        elif not rating and not feedback:168            raise ValueError("rating cannot be None when feedback not exists")169        else:170            feedback = MessageFeedback(171                app_id=app_model.id,172                conversation_id=message.conversation_id,173                message_id=message.id,174                rating=rating,175                from_source=("user" if isinstance(user, EndUser) else "admin"),176                from_end_user_id=(user.id if isinstance(user, EndUser) else None),177                from_account_id=(user.id if isinstance(user, Account) else None),178            )179            db.session.add(feedback)180 181        db.session.commit()182 183        return feedback184 185    @classmethod186    def get_message(cls, app_model: App, user: Optional[Union[Account, EndUser]], message_id: str):187        message = (188            db.session.query(Message)189            .filter(190                Message.id == message_id,191                Message.app_id == app_model.id,192                Message.from_source == ("api" if isinstance(user, EndUser) else "console"),193                Message.from_end_user_id == (user.id if isinstance(user, EndUser) else None),194                Message.from_account_id == (user.id if isinstance(user, Account) else None),195            )196            .first()197        )198 199        if not message:200            raise MessageNotExistsError()201 202        return message203 204    @classmethod205    def get_suggested_questions_after_answer(206        cls, app_model: App, user: Optional[Union[Account, EndUser]], message_id: str, invoke_from: InvokeFrom207    ) -> list[Message]:208        if not user:209            raise ValueError("user cannot be None")210 211        message = cls.get_message(app_model=app_model, user=user, message_id=message_id)212 213        conversation = ConversationService.get_conversation(214            app_model=app_model, conversation_id=message.conversation_id, user=user215        )216 217        if not conversation:218            raise ConversationNotExistsError()219 220        if conversation.status != "normal":221            raise ConversationCompletedError()222 223        model_manager = ModelManager()224 225        if app_model.mode == AppMode.ADVANCED_CHAT.value:226            workflow_service = WorkflowService()227            if invoke_from == InvokeFrom.DEBUGGER:228                workflow = workflow_service.get_draft_workflow(app_model=app_model)229            else:230                workflow = workflow_service.get_published_workflow(app_model=app_model)231 232            if workflow is None:233                return []234 235            app_config = AdvancedChatAppConfigManager.get_app_config(app_model=app_model, workflow=workflow)236 237            if not app_config.additional_features.suggested_questions_after_answer:238                raise SuggestedQuestionsAfterAnswerDisabledError()239 240            model_instance = model_manager.get_default_model_instance(241                tenant_id=app_model.tenant_id, model_type=ModelType.LLM242            )243        else:244            if not conversation.override_model_configs:245                app_model_config = (246                    db.session.query(AppModelConfig)247                    .filter(248                        AppModelConfig.id == conversation.app_model_config_id, AppModelConfig.app_id == app_model.id249                    )250                    .first()251                )252            else:253                conversation_override_model_configs = json.loads(conversation.override_model_configs)254                app_model_config = AppModelConfig(255                    id=conversation.app_model_config_id,256                    app_id=app_model.id,257                )258 259                app_model_config = app_model_config.from_model_config_dict(conversation_override_model_configs)260 261            suggested_questions_after_answer = app_model_config.suggested_questions_after_answer_dict262            if suggested_questions_after_answer.get("enabled", False) is False:263                raise SuggestedQuestionsAfterAnswerDisabledError()264 265            model_instance = model_manager.get_model_instance(266                tenant_id=app_model.tenant_id,267                provider=app_model_config.model_dict["provider"],268                model_type=ModelType.LLM,269                model=app_model_config.model_dict["name"],270            )271 272        # get memory of conversation (read-only)273        memory = TokenBufferMemory(conversation=conversation, model_instance=model_instance)274 275        histories = memory.get_history_prompt_text(276            max_token_limit=3000,277            message_limit=3,278        )279 280        with measure_time() as timer:281            questions = LLMGenerator.generate_suggested_questions_after_answer(282                tenant_id=app_model.tenant_id, histories=histories283            )284 285        # get tracing instance286        trace_manager = TraceQueueManager(app_id=app_model.id)287        trace_manager.add_trace_task(288            TraceTask(289                TraceTaskName.SUGGESTED_QUESTION_TRACE, message_id=message_id, suggested_question=questions, timer=timer290            )291        )292 293        return questions294