Merge pull request #712 from RockChinQ/feat/component-extensibility

Feat: 更多组件的可扩展性

pkg/pipeline/cntfilter/cntfilter.py

@@ -7,7 +7,7 @@ from ...core import app
 from .. import stage, entities, stagemgr
 from ...core import entities as core_entities
 from ...config import manager as cfg_mgr
-from . import filter, entities as filter_entities
+from . import filter as filter_model, entities as filter_entities
 from .filters import cntignore, banwords, baiduexamine
 
 
@@ -16,20 +16,29 @@ from .filters import cntignore, banwords, baiduexamine
 class ContentFilterStage(stage.PipelineStage):
     """内容过滤阶段"""
 
-    filter_chain: list[filter.ContentFilter]
+    filter_chain: list[filter_model.ContentFilter]
 
     def __init__(self, ap: app.Application):
         self.filter_chain = []
         super().__init__(ap)
 
     async def initialize(self):
-        self.filter_chain.append(cntignore.ContentIgnore(self.ap))
+
+        filters_required = [
+            "content-filter"
+        ]
 
         if self.ap.pipeline_cfg.data['check-sensitive-words']:
-            self.filter_chain.append(banwords.BanWordFilter(self.ap))
-        
+            filters_required.append("ban-word-filter")
+
         if self.ap.pipeline_cfg.data['baidu-cloud-examine']['enable']:
-            self.filter_chain.append(baiduexamine.BaiduCloudExamine(self.ap))
+            filters_required.append("baidu-cloud-examine")
+
+        for filter in filter_model.preregistered_filters:
+            if filter.name in filters_required:
+                self.filter_chain.append(
+                    filter(self.ap)
+                )
 
         for filter in self.filter_chain:
             await filter.initialize()

pkg/pipeline/cntfilter/filter.py

@@ -1,12 +1,42 @@
 # 内容过滤器的抽象类
 from __future__ import annotations
 import abc
+import typing
 
 from ...core import app
 from . import entities
 
 
+preregistered_filters: list[typing.Type[ContentFilter]] = []
+
+
+def filter_class(
+    name: str
+) -> typing.Callable[[typing.Type[ContentFilter]], typing.Type[ContentFilter]]:
+    """内容过滤器类装饰器
+
+    Args:
+        name (str): 过滤器名称
+
+    Returns:
+        typing.Callable[[typing.Type[ContentFilter]], typing.Type[ContentFilter]]: 装饰器
+    """
+    def decorator(cls: typing.Type[ContentFilter]) -> typing.Type[ContentFilter]:
+        assert issubclass(cls, ContentFilter)
+
+        cls.name = name
+
+        preregistered_filters.append(cls)
+
+        return cls
+
+    return decorator
+
+
 class ContentFilter(metaclass=abc.ABCMeta):
+    """内容过滤器抽象类"""
+
+    name: str
 
     ap: app.Application
 

pkg/pipeline/cntfilter/filters/baiduexamine.py

@@ -10,6 +10,7 @@ BAIDU_EXAMINE_URL = "https://aip.baidubce.com/rest/2.0/solution/v1/text_censor/v
 BAIDU_EXAMINE_TOKEN_URL = "https://aip.baidubce.com/oauth/2.0/token"
 
 
+@filter_model.filter_class("baidu-cloud-examine")
 class BaiduCloudExamine(filter_model.ContentFilter):
     """百度云内容审核"""
 

pkg/pipeline/cntfilter/filters/banwords.py

@@ -6,6 +6,7 @@ from .. import entities
 from ....config import manager as cfg_mgr
 
 
+@filter_model.filter_class("ban-word-filter")
 class BanWordFilter(filter_model.ContentFilter):
     """根据内容禁言"""
 

pkg/pipeline/cntfilter/filters/cntignore.py

@@ -5,6 +5,7 @@ from .. import entities
 from .. import filter as filter_model
 
 
+@filter_model.filter_class("content-ignore")
 class ContentIgnore(filter_model.ContentFilter):
     """根据内容忽略消息"""
 

pkg/pipeline/longtext/longtext.py

@@ -45,11 +45,14 @@ class LongTextProcessStage(stage.PipelineStage):
                 self.ap.logger.error("加载字体文件失败({})，更换为转发消息组件以发送长消息，您可以在config.py中调整相关设置。".format(use_font))
 
                 self.ap.platform_cfg.data['long-text-process']['strategy'] = "forward"
