import hashlib import logging from threading import Thread, Timer from typing import Union from flask import Flask, current_app from sqlalchemy import select from sqlalchemy.orm import Session from configs import dify_config from core.app.entities.app_invoke_entities import ( AdvancedChatAppGenerateEntity, AgentChatAppGenerateEntity, ChatAppGenerateEntity, CompletionAppGenerateEntity, ) from core.app.entities.queue_entities import ( QueueAnnotationReplyEvent, QueueMessageFileEvent, QueueRetrieverResourcesEvent, ) from core.app.entities.task_entities import ( AnnotationReply, AnnotationReplyAccount, EasyUITaskState, MessageFileStreamResponse, MessageReplaceStreamResponse, MessageStreamResponse, StreamEvent, WorkflowTaskState, ) from core.db.session_factory import session_factory from core.llm_generator.llm_generator import LLMGenerator from core.tools.signature import sign_tool_file from extensions.ext_database import db from extensions.ext_redis import redis_client from models.enums import MessageFileBelongsTo from models.model import AppMode, Conversation, MessageAnnotation, MessageFile from services.annotation_service import AppAnnotationService logger = logging.getLogger(__name__) class MessageCycleManager: def __init__( self, *, application_generate_entity: Union[ ChatAppGenerateEntity, CompletionAppGenerateEntity, AgentChatAppGenerateEntity, AdvancedChatAppGenerateEntity, ], task_state: Union[EasyUITaskState, WorkflowTaskState], ): self._application_generate_entity = application_generate_entity self._task_state = task_state self._message_has_file: set[str] = set() def get_message_event_type(self, message_id: str) -> StreamEvent: # Fast path: cached determination from prior QueueMessageFileEvent if message_id in self._message_has_file: return StreamEvent.MESSAGE_FILE # Use SQLAlchemy 2.x style session.scalar(select(...)) with session_factory.create_session() as session: message_file = session.scalar( select(MessageFile) .where( MessageFile.message_id == message_id, ) .where(MessageFile.belongs_to == "assistant") ) if message_file: self._message_has_file.add(message_id) return StreamEvent.MESSAGE_FILE return StreamEvent.MESSAGE def generate_conversation_name(self, *, conversation_id: str, query: str) -> Thread | None: """ Generate conversation name. :param conversation_id: conversation id :param query: query :return: thread """ if isinstance(self._application_generate_entity, CompletionAppGenerateEntity): return None is_first_message = self._application_generate_entity.is_new_conversation extras = self._application_generate_entity.extras auto_generate_conversation_name = extras.get("auto_generate_conversation_name", True) thread: Thread | None = None if auto_generate_conversation_name and is_first_message: # start generate thread # time.sleep not block other logic thread = Timer( 1, self._generate_conversation_name_worker, kwargs={ "flask_app": current_app._get_current_object(), # type: ignore "conversation_id": conversation_id, "query": query, }, ) thread.daemon = True thread.start() if is_first_message: self._application_generate_entity.is_new_conversation = False return thread def _generate_conversation_name_worker(self, flask_app: Flask, conversation_id: str, query: str): with flask_app.app_context(): # get conversation and message stmt = select(Conversation).where(Conversation.id == conversation_id) conversation = db.session.scalar(stmt) if not conversation: return if conversation.mode != AppMode.COMPLETION: app_model = conversation.app if not app_model: return # generate conversation name query_hash = hashlib.md5(query.encode()).hexdigest()[:16] cache_key = f"conv_name:{conversation_id}:{query_hash}" cached_name = redis_client.get(cache_key) if cached_name: name = cached_name.decode("utf-8") else: try: name = LLMGenerator.generate_conversation_name( app_model.tenant_id, query, conversation_id, conversation.app_id ) redis_client.setex(cache_key, 3600, name) except Exception: if dify_config.DEBUG: logger.exception("generate conversation name failed, conversation_id: %s", conversation_id) name = query[:47] + "..." if len(query) > 50 else query conversation.name = name db.session.commit() db.session.close() def handle_annotation_reply(self, event: QueueAnnotationReplyEvent) -> MessageAnnotation | None: """ Handle annotation reply. :param event: event :return: """ annotation = AppAnnotationService.get_annotation_by_id(event.message_annotation_id) if annotation: account = annotation.account self._task_state.metadata.annotation_reply = AnnotationReply( id=annotation.id, account=AnnotationReplyAccount( id=annotation.account_id, name=account.name if account else "Dify user", ), ) return annotation return None def handle_retriever_resources(self, event: QueueRetrieverResourcesEvent): """ Handle retriever resources. :param event: event :return: """ if not self._application_generate_entity.app_config.additional_features: raise ValueError("Additional features not found") if self._application_generate_entity.app_config.additional_features.show_retrieve_source: merged_resources = [r for r in self._task_state.metadata.retriever_resources or [] if r] existing_ids = {(r.dataset_id, r.document_id) for r in merged_resources if r.dataset_id and r.document_id} # Add new unique resources from the event for resource in event.retriever_resources or []: if not resource: continue is_duplicate = ( resource.dataset_id and resource.document_id and (resource.dataset_id, resource.document_id) in existing_ids ) if not is_duplicate: merged_resources.append(resource) for i, resource in enumerate(merged_resources, 1): resource.position = i self._task_state.metadata.retriever_resources = merged_resources def message_file_to_stream_response(self, event: QueueMessageFileEvent) -> MessageFileStreamResponse | None: """ Message file to stream response. :param event: event :return: """ with Session(db.engine, expire_on_commit=False) as session: message_file = session.scalar(select(MessageFile).where(MessageFile.id == event.message_file_id)) if message_file and message_file.url is not None: self._message_has_file.add(message_file.message_id) # get tool file id tool_file_id = message_file.url.split("/")[-1] # trim extension tool_file_id = tool_file_id.split(".")[0] # get extension if "." in message_file.url: extension = f".{message_file.url.split('.')[-1]}" if len(extension) > 10: extension = ".bin" else: extension = ".bin" # add sign url to local file if message_file.url.startswith("http"): url = message_file.url else: url = sign_tool_file(tool_file_id=tool_file_id, extension=extension) return MessageFileStreamResponse( task_id=self._application_generate_entity.task_id, id=message_file.id, type=message_file.type, belongs_to=message_file.belongs_to or MessageFileBelongsTo.USER, url=url, ) return None def message_to_stream_response( self, answer: str, message_id: str, from_variable_selector: list[str] | None = None, event_type: StreamEvent | None = None, ) -> MessageStreamResponse: """ Message to stream response. :param answer: answer :param message_id: message id :return: """ return MessageStreamResponse( task_id=self._application_generate_entity.task_id, id=message_id, answer=answer, from_variable_selector=from_variable_selector, event=event_type or StreamEvent.MESSAGE, ) def message_replace_to_stream_response(self, answer: str, reason: str = "") -> MessageReplaceStreamResponse: """ Message replace to stream response. :param answer: answer :return: """ return MessageReplaceStreamResponse( task_id=self._application_generate_entity.task_id, answer=answer, reason=reason )