diff --git a/apps/application/models/application_chat.py b/apps/application/models/application_chat.py index 16fb6d432ea..d617b65649a 100644 --- a/apps/application/models/application_chat.py +++ b/apps/application/models/application_chat.py @@ -24,6 +24,7 @@ class ChatUserType(models.TextChoices): SYSTEM_API_KEY = "SYSTEM_API_KEY", "系统API_KEY" APPLICATION_API_KEY = "APPLICATION_API_KEY", "应用API_KEY" PLATFORM_USER = "PLATFORM_USER", "平台用户" + SYSTEM_USER = "SYSTEM_USER", '系统用户' class ExecuteType(models.TextChoices): diff --git a/apps/application/urls.py b/apps/application/urls.py index 976fb1fbb0c..dc14accdbd5 100644 --- a/apps/application/urls.py +++ b/apps/application/urls.py @@ -44,6 +44,7 @@ path('workspace//application//mcp_tools', views.McpServers.as_view()), path('workspace//application//model//prompt_generate', views.PromptGenerateView.as_view()), path('chat_message/', views.ChatView.as_view()), + path('chat_message//cancel', views.CancelWorkflowView.as_view()), path('workspace//application//historical_conversation//', views.DebugHistoricalConversation.PageView.as_view()), path('workspace//application//historical_conversation/', views.DebugHistoricalConversation.Operate.as_view()), path('workspace//application//historical_conversation_record///', views.DebugHistoricalConversation.RecordPageView.as_view()), diff --git a/apps/application/views/application_chat.py b/apps/application/views/application_chat.py index c0b0c5bb5c8..8b31ea80f26 100644 --- a/apps/application/views/application_chat.py +++ b/apps/application/views/application_chat.py @@ -18,7 +18,7 @@ ApplicationChatExportAPI from application.models import ChatUserType, Application, ChatSourceChoices from application.serializers.application_chat import ApplicationChatQuerySerializers -from chat.api.chat_api import ChatAPI, PromptGenerateAPI +from chat.api.chat_api import ChatAPI, PromptGenerateAPI, PageHistoricalConversationAPI, HistoricalConversationRecordAPI from chat.api.chat_authentication_api import ChatOpenAPI from chat.serializers.chat import OpenChatSerializers, ChatSerializers, DebugChatSerializers, PromptGenerateSerializer from common.auth import TokenAuth @@ -137,11 +137,10 @@ def get(self, request: Request, workspace_id: str, application_id: str): ip_address = _get_ip_address(request) return result.success(OpenChatSerializers( data={'workspace_id': workspace_id, 'application_id': application_id, - 'chat_user_id': str(uuid.uuid7()), 'chat_user_type': ChatUserType.ANONYMOUS_USER, + 'chat_user_id': str(request.user.id), 'chat_user_type': ChatUserType.SYSTEM_USER, 'ip_address': ip_address, 'source': { - 'type': ChatSourceChoices.ONLINE.value} - , + 'type': ChatSourceChoices.ONLINE.value}, 'debug': True}).open()) @@ -162,6 +161,27 @@ def post(self, request: Request, chat_id: str): return DebugChatSerializers(data={'chat_id': chat_id}).chat(request.data) +class CancelWorkflowView(APIView): + authentication_classes = [TokenAuth] + + @extend_schema( + methods=['POST'], + description=_("Cancel running workflow"), + summary=_("Cancel running workflow"), + operation_id=_("Cancel running workflow"), # type: ignore + tags=[_('Application')] # type: ignore + ) + def post(self, request: Request, chat_id: str): + from application.workflow.workflow_run_registry import WorkflowRunRegistry, CancelResult + result_enum = WorkflowRunRegistry.cancel_by_chat_id(chat_id) + if result_enum == CancelResult.CANCELLED: + return result.success({'status': 'cancelled', 'chat_id': chat_id}) + elif result_enum == CancelResult.NOT_FOUND: + return result.success({'status': 'not_found', 'chat_id': chat_id}) + else: + return result.fail(500, _('Failed to cancel workflow')) + + class PromptGenerateView(APIView): authentication_classes = [TokenAuth] @@ -194,23 +214,47 @@ class DebugHistoricalConversation(APIView): class PageView(APIView): authentication_classes = [TokenAuth] + @extend_schema( + methods=['GET'], + description=_("Get historical conversation by page"), + summary=_("Get historical conversation by page"), + operation_id=_("Get historical conversation by page"), # type: ignore + parameters=PageHistoricalConversationAPI.get_parameters(), + responses=PageHistoricalConversationAPI.get_response(), + tags=[_('Chat')] # type: ignore + ) + @has_permissions(PermissionConstants.APPLICATION_READ.get_workspace_application_permission(), + PermissionConstants.APPLICATION_READ.get_workspace_permission_workspace_manage_role(), + ViewPermission([RoleConstants.USER.get_workspace_role()], + [PermissionConstants.APPLICATION.get_workspace_application_permission()], + CompareConstants.AND), + RoleConstants.WORKSPACE_MANAGE.get_workspace_role()) def get(self, request: Request, workspace_id: str, application_id: str, current_page: int, page_size: int): - from chat.serializers.chat_record import HistoricalConversationSerializer, page_search, HistoryChatModel - from django.db.models import QuerySet - from application.models import Chat - - queryset = QuerySet(Chat).filter( - application_id=application_id, - is_deleted=False - ).order_by('-update_time', 'id') - - return result.success( - page_search(current_page, page_size, queryset, lambda r: HistoryChatModel(r).data) - ) + from chat.serializers.chat_record import HistoricalConversationSerializer + return result.success(HistoricalConversationSerializer( + data={ + 'application_id': application_id, + 'chat_user_id': str(request.user.id), + }).page(current_page, page_size)) class RecordPageView(APIView): authentication_classes = [TokenAuth] + @extend_schema( + methods=['GET'], + description=_("Get historical conversation records"), + summary=_("Get historical conversation records"), + operation_id=_("Get historical conversation records"), # type: ignore + parameters=HistoricalConversationRecordAPI.get_parameters(), + responses=HistoricalConversationRecordAPI.get_response(), + tags=[_('Chat')] # type: ignore + ) + @has_permissions(PermissionConstants.APPLICATION_READ.get_workspace_application_permission(), + PermissionConstants.APPLICATION_READ.get_workspace_permission_workspace_manage_role(), + ViewPermission([RoleConstants.USER.get_workspace_role()], + [PermissionConstants.APPLICATION.get_workspace_application_permission()], + CompareConstants.AND), + RoleConstants.WORKSPACE_MANAGE.get_workspace_role()) def get(self, request: Request, workspace_id: str, application_id: str, chat_id: str, current_page: int, page_size: int): from chat.serializers.chat_record import HistoricalConversationRecordSerializer @@ -219,24 +263,10 @@ def get(self, request: Request, workspace_id: str, application_id: str, chat_id: data={ 'application_id': application_id, 'chat_id': chat_id, - 'chat_user_id': str(uuid.uuid7()), + 'chat_user_id': str(request.user.id), } ) - try: - return result.success(serializer.page(current_page, page_size)) - except Exception: - from django.db.models import QuerySet - from application.models import ChatRecord - from common.db.search import page_search - from application.serializers.application_chat_record import ChatRecordSerializerModel - - queryset = QuerySet(ChatRecord).filter( - chat_id=chat_id, - ).order_by('create_time', 'id') - - return result.success( - page_search(current_page, page_size, queryset, lambda r: ChatRecordSerializerModel(r).data) - ) + return result.success(serializer.page(current_page, page_size)) class Operate(APIView): authentication_classes = [TokenAuth] diff --git a/apps/application/workflow/i_node.py b/apps/application/workflow/i_node.py index b9906956548..c67ef889614 100644 --- a/apps/application/workflow/i_node.py +++ b/apps/application/workflow/i_node.py @@ -18,10 +18,16 @@ from common.utils.logger import maxkb_logger +class CancelledException(Exception): + """工作流取消异常""" + pass + + class Signal(str, Enum): BREAK = 'BREAK' CONTINUE = 'CONTINUE' FORM = 'FORM' + CANCELLED = "CANCELLED" class INode: @@ -88,6 +94,8 @@ def run(self): self.status = Status.RUNNING try: self._run() + except CancelledException: + self.complete(Status.CANCELLED) except Exception as e: self.complete(Status.FAIL, error=e) @@ -120,7 +128,8 @@ def complete(self, status, anchors=None, error=None, signal: Optional[Signal] = self.workflow_manage.signal = signal anchors = [] if anchors is None: - anchors = [self.success_anchor() if status == Status.SUCCESS else self.fail_anchor()] + anchors = [self.success_anchor() if [Status.SUCCESS, Status.CANCELLED].__contains__( + status) else self.fail_anchor()] self._dispatch(anchors) self.workflow_manage.assertion_end(error) @@ -201,3 +210,18 @@ def get_next_nodes(self, wf): def write(self, message: Content): self.workflow_manage.write(message) + + def cancel(self): + """ + 取消运行 + @return: + """ + self.status = Status.CANCELLED + + def _check_cancelled(self): + """ + 检查是否已取消，如果已取消则抛出 CancelledException + @return: + """ + if self.status == Status.CANCELLED or self.workflow_manage.signal == Signal.CANCELLED: + raise CancelledException() diff --git a/apps/application/workflow/nodes/ai_chat_node/ai_chat_node.py b/apps/application/workflow/nodes/ai_chat_node/ai_chat_node.py index bbb92d24512..c24ad22d283 100644 --- a/apps/application/workflow/nodes/ai_chat_node/ai_chat_node.py +++ b/apps/application/workflow/nodes/ai_chat_node/ai_chat_node.py @@ -285,6 +285,7 @@ def _stream_response(self, response, chat_model, message_list, question, response_reasoning_content = False for chunk in response: + self._check_cancelled() reasoning_chunk = reasoning.get_reasoning_content(chunk) content_chunk = reasoning_chunk.get('content') if 'reasoning_content' in chunk.additional_kwargs: @@ -472,6 +473,7 @@ def _handle_mcp( answer = '' tool_calls_map = {} for chunk in r: + self._check_cancelled() if isinstance(chunk, ToolMessage): tool_call = tool_calls_map.get(chunk.tool_call_id, {}) self.write(ToolContent( diff --git a/apps/application/workflow/nodes/image_generate_node/image_generate_node.py b/apps/application/workflow/nodes/image_generate_node/image_generate_node.py index 5626194c895..151060b038a 100644 --- a/apps/application/workflow/nodes/image_generate_node/image_generate_node.py +++ b/apps/application/workflow/nodes/image_generate_node/image_generate_node.py @@ -93,6 +93,7 @@ def execute(self): self.write_context('negative_prompt', self.workflow_manage.generate_prompt(negative_prompt or '')) self.write_context('dialogue_type', dialogue_type) + self._check_cancelled() image_urls = tti_model.generate_image(question, negative_prompt) file_urls = [] diff --git a/apps/application/workflow/nodes/image_to_video_node/image_to_video_node.py b/apps/application/workflow/nodes/image_to_video_node/image_to_video_node.py index 88ca29cb67e..941be8779ac 100644 --- a/apps/application/workflow/nodes/image_to_video_node/image_to_video_node.py +++ b/apps/application/workflow/nodes/image_to_video_node/image_to_video_node.py @@ -120,6 +120,7 @@ def execute(self): first_frame_url = self._get_file_base64(first_frame_url) last_frame_url = self._get_file_base64(last_frame_url) + self._check_cancelled() video_urls = ttv_model.generate_video(question, negative_prompt, first_frame_url, last_frame_url) maxkb_logger.info(f'[ImageToVideoNode] generate_video result: {video_urls is not None}, node_id={self.get_node_id()}') diff --git a/apps/application/workflow/nodes/image_understand_node/image_understand_node.py b/apps/application/workflow/nodes/image_understand_node/image_understand_node.py index 8531a5322d8..72760d7d79c 100644 --- a/apps/application/workflow/nodes/image_understand_node/image_understand_node.py +++ b/apps/application/workflow/nodes/image_understand_node/image_understand_node.py @@ -136,6 +136,7 @@ def _stream_response(self, response, chat_model, message_list, question, response_reasoning_content = False for chunk in response: + self._check_cancelled() reasoning_chunk = reasoning.get_reasoning_content(chunk) content_chunk = reasoning_chunk.get('content') if 'reasoning_content' in chunk.additional_kwargs: diff --git a/apps/application/workflow/nodes/intent_node/intent_node.py b/apps/application/workflow/nodes/intent_node/intent_node.py index f80cabe19f1..b1db2d870ef 100644 --- a/apps/application/workflow/nodes/intent_node/intent_node.py +++ b/apps/application/workflow/nodes/intent_node/intent_node.py @@ -100,6 +100,7 @@ def execute(self): self.write_context('message_list', message_list) try: + self._check_cancelled() r = chat_model.invoke(message_list) classification_result = r.content.strip() matched_branch = self._parse_classification_result(classification_result, branch) diff --git a/apps/application/workflow/nodes/mcp_node/mcp_node.py b/apps/application/workflow/nodes/mcp_node/mcp_node.py index c50986768b5..42a3dbc31ee 100644 --- a/apps/application/workflow/nodes/mcp_node/mcp_node.py +++ b/apps/application/workflow/nodes/mcp_node/mcp_node.py @@ -57,6 +57,8 @@ def execute(self): params = json.loads(json.dumps(tool_params)) params = self._handle_variables(params) + self._check_cancelled() + async def call_tool(t, a): client = MultiServerMCPClient(servers) async with client.session(mcp_server) as s: diff --git a/apps/application/workflow/nodes/parameter_extraction_node/parameter_extraction_node.py b/apps/application/workflow/nodes/parameter_extraction_node/parameter_extraction_node.py index e6ded6cb9e4..b9702378d6b 100644 --- a/apps/application/workflow/nodes/parameter_extraction_node/parameter_extraction_node.py +++ b/apps/application/workflow/nodes/parameter_extraction_node/parameter_extraction_node.py @@ -131,6 +131,7 @@ def execute(self): self.write_context('request', input_variable_str) content = _generate_content(input_variable_str, variable_list) + self._check_cancelled() response = chat_model.invoke([HumanMessage(content=content)]) result = _json_loads(response.content, variable_list) diff --git a/apps/application/workflow/nodes/question_node/question_node.py b/apps/application/workflow/nodes/question_node/question_node.py index 07c4f7881d0..fad17aa8103 100644 --- a/apps/application/workflow/nodes/question_node/question_node.py +++ b/apps/application/workflow/nodes/question_node/question_node.py @@ -116,6 +116,7 @@ def execute(self): answer = '' for chunk in response: + self._check_cancelled() answer += chunk.content message_tokens = chat_model.get_num_tokens_from_messages(message_list) diff --git a/apps/application/workflow/nodes/reranker_node/reranker_node.py b/apps/application/workflow/nodes/reranker_node/reranker_node.py index 549f28569f9..e6dc06a8faf 100644 --- a/apps/application/workflow/nodes/reranker_node/reranker_node.py +++ b/apps/application/workflow/nodes/reranker_node/reranker_node.py @@ -143,6 +143,7 @@ def execute(self): workspace_id = workflow_params.get('workspace_id') reranker_model = get_model_instance_by_model_workspace_id(reranker_model_id, workspace_id, top_n=top_n) + self._check_cancelled() result = reranker_model.compress_documents(documents, question) similarity = reranker_setting.get('similarity', 0.6) diff --git a/apps/application/workflow/nodes/search_knowledge_node/search_knowledge_node.py b/apps/application/workflow/nodes/search_knowledge_node/search_knowledge_node.py index f5ced062bc9..051fec90a68 100644 --- a/apps/application/workflow/nodes/search_knowledge_node/search_knowledge_node.py +++ b/apps/application/workflow/nodes/search_knowledge_node/search_knowledge_node.py @@ -180,6 +180,7 @@ def execute(self): return model_id = _get_embedding_id(knowledge_id_list) + self._check_cancelled() embedding_model = get_model_instance_by_model_workspace_id(model_id, workspace_id) embedding_value = embedding_model.embed_query(question) vector = VectorStore.get_embedding_vector() @@ -187,6 +188,7 @@ def execute(self): exclude_document_id_list = [str(document.id) for document in QuerySet(Document).filter(knowledge_id__in=knowledge_id_list, is_active=False)] + self._check_cancelled() embedding_list = vector.query(question, embedding_value, knowledge_id_list, document_id_list, exclude_document_id_list, exclude_paragraph_id_list, True, knowledge_setting.get('top_n'), knowledge_setting.get('similarity'), diff --git a/apps/application/workflow/nodes/speech_to_text_node/speech_to_text_node.py b/apps/application/workflow/nodes/speech_to_text_node/speech_to_text_node.py index a28df6ffcf6..630e85fc7fd 100644 --- a/apps/application/workflow/nodes/speech_to_text_node/speech_to_text_node.py +++ b/apps/application/workflow/nodes/speech_to_text_node/speech_to_text_node.py @@ -94,6 +94,7 @@ def execute(self): self.write_context('audio_list', audio_list) + self._check_cancelled() result = _process_audio_items(audio_list, stt_model) content = [] result_content = [] diff --git a/apps/application/workflow/nodes/text_to_speech_node/text_to_speech_node.py b/apps/application/workflow/nodes/text_to_speech_node/text_to_speech_node.py index ce2fa412918..b5d8fd9bdae 100644 --- a/apps/application/workflow/nodes/text_to_speech_node/text_to_speech_node.py +++ b/apps/application/workflow/nodes/text_to_speech_node/text_to_speech_node.py @@ -88,6 +88,7 @@ def execute(self): temp_files = [] for chunk in content_chunks: + self._check_cancelled() self.write_context('content', chunk) workspace_id = workflow_params.get('workspace_id') model = get_model_instance_by_model_workspace_id( diff --git a/apps/application/workflow/nodes/text_to_video_node/text_to_video_node.py b/apps/application/workflow/nodes/text_to_video_node/text_to_video_node.py index a8bb65791fc..c4f13066f8b 100644 --- a/apps/application/workflow/nodes/text_to_video_node/text_to_video_node.py +++ b/apps/application/workflow/nodes/text_to_video_node/text_to_video_node.py @@ -98,6 +98,7 @@ def execute(self): self.write_context('dialogue_type', dialogue_type) self.write_context('negative_prompt', self.workflow_manage.generate_prompt(negative_prompt)) + self._check_cancelled() video_urls = ttv_model.generate_video(question, negative_prompt) maxkb_logger.info(f'[TextToVideoNode] generate_video result: {video_urls is not None}, node_id={self.get_node_id()}') diff --git a/apps/application/workflow/workflow_manage.py b/apps/application/workflow/workflow_manage.py index fd989bca779..17beec7edc3 100644 --- a/apps/application/workflow/workflow_manage.py +++ b/apps/application/workflow/workflow_manage.py @@ -14,7 +14,7 @@ from langchain_core.prompts import PromptTemplate from application.workflow.common import Workflow, WorkflowType, Node, get_node_parameters -from application.workflow.i_node import INode +from application.workflow.i_node import INode, Signal from application.workflow.message.struct.content import Content from application.workflow.status import Status @@ -77,6 +77,8 @@ def next_nodes(self, nodes: Optional[List[Node]]): @param nodes: 执行下面要执行的节点 @return: """ + if [Signal.FORM, Signal.CANCELLED].__contains__(self.signal): + return if nodes is None or len(nodes) == 0: return # 需要校验是否可执行 @@ -247,3 +249,8 @@ def from_context(cls, chat_record_id, workflow, parameters, workflow_type, import traceback traceback.print_exc() return None + + def cancel(self): + self.signal = Signal.CANCELLED + for node in self.nodes: + node.cancel() diff --git a/apps/application/workflow/workflow_run_registry.py b/apps/application/workflow/workflow_run_registry.py new file mode 100644 index 00000000000..5e240eb84c4 --- /dev/null +++ b/apps/application/workflow/workflow_run_registry.py @@ -0,0 +1,168 @@ +# coding=utf-8 +""" + @project: MaxKB + @file： workflow_run_registry.py + @desc: 工作流运行注册表，用于管理和取消正在运行的工作流实例 +""" +import threading +from enum import Enum + +from common.utils.logger import maxkb_logger + + +class CancelResult(Enum): + """取消操作结果""" + CANCELLED = "CANCELLED" + NOT_FOUND = "NOT_FOUND" + FAILED = "FAILED" + + +class WorkflowRunRegistry: + _lock = threading.Lock() + _running = {} # {chat_record_id: WorkflowManage} + _chat_to_records = {} # {chat_id: set[chat_record_id]} + + @classmethod + def register(cls, chat_record_id: str, chat_id: str, workflow_manage) -> None: + """ + 注册一个正在运行的工作流实例 + @param chat_record_id: 聊天记录ID + @param chat_id: 聊天ID + @param workflow_manage: WorkflowManage 实例 + """ + if not chat_record_id or not workflow_manage: + return + with cls._lock: + cls._running[str(chat_record_id)] = workflow_manage + if chat_id: + if chat_id not in cls._chat_to_records: + cls._chat_to_records[chat_id] = set() + cls._chat_to_records[chat_id].add(str(chat_record_id)) + maxkb_logger.debug(f"Workflow registered: {chat_record_id}, total running: {len(cls._running)}") + + @classmethod + def unregister(cls, chat_record_id: str, chat_id: str = None) -> None: + """ + 注销一个工作流实例（无论成功/失败/取消都应调用） + @param chat_record_id: 聊天记录ID + @param chat_id: 聊天ID + """ + if not chat_record_id: + return + with cls._lock: + removed = cls._running.pop(str(chat_record_id), None) + if chat_id and chat_id in cls._chat_to_records: + cls._chat_to_records[chat_id].discard(str(chat_record_id)) + if not cls._chat_to_records[chat_id]: + del cls._chat_to_records[chat_id] + if removed is not None: + maxkb_logger.debug(f"Workflow unregistered: {chat_record_id}, total running: {len(cls._running)}") + + @classmethod + def cancel_by_chat_id(cls, chat_id: str) -> CancelResult: + """ + 取消某个聊天下所有运行中的工作流 + @param chat_id: 聊天ID + @return: CancelResult + """ + if not chat_id: + return CancelResult.NOT_FOUND + + with cls._lock: + record_ids = list(cls._chat_to_records.get(chat_id, set())) + + if not record_ids: + maxkb_logger.info(f"Cancel requested but no running workflow found for chat: {chat_id}") + return CancelResult.NOT_FOUND + + cancelled_count = 0 + failed_count = 0 + for record_id in record_ids: + with cls._lock: + wm = cls._running.get(record_id) + if wm: + try: + wm.cancel() + cancelled_count += 1 + maxkb_logger.info(f"Cancel signal sent to workflow: {record_id}") + except Exception as e: + failed_count += 1 + maxkb_logger.error(f"Failed to cancel workflow: {record_id}, error: {e}") + + if failed_count > 0 and cancelled_count == 0: + return CancelResult.FAILED + return CancelResult.CANCELLED + + @classmethod + def cancel_by_record_id(cls, chat_record_id: str) -> CancelResult: + """ + 取消某个特定的工作流 + @param chat_record_id: 聊天记录ID + @return: CancelResult + """ + if not chat_record_id: + return CancelResult.NOT_FOUND + + with cls._lock: + wm = cls._running.get(str(chat_record_id)) + + if wm is None: + maxkb_logger.info(f"Cancel requested but workflow not found (may already finished): {chat_record_id}") + return CancelResult.NOT_FOUND + + try: + wm.cancel() + maxkb_logger.info(f"Cancel signal sent to workflow: {chat_record_id}") + return CancelResult.CANCELLED + except Exception as e: + maxkb_logger.error(f"Failed to cancel workflow: {chat_record_id}, error: {e}") + return CancelResult.FAILED + + @classmethod + def get(cls, chat_record_id: str): + """ + 获取正在运行的工作流实例 + @param chat_record_id: 聊天记录ID + @return: WorkflowManage 实例或 None + """ + if not chat_record_id: + return None + return cls._running.get(str(chat_record_id)) + + @classmethod + def is_running(cls, chat_record_id: str) -> bool: + """ + 检查工作流是否正在运行 + @param chat_record_id: 聊天记录ID + @return: 是否正在运行 + """ + return chat_record_id is not None and str(chat_record_id) in cls._running + + @classmethod + def is_chat_running(cls, chat_id: str) -> bool: + """ + 检查某个聊天是否有正在运行的工作流 + @param chat_id: 聊天ID + @return: 是否有正在运行的工作流 + """ + if not chat_id: + return False + with cls._lock: + return chat_id in cls._chat_to_records and len(cls._chat_to_records[chat_id]) > 0 + + @classmethod + def running_count(cls) -> int: + """ + 获取正在运行的工作流数量 + @return: 数量 + """ + return len(cls._running) + + @classmethod + def running_ids(cls) -> list: + """ + 获取所有正在运行的工作流ID列表 + @return: ID列表 + """ + with cls._lock: + return list(cls._running.keys()) diff --git a/apps/chat/serializers/chat.py b/apps/chat/serializers/chat.py index 4aa6401ac94..18b2f1928e7 100644 --- a/apps/chat/serializers/chat.py +++ b/apps/chat/serializers/chat.py @@ -29,6 +29,7 @@ from application.workflow.common import WorkflowType, new_instance from application.workflow.message.aggregator import AggregationManager from application.workflow.workflow_manage import WorkflowManage, CallBack +from application.workflow.workflow_run_registry import WorkflowRunRegistry from application.workflow.nodes import get_start_node from application.workflow.message.struct.text_content import TextContent from application.workflow.message.struct.reasoning_content import ReasoningContent @@ -510,6 +511,8 @@ def position_to_dict(pos): })) def on_complete(wf_manage, error): + # 注销工作流实例 + WorkflowRunRegistry.unregister(chat_record_id_str, str(chat_info.chat_id)) if error: result_queue.put(('error', error)) self._save_chat_record(chat_info, chat_info.chat_id, chat_record_id_str, @@ -543,6 +546,9 @@ def get_start_node_fn(wf, wm): work_flow_manage.start_node.workflow_manage = work_flow_manage + # 注册工作流实例到注册表 + WorkflowRunRegistry.register(chat_record_id_str, str(chat_info.chat_id), work_flow_manage) + chat_info.set_chat(message) if stream: diff --git a/apps/chat/urls.py b/apps/chat/urls.py index 5a386bda107..638ad829e60 100644 --- a/apps/chat/urls.py +++ b/apps/chat/urls.py @@ -14,6 +14,7 @@ path('profile', views.AuthProfile.as_view()), path('application/profile', views.ApplicationProfile.as_view(), name='profile'), path('chat_message/', views.ChatView.as_view(), name='chat'), + path('chat_message//cancel', views.CancelWorkflowView.as_view(), name='cancel_workflow'), path('open', views.OpenView.as_view(), name='open'), path('text_to_speech', views.TextToSpeech.as_view()), path('speech_to_text', views.SpeechToText.as_view()), diff --git a/apps/chat/views/chat.py b/apps/chat/views/chat.py index 8f1edfc10a6..07296313d51 100644 --- a/apps/chat/views/chat.py +++ b/apps/chat/views/chat.py @@ -13,6 +13,8 @@ from django.http import HttpResponse, StreamingHttpResponse from django.utils.translation import gettext_lazy as _ from drf_spectacular.utils import extend_schema +from drf_spectacular.types import OpenApiTypes +from drf_spectacular.utils import OpenApiParameter from rest_framework.parsers import MultiPartParser from rest_framework.request import Request from rest_framework.views import APIView @@ -226,6 +228,32 @@ def get(self, request: Request): 'debug': False}).open()) +class CancelWorkflowView(APIView): + authentication_classes = [ChatTokenAuth] + + @extend_schema( + methods=['POST'], + description=_("Cancel running workflow"), + summary=_("Cancel running workflow"), + operation_id=_("Cancel running workflow"), # type: ignore + parameters=[ + OpenApiParameter(name='chat_id', type=OpenApiTypes.UUID, location=OpenApiParameter.PATH, + description=_('Chat ID')), + ], + responses=None, + tags=[_('Chat')] # type: ignore + ) + def post(self, request: Request, chat_id: str): + from application.workflow.workflow_run_registry import WorkflowRunRegistry, CancelResult + result_enum = WorkflowRunRegistry.cancel_by_chat_id(chat_id) + if result_enum == CancelResult.CANCELLED: + return result.success({'status': 'cancelled', 'chat_id': chat_id}) + elif result_enum == CancelResult.NOT_FOUND: + return result.success({'status': 'not_found', 'chat_id': chat_id}) + else: + return result.fail(500, _('Failed to cancel workflow')) + + class CaptchaView(APIView): @extend_schema(methods=['GET'], summary=_("Get Chat captcha"),