-        
-        if config['strategy'] == 'image':
-            self.strategy_impl = image.Text2ImageStrategy(self.ap)
-        elif config['strategy'] == 'forward':
-            self.strategy_impl = forward.ForwardComponentStrategy(self.ap)
+
+        for strategy_cls in strategy.preregistered_strategies:
+            if strategy_cls.name == config['strategy']:
+                self.strategy_impl = strategy_cls(self.ap)
+                break
+        else:
+            raise ValueError(f"未找到名为 {config['strategy']} 的长消息处理策略")
+
         await self.strategy_impl.initialize()
     
     async def process(self, query: core_entities.Query, stage_inst_name: str) -> entities.StageProcessResult:

pkg/pipeline/longtext/strategies/forward.py

@@ -36,6 +36,7 @@ class Forward(MessageComponent):
         return '[聊天记录]'
 
 
+@strategy_model.strategy_class("forward")
 class ForwardComponentStrategy(strategy_model.LongTextStrategy):
 
     async def process(self, message: str, query: core_entities.Query) -> list[MessageComponent]:

pkg/pipeline/longtext/strategies/image.py

@@ -15,6 +15,7 @@ from .. import strategy as strategy_model
 from ....core import entities as core_entities
 
 
+@strategy_model.strategy_class("image")
 class Text2ImageStrategy(strategy_model.LongTextStrategy):
 
     text_render_font: ImageFont.FreeTypeFont

pkg/pipeline/longtext/strategy.py

@@ -9,7 +9,30 @@ from ...core import app
 from ...core import entities as core_entities
 
 
+preregistered_strategies: list[typing.Type[LongTextStrategy]] = []
+
+
+def strategy_class(
+    name: str
+) -> typing.Callable[[typing.Type[LongTextStrategy]], typing.Type[LongTextStrategy]]:
+    def decorator(cls: typing.Type[LongTextStrategy]) -> typing.Type[LongTextStrategy]:
+        assert issubclass(cls, LongTextStrategy)
+
+        cls.name = name
+
+        preregistered_strategies.append(cls)
+
+        return cls
+
+    return decorator
+
+
 class LongTextStrategy(metaclass=abc.ABCMeta):
+    """长文本处理策略抽象类
+    """
+
+    name: str
+
     ap: app.Application
 
     def __init__(self, ap: app.Application):

pkg/pipeline/preproc/preproc.py

@@ -51,28 +51,6 @@ class PreProcessor(stage.PipelineStage):
         query.prompt.messages = event_ctx.event.default_prompt
         query.messages = event_ctx.event.prompt
 
-        # 根据模型max_tokens剪裁
-        max_tokens = min(query.use_model.max_tokens, self.ap.pipeline_cfg.data['submit-messages-tokens'])
-
-        test_messages = query.prompt.messages + query.messages + [query.user_message]
-
-        while await query.use_model.tokenizer.count_token(test_messages, query.use_model) > max_tokens:
-            # 前文都pop完了，还是大于max_tokens，由于prompt和user_messages不能删减，报错
-            if len(query.prompt.messages) == 0:
-                return entities.StageProcessResult(
-                    result_type=entities.ResultType.INTERRUPT,
-                    new_query=query,
-                    user_notice='输入内容过长，请减少情景预设或者输入内容长度',
-                    console_notice='输入内容过长，请减少情景预设或者输入内容长度，或者增大配置文件中的 submit-messages-tokens 项（但不能超过所用模型最大tokens数）'
-                )
-
-            query.messages.pop(0) # pop第一个肯定是role=user的
-            # 继续pop到第二个role=user前一个
-            while len(query.messages) > 0 and query.messages[0].role != 'user':
-                query.messages.pop(0)
-
-            test_messages = query.prompt.messages + query.messages + [query.user_message]
-
         return entities.StageProcessResult(
             result_type=entities.ResultType.CONTINUE,
             new_query=query

pkg/pipeline/process/handlers/chat.py

@@ -21,8 +21,6 @@ class ChatMessageHandler(handler.MessageHandler):
     ) -> typing.AsyncGenerator[entities.StageProcessResult, None]:
         """处理
         """
-        # 取session
-        # 取conversation
         # 调API
         #   生成器
 

pkg/pipeline/ratelimit/algo.py

@@ -1,11 +1,26 @@
 from __future__ import annotations
 import abc
