feat: 消息截断器

pkg/pipeline/msgtrun/msgtrun.py

@@ -0,0 +1,35 @@
+from __future__ import annotations
+
+from .. import stage, entities, stagemgr
+from ...core import entities as core_entities
+from . import truncator
+from .truncators import round
+
+
+@stage.stage_class("ConversationMessageTruncator")
+class ConversationMessageTruncator(stage.PipelineStage):
+    """会话消息截断器
+
+    用于截断会话消息链，以适应平台消息长度限制。
+    """
+    trun: truncator.Truncator
+
+    async def initialize(self):
+        use_method = self.ap.pipeline_cfg.data['msg-truncate']['method']
+
+        for trun in truncator.preregistered_truncators:
+            if trun.name == use_method:
+                self.trun = trun(self.ap)
+                break
+        else:
+            raise ValueError(f"未知的截断器: {use_method}")
+
+    async def process(self, query: core_entities.Query, stage_inst_name: str) -> entities.StageProcessResult:
+        """处理
+        """
+        query = await self.trun.truncate(query)
+
+        return entities.StageProcessResult(
+            result_type=entities.ResultType.CONTINUE,
+            new_query=query
+        )

pkg/pipeline/msgtrun/truncator.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+import typing
+import abc
+
+from ...core import entities as core_entities, app
+
+
+preregistered_truncators: list[typing.Type[Truncator]] = []
+
+
+def truncator_class(
+    name: str
+) -> typing.Callable[[typing.Type[Truncator]], typing.Type[Truncator]]:
+    """截断器类装饰器
+
+    Args:
+        name (str): 截断器名称
+
+    Returns:
+        typing.Callable[[typing.Type[Truncator]], typing.Type[Truncator]]: 装饰器
+    """
+    def decorator(cls: typing.Type[Truncator]) -> typing.Type[Truncator]:
+        assert issubclass(cls, Truncator)
+
+        cls.name = name
+
+        preregistered_truncators.append(cls)
+
+        return cls
+
+    return decorator
+
+
+class Truncator(abc.ABC):
+    """消息截断器基类
+    """
+
+    name: str
+
+    ap: app.Application
+    
+    def __init__(self, ap: app.Application):
+        self.ap = ap
+
+    async def initialize(self):
+        pass
+
+    @abc.abstractmethod
+    async def truncate(self, query: core_entities.Query) -> core_entities.Query:
+        """截断
+
+        一般只需要操作query.messages，也可以扩展操作query.prompt, query.user_message。
+        请勿操作其他字段。
+        """
+        pass

pkg/pipeline/msgtrun/truncators/round.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from .. import truncator
+from ....core import entities as core_entities
+
+
+@truncator.truncator_class("round")
+class RoundTruncator(truncator.Truncator):
+    """前文回合数阶段器
+    """
+
+    async def truncate(self, query: core_entities.Query) -> core_entities.Query:
+        """截断
+        """
+        max_round = self.ap.pipeline_cfg.data['msg-truncate']['round']['max-round']
+
+        temp_messages = []
+
+        current_round = 0
+
+        # 从后往前遍历
+        for msg in query.messages[::-1]:
+            if current_round < max_round:
+                temp_messages.append(msg)
+                if msg.role == 'user':
+                    current_round += 1
+            else:
+                break
+        
+        query.messages = temp_messages[::-1]
+
+        return query