feat: Support workflow stoppage

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):

apps/application/urls.py

@@ -44,6 +44,7 @@
     path('workspace/<str:workspace_id>/application/<str:application_id>/mcp_tools', views.McpServers.as_view()),
     path('workspace/<str:workspace_id>/application/<str:application_id>/model/<str:model_id>/prompt_generate', views.PromptGenerateView.as_view()),
     path('chat_message/<str:chat_id>', views.ChatView.as_view()),
+    path('chat_message/<str:chat_id>/cancel', views.CancelWorkflowView.as_view()),
     path('workspace/<str:workspace_id>/application/<str:application_id>/historical_conversation/<int:current_page>/<int:page_size>', views.DebugHistoricalConversation.PageView.as_view()),
     path('workspace/<str:workspace_id>/application/<str:application_id>/historical_conversation/<str:chat_id>', views.DebugHistoricalConversation.Operate.as_view()),
     path('workspace/<str:workspace_id>/application/<str:application_id>/historical_conversation_record/<str:chat_id>/<int:current_page>/<int:page_size>', views.DebugHistoricalConversation.RecordPageView.as_view()),

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]

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()

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(

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 = []

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()}')
 

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:

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)

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:

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)
 

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)

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)

apps/application/workflow/nodes/search_knowledge_node/search_knowledge_node.py

@@ -180,13 +180,15 @@ 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()
 
         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'),

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 = []

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(

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()}')
 

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()