From 5536c5bf7697c0a61cae1f7f538439bb5451948c Mon Sep 17 00:00:00 2001
From: tcmofashi <mofashiforzbx@qq.com>
Date: Fri, 23 May 2025 18:10:37 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E6=96=B0=E5=A2=9Emaim=5Fmessage?=
 =?UTF-8?q?=E9=85=8D=E7=BD=AE=EF=BC=8C=E6=94=AF=E6=8C=81wss=EF=BC=8Ctcp?=
 =?UTF-8?q?=E6=A8=A1=E5=BC=8F=EF=BC=8C=E6=94=AF=E6=8C=81token?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/common/message/api.py         | 34 ++++++++++++++++++++++++++++---
 src/config/config.py              |  2 ++
 src/config/official_configs.py    | 29 ++++++++++++++++++++++++++
 template/bot_config_template.toml | 15 +++++++++++++-
 4 files changed, 76 insertions(+), 4 deletions(-)

diff --git a/src/common/message/api.py b/src/common/message/api.py
index f11a78d30..7f8ffe7f9 100644
--- a/src/common/message/api.py
+++ b/src/common/message/api.py
@@ -3,25 +3,53 @@ import os
 import importlib.metadata
 from maim_message import MessageServer
 from src.common.logger_manager import get_logger
+from src.config.config import global_config
 
 
 # 检查maim_message版本
 try:
     maim_message_version = importlib.metadata.version("maim_message")
-    version_compatible = [int(x) for x in maim_message_version.split(".")] >= [0, 3, 0]
+    version_compatible = [int(x) for x in maim_message_version.split(".")] >= [0, 3, 3]
 except (importlib.metadata.PackageNotFoundError, ValueError):
     version_compatible = False
 
-# 根据版本决定是否使用自定义logger
+# 读取配置项
+maim_message_config = global_config.maim_message
+
+# 设置基本参数
 kwargs = {
     "host": os.environ["HOST"],
     "port": int(os.environ["PORT"]),
     "app": global_server.get_app(),
 }
 
-# 只有在版本 >= 0.3.0 时才使用自定义logger
+# 只有在版本 >= 0.3.0 时才使用高级特性
 if version_compatible:
+    # 添加自定义logger
     maim_message_logger = get_logger("maim_message")
     kwargs["custom_logger"] = maim_message_logger
 
+    # 添加token认证
+    if maim_message_config.auth_token:
+        if len(maim_message_config.auth_token) > 0:
+            kwargs["enable_token"] = True
+
+    if maim_message_config.use_custom:
+        # 添加WSS模式支持
+        del kwargs["app"]
+        kwargs["host"] = maim_message_config.host
+        kwargs["port"] = maim_message_config.port
+        kwargs["mode"] = maim_message_config.mode
+        if maim_message_config.use_wss:
+            if maim_message_config.cert_file:
+                kwargs["ssl_certfile"] = maim_message_config.cert_file
+            if maim_message_config.key_file:
+                kwargs["ssl_keyfile"] = maim_message_config.key_file
+        kwargs["enable_custom_uvicorn_logger"] = False
+
+
 global_api = MessageServer(**kwargs)
+
+if version_compatible and maim_message_config.auth_token:
+    for token in maim_message_config.auth_token:
+        global_api.add_valid_token(token)
diff --git a/src/config/config.py b/src/config/config.py
index b3a09850c..29f93a9e2 100644
--- a/src/config/config.py
+++ b/src/config/config.py
@@ -31,6 +31,7 @@ from src.config.official_configs import (
     ModelConfig,
     FocusChatProcessorConfig,
     MessageReceiveConfig,
+    MaimMessageConfig,
 )
 
 install(extra_lines=3)
@@ -157,6 +158,7 @@ class Config(ConfigBase):
     telemetry: TelemetryConfig
     experimental: ExperimentalConfig
     model: ModelConfig
+    maim_message: MaimMessageConfig
 
 
 def load_config(config_path: str) -> Config:
diff --git a/src/config/official_configs.py b/src/config/official_configs.py
index f75afcf6a..e0cf931a2 100644
--- a/src/config/official_configs.py
+++ b/src/config/official_configs.py
@@ -355,6 +355,35 @@ class ExperimentalConfig(ConfigBase):
     """是否启用PFC"""
 
 
+@dataclass
+class MaimMessageConfig(ConfigBase):
+    """maim_message配置类"""
+
+    use_custom: bool = False
+    """是否使用自定义的maim_message配置"""
+
+    host: str = "127.0.0.1"
+    """主机地址"""
+
+    port: int = 8090
+    """"端口号"""
+
+    mode: str = "ws"
+    """连接模式，支持ws和tcp"""
+
+    use_wss: bool = False
+    """是否使用WSS安全连接"""
+
+    cert_file: str = ""
+    """SSL证书文件路径，仅在use_wss=True时有效"""
+
+    key_file: str = ""
+    """SSL密钥文件路径，仅在use_wss=True时有效"""
+
+    auth_token: list[str] = field(default_factory=lambda: [])
+    """认证令牌，用于API验证，为空则不启用验证"""
+
+
 @dataclass
 class ModelConfig(ConfigBase):
     """模型配置类"""
diff --git a/template/bot_config_template.toml b/template/bot_config_template.toml
index cb13de768..f3822729e 100644
--- a/template/bot_config_template.toml
+++ b/template/bot_config_template.toml
@@ -169,6 +169,18 @@ max_length = 256 # 回复允许的最大长度
 max_sentence_num = 4 # 回复允许的最大句子数
 enable_kaomoji_protection = false # 是否启用颜文字保护
 
+[maim_message]
+auth_token = [] # 认证令牌，用于API验证，为空则不启用验证
+# 以下项目若要使用需要打开use_custom，并单独配置maim_message的服务器
+use_custom = false # 是否启用自定义的maim_message服务器，注意这需要设置新的端口，不能与.env重复
+host="127.0.0.1"
+port=8090
+mode="ws" # 支持ws和tcp两种模式
+use_wss = false # 是否使用WSS安全连接，只支持ws模式
+cert_file = "" # SSL证书文件路径，仅在use_wss=true时有效
+key_file = "" # SSL密钥文件路径，仅在use_wss=true时有效
+
+
 [telemetry] #发送统计信息，主要是看全球有多少只麦麦
 enable = true
 
@@ -292,4 +304,5 @@ pri_out = 8
 name = "Qwen/Qwen2.5-32B-Instruct"
 provider = "SILICONFLOW"
 pri_in = 1.26
-pri_out = 1.26
\ No newline at end of file
+pri_out = 1.26
+