Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 9 additions & 3 deletions apps/application/workflow/workflow_run_registry.py
Original file line number Diff line number Diff line change
@@ -1,17 +1,20 @@
# coding=utf-8
"""
@project: MaxKB
@file: workflow_run_registry.py
@desc: 工作流运行注册表,用于管理和取消正在运行的工作流实例
@project: MaxKB
@file: workflow_run_registry.py
@desc: 工作流运行注册表,用于管理和取消正在运行的工作流实例
"""

import threading
from enum import Enum

from common.utils.logger import maxkb_logger
from common.bus import on


class CancelResult(Enum):
"""取消操作结果"""

CANCELLED = "CANCELLED"
NOT_FOUND = "NOT_FOUND"
FAILED = "FAILED"
Expand Down Expand Up @@ -166,3 +169,6 @@ def running_ids(cls) -> list:
"""
with cls._lock:
return list(cls._running.keys())


on("cancel_chat", lambda args: WorkflowRunRegistry.cancel_by_chat_id(args.get("chat_id")))
13 changes: 3 additions & 10 deletions apps/chat/views/v3/chat.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@
from rest_framework.parsers import MultiPartParser
from rest_framework.request import Request
from rest_framework.views import APIView

from common.bus import emit
from application.api.application_api import SpeechToTextAPI, TextToSpeechAPI
from application.models import ChatUserType, ChatSourceChoices, ApplicationAccessToken, ChatRecord
from chat.api.chat_api import ChatAPI
Expand Down Expand Up @@ -285,15 +285,8 @@ class CancelWorkflowView(APIView):
tags=[_("V3 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"))
emit("cancel_chat", {"chat_id": str(chat_id)})
return result.success({"status": "cancelled", "chat_id": chat_id})


class ResumeStreamView(APIView):
Expand Down
1 change: 1 addition & 0 deletions apps/common/bus/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
from common.bus.message_bus import on, emit
109 changes: 109 additions & 0 deletions apps/common/bus/message_bus.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,109 @@
# coding=utf-8
"""
@project: MaxKB
@file: message_bus.py
@desc: 基于 Redis Pub/Sub 的跨进程、跨系统消息总线。提供 emit / on 两个接口。
"""

import json
import threading
from typing import Callable, Dict, List

from common.utils.logger import maxkb_logger

CHANNEL = "maxkb:bus"

_handlers: Dict[str, List[Callable]] = {}
_handlers_lock = threading.Lock()

_subscriber_started = False
_subscriber_thread = None
_subscriber_lock = threading.Lock()


def _redis_client():
"""取底层 redis-py 客户端(含 pubsub 能力),已由 Django settings 配置支持 sentinel/普通 Redis。"""
from django_redis import get_redis_connection

return get_redis_connection("default")


def _start_subscriber() -> None:
"""懒启动:每个进程只起一条 daemon 订阅线程;线程因建连失败退出后,下次调用会重新拉起。"""
global _subscriber_started, _subscriber_thread
if _subscriber_started and _subscriber_thread is not None and _subscriber_thread.is_alive():
return
with _subscriber_lock:
if _subscriber_started and _subscriber_thread is not None and _subscriber_thread.is_alive():
return
_subscriber_started = True
_subscriber_thread = threading.Thread(target=_listen, name="maxkb-bus", daemon=True)
_subscriber_thread.start()


def _listen() -> None:
# 建连失败(Redis 不可用)则记日志退出,不自动重连;由 emit 检测线程存活后重新拉起。
# 建连成功后的消息循环由 redis-py 内部维持,无需自行重连。
ps = None
try:
ps = _redis_client().pubsub(ignore_subscribe_messages=True)
ps.subscribe(CHANNEL)
for message in ps.listen():
_dispatch(message)
except Exception as e:
maxkb_logger.error(f"Message bus subscriber failed: {e}")
finally:
if ps is not None:
try:
ps.close()
except Exception:
pass


def _dispatch(message) -> None:
try:
data = message.get("data")
if isinstance(data, (bytes, bytearray)):
data = data.decode("utf-8")
body = json.loads(data)
event = body.get("event")
if not event:
return
args = body.get("args")
for handler in list(_handlers.get(event, ())):
try:
handler(args)
except Exception as e:
maxkb_logger.error(f"Message bus handler error for event[{event}]: {e}")
except Exception as e:
maxkb_logger.error("Failed to dispatch message bus event: %s", e)


def on(event: str, handler: Callable) -> None:
"""
注册一个事件监听,本进程内生效;收到该事件时异步调用 handler(args)。
:param event: 事件名称
:param handler: 回调函数 def handler(args):
pass
"""
with _handlers_lock:
_handlers.setdefault(event, []).append(handler)
_start_subscriber()


def emit(event: str, args=None) -> bool:
"""
向所有进程广播一个事件。Redis 不可用时返回 False。
:param event: 事件名称
:param args: 事件参数
"""
_start_subscriber()
try:
_redis_client().publish(
CHANNEL,
json.dumps({"event": event, "args": args}, ensure_ascii=False),
)
return True
except Exception as e:
maxkb_logger.warning(f"Message bus emit[{event}] failed: {e}")
return False
Loading