Underground-Digital/Workflow-Engine

1import logging2from threading import Thread3from typing import Optional, Union4 5from flask import Flask, current_app6 7from configs import dify_config8from core.app.entities.app_invoke_entities import (9    AdvancedChatAppGenerateEntity,10    AgentChatAppGenerateEntity,11    ChatAppGenerateEntity,12    CompletionAppGenerateEntity,13)14from core.app.entities.queue_entities import (15    QueueAnnotationReplyEvent,16    QueueMessageFileEvent,17    QueueRetrieverResourcesEvent,18)19from core.app.entities.task_entities import (20    EasyUITaskState,21    MessageFileStreamResponse,22    MessageReplaceStreamResponse,23    MessageStreamResponse,24    WorkflowTaskState,25)26from core.llm_generator.llm_generator import LLMGenerator27from core.tools.tool_file_manager import ToolFileManager28from extensions.ext_database import db29from models.model import AppMode, Conversation, MessageAnnotation, MessageFile30from services.annotation_service import AppAnnotationService31 32 33class MessageCycleManage:34    _application_generate_entity: Union[35        ChatAppGenerateEntity, CompletionAppGenerateEntity, AgentChatAppGenerateEntity, AdvancedChatAppGenerateEntity36    ]37    _task_state: Union[EasyUITaskState, WorkflowTaskState]38 39    def _generate_conversation_name(self, conversation: Conversation, query: str) -> Optional[Thread]:40        """41        Generate conversation name.42        :param conversation: conversation43        :param query: query44        :return: thread45        """46        if isinstance(self._application_generate_entity, CompletionAppGenerateEntity):47            return None48 49        is_first_message = self._application_generate_entity.conversation_id is None50        extras = self._application_generate_entity.extras51        auto_generate_conversation_name = extras.get("auto_generate_conversation_name", True)52 53        if auto_generate_conversation_name and is_first_message:54            # start generate thread55            thread = Thread(56                target=self._generate_conversation_name_worker,57                kwargs={58                    "flask_app": current_app._get_current_object(),  # type: ignore59                    "conversation_id": conversation.id,60                    "query": query,61                },62            )63 64            thread.start()65 66            return thread67 68        return None69 70    def _generate_conversation_name_worker(self, flask_app: Flask, conversation_id: str, query: str):71        with flask_app.app_context():72            # get conversation and message73            conversation = db.session.query(Conversation).filter(Conversation.id == conversation_id).first()74 75            if not conversation:76                return77 78            if conversation.mode != AppMode.COMPLETION.value:79                app_model = conversation.app80                if not app_model:81                    return82 83                # generate conversation name84                try:85                    name = LLMGenerator.generate_conversation_name(app_model.tenant_id, query)86                    conversation.name = name87                except Exception as e:88                    if dify_config.DEBUG:89                        logging.exception(f"generate conversation name failed: {e}")90                    pass91 92                db.session.merge(conversation)93                db.session.commit()94                db.session.close()95 96    def _handle_annotation_reply(self, event: QueueAnnotationReplyEvent) -> Optional[MessageAnnotation]:97        """98        Handle annotation reply.99        :param event: event100        :return:101        """102        annotation = AppAnnotationService.get_annotation_by_id(event.message_annotation_id)103        if annotation:104            account = annotation.account105            self._task_state.metadata["annotation_reply"] = {106                "id": annotation.id,107                "account": {"id": annotation.account_id, "name": account.name if account else "Dify user"},108            }109 110            return annotation111 112        return None113 114    def _handle_retriever_resources(self, event: QueueRetrieverResourcesEvent) -> None:115        """116        Handle retriever resources.117        :param event: event118        :return:119        """120        if self._application_generate_entity.app_config.additional_features.show_retrieve_source:121            self._task_state.metadata["retriever_resources"] = event.retriever_resources122 123    def _message_file_to_stream_response(self, event: QueueMessageFileEvent) -> Optional[MessageFileStreamResponse]:124        """125        Message file to stream response.126        :param event: event127        :return:128        """129        message_file = db.session.query(MessageFile).filter(MessageFile.id == event.message_file_id).first()130 131        if message_file:132            # get tool file id133            tool_file_id = message_file.url.split("/")[-1]134            # trim extension135            tool_file_id = tool_file_id.split(".")[0]136 137            # get extension138            if "." in message_file.url:139                extension = f'.{message_file.url.split(".")[-1]}'140                if len(extension) > 10:141                    extension = ".bin"142            else:143                extension = ".bin"144            # add sign url to local file145            if message_file.url.startswith("http"):146                url = message_file.url147            else:148                url = ToolFileManager.sign_file(tool_file_id=tool_file_id, extension=extension)149 150            return MessageFileStreamResponse(151                task_id=self._application_generate_entity.task_id,152                id=message_file.id,153                type=message_file.type,154                belongs_to=message_file.belongs_to or "user",155                url=url,156            )157 158        return None159 160    def _message_to_stream_response(161        self, answer: str, message_id: str, from_variable_selector: Optional[list[str]] = None162    ) -> MessageStreamResponse:163        """164        Message to stream response.165        :param answer: answer166        :param message_id: message id167        :return:168        """169        return MessageStreamResponse(170            task_id=self._application_generate_entity.task_id,171            id=message_id,172            answer=answer,173            from_variable_selector=from_variable_selector,174        )175 176    def _message_replace_to_stream_response(self, answer: str) -> MessageReplaceStreamResponse:177        """178        Message replace to stream response.179        :param answer: answer180        :return:181        """182        return MessageReplaceStreamResponse(task_id=self._application_generate_entity.task_id, answer=answer)183