+import typing
 
 from ...core import app
 
 
+preregistered_algos: list[typing.Type[ReteLimitAlgo]] = []
+
+def algo_class(name: str):
+    
+    def decorator(cls: typing.Type[ReteLimitAlgo]) -> typing.Type[ReteLimitAlgo]:
+        cls.name = name
+        preregistered_algos.append(cls)
+        return cls
+    
+    return decorator
+
+
 class ReteLimitAlgo(metaclass=abc.ABCMeta):
     
+    name: str = None
+
     ap: app.Application
 
     def __init__(self, ap: app.Application):

pkg/pipeline/ratelimit/algos/fixedwin.py

@@ -19,6 +19,7 @@ class SessionContainer:
         self.records = {}
 
 
+@algo.algo_class("fixwin")
 class FixedWindowAlgo(algo.ReteLimitAlgo):
 
     containers_lock: asyncio.Lock

pkg/pipeline/ratelimit/ratelimit.py

@@ -16,7 +16,19 @@ class RateLimit(stage.PipelineStage):
     algo: algo.ReteLimitAlgo
 
     async def initialize(self):
-        self.algo = fixedwin.FixedWindowAlgo(self.ap)
+
+        algo_name = self.ap.pipeline_cfg.data['rate-limit']['algo']
+
+        algo_class = None
+
+        for algo_cls in algo.preregistered_algos:
+            if algo_cls.name == algo_name:
+                algo_class = algo_cls
+                break
+        else:
+            raise ValueError(f'未知的限速算法: {algo_name}')
+
+        self.algo = algo_class(self.ap)
         await self.algo.initialize()
 
     async def process(

pkg/pipeline/resprule/resprule.py

@@ -21,15 +21,13 @@ class GroupRespondRuleCheckStage(stage.PipelineStage):
     async def initialize(self):
         """初始化检查器
         """
-        self.rule_matchers = [
-            atbot.AtBotRule(self.ap),
-            prefix.PrefixRule(self.ap),
-            regexp.RegExpRule(self.ap),
-            random.RandomRespRule(self.ap),
-        ]
 
-        for rule_matcher in self.rule_matchers:
-            await rule_matcher.initialize()
+        self.rule_matchers = []
+
+        for rule_matcher in rule.preregisetered_rules:
+            rule_inst = rule_matcher(self.ap)
+            await rule_inst.initialize()
+            self.rule_matchers.append(rule_inst)
 
     async def process(self, query: core_entities.Query, stage_inst_name: str) -> entities.StageProcessResult:
         

pkg/pipeline/resprule/rule.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 import abc
+import typing
 
 import mirai
 
@@ -7,9 +8,20 @@ from ...core import app, entities as core_entities
 from . import entities
 
 
+preregisetered_rules: list[typing.Type[GroupRespondRule]] = []
+
+def rule_class(name: str):
+    def decorator(cls: typing.Type[GroupRespondRule]) -> typing.Type[GroupRespondRule]:
+        cls.name = name
+        preregisetered_rules.append(cls)
+        return cls
+    return decorator
+
+
 class GroupRespondRule(metaclass=abc.ABCMeta):
     """群组响应规则的抽象类
     """
+    name: str
 
     ap: app.Application
 

pkg/pipeline/resprule/rules/atbot.py

@@ -7,6 +7,7 @@ from .. import entities
 from ....core import entities as core_entities
 
 
+@rule_model.rule_class("at-bot")
 class AtBotRule(rule_model.GroupRespondRule):
 
     async def match(

pkg/pipeline/resprule/rules/prefix.py

@@ -5,6 +5,7 @@ from .. import entities
 from ....core import entities as core_entities
 
 
+@rule_model.rule_class("prefix")
 class PrefixRule(rule_model.GroupRespondRule):
 
     async def match(

pkg/pipeline/resprule/rules/random.py

@@ -7,6 +7,7 @@ from .. import entities
 from ....core import entities as core_entities
 
 
+@rule_model.rule_class("random")
 class RandomRespRule(rule_model.GroupRespondRule):
 
     async def match(

pkg/pipeline/resprule/rules/regexp.py

@@ -7,6 +7,7 @@ from .. import entities
 from ....core import entities as core_entities
 
 
+@rule_model.rule_class("regexp")
 class RegExpRule(rule_model.GroupRespondRule):
 
     async def match(