Kris
7a972055b7
新增大量渠道适配器相关的协议、策略、工具类与基础设施代码，包括： 1. 多协议定义：认证、消息、配置、网关等核心接口 2. 策略模块：上下文、群聊、去重、防抖等业务策略 3. 工具集：重试、去重、文本分块、消息格式化等SDK工具 4. 基础设施：外部进程管理、事件广播、熔断机制等 5. 账户与管道系统：账户管理、消息处理管道实现 6. 运行时服务：状态收集、维护任务、日志等后台服务
261 lines
7.8 KiB
Python
261 lines
7.8 KiB
Python
from __future__ import annotations
|
||
|
||
import re
|
||
from collections.abc import Iterator
|
||
from enum import StrEnum
|
||
from dataclasses import dataclass
|
||
|
||
|
||
class ChunkMode(StrEnum):
|
||
LENGTH = "length"
|
||
NEWLINE = "newline"
|
||
SENTENCE = "sentence"
|
||
MARKDOWN = "markdown"
|
||
|
||
|
||
@dataclass
|
||
class TextChunk:
|
||
text: str
|
||
index: int
|
||
is_last: bool
|
||
|
||
|
||
_NATURAL_BOUNDARIES = [
|
||
"\n\n",
|
||
"\n",
|
||
"。",
|
||
"！",
|
||
"？",
|
||
"；",
|
||
"：",
|
||
". ",
|
||
"! ",
|
||
"? ",
|
||
]
|
||
|
||
_MD_FENCE = re.compile(r"^```")
|
||
_MD_TABLE_SEP = re.compile(r"^\|[-:| ]+\|")
|
||
_MD_ORDERED_LIST = re.compile(r"^\d+\.\s")
|
||
_MD_UNORDERED_LIST = re.compile(r"^[-*+]\s")
|
||
_MD_HEADING = re.compile(r"^#{1,6}\s")
|
||
|
||
|
||
class TextChunker:
|
||
def __init__(self, max_length: int = 2000, mode: ChunkMode = ChunkMode.NEWLINE):
|
||
self.max_length = max_length
|
||
self.mode: ChunkMode = mode
|
||
|
||
def chunk(self, text: str) -> list[str]:
|
||
if len(text) <= self.max_length:
|
||
return [text] if text else []
|
||
if self.mode == ChunkMode.LENGTH:
|
||
return [c for c in self._chunk_by_length(text) if c]
|
||
if self.mode == ChunkMode.SENTENCE:
|
||
return [c for c in self._chunk_by_sentence(text) if c]
|
||
if self.mode == ChunkMode.MARKDOWN:
|
||
return [c for c in self._chunk_markdown(text) if c]
|
||
return [c for c in self._chunk_by_newline(text) if c]
|
||
|
||
def chunk_with_index(self, text: str) -> Iterator[tuple[int, int, str]]:
|
||
chunks = self.chunk(text)
|
||
total = len(chunks)
|
||
for i, chunk in enumerate(chunks):
|
||
yield i + 1, total, chunk
|
||
|
||
def chunk_as_objects(self, text: str) -> list[TextChunk]:
|
||
chunks = self.chunk(text)
|
||
return [TextChunk(text=c, index=i, is_last=(i == len(chunks) - 1)) for i, c in enumerate(chunks)]
|
||
|
||
def _chunk_by_length(self, text: str) -> list[str]:
|
||
chunks: list[str] = []
|
||
remaining = text
|
||
while len(remaining) > self.max_length:
|
||
split_at = self._find_split_point(remaining)
|
||
chunks.append(remaining[:split_at].rstrip())
|
||
remaining = remaining[split_at:].lstrip()
|
||
if remaining:
|
||
chunks.append(remaining)
|
||
return chunks
|
||
|
||
def _chunk_by_newline(self, text: str) -> list[str]:
|
||
paragraphs = re.split(r"\n{2,}", text)
|
||
chunks: list[str] = []
|
||
for para in paragraphs:
|
||
if len(para) <= self.max_length:
|
||
if para:
|
||
chunks.append(para)
|
||
continue
|
||
for sub in self._chunk_by_length(para):
|
||
chunks.append(sub)
|
||
return chunks
|
||
|
||
def _chunk_by_sentence(self, text: str) -> list[str]:
|
||
sentences = re.split(r"(?<=[。！？.!?])\s*", text)
|
||
chunks: list[str] = []
|
||
buffer = ""
|
||
for sentence in sentences:
|
||
if not sentence.strip():
|
||
buffer += sentence
|
||
continue
|
||
candidate = buffer + sentence
|
||
if len(candidate) <= self.max_length:
|
||
buffer = candidate
|
||
else:
|
||
if buffer:
|
||
chunks.append(buffer)
|
||
if len(sentence) > self.max_length:
|
||
for sub in self._chunk_by_length(sentence):
|
||
chunks.append(sub)
|
||
buffer = ""
|
||
else:
|
||
buffer = sentence
|
||
if buffer:
|
||
chunks.append(buffer)
|
||
return chunks
|
||
|
||
def _chunk_markdown(self, text: str) -> list[str]:
|
||
segments = self._segment_markdown(text)
|
||
chunks: list[str] = []
|
||
buffer = ""
|
||
for seg in segments:
|
||
candidate = buffer + ("\n" if buffer else "") + seg if buffer else seg
|
||
if len(candidate) <= self.max_length:
|
||
buffer = candidate
|
||
else:
|
||
if buffer:
|
||
chunks.append(buffer)
|
||
if len(seg) > self.max_length:
|
||
for sub in self._chunk_overlong_md(seg):
|
||
chunks.append(sub)
|
||
buffer = ""
|
||
else:
|
||
buffer = seg
|
||
if buffer:
|
||
chunks.append(buffer)
|
||
return chunks
|
||
|
||
def _segment_markdown(self, text: str) -> list[str]:
|
||
lines = text.split("\n")
|
||
segments: list[str] = []
|
||
current: list[str] = []
|
||
in_fence = False
|
||
in_table = False
|
||
in_list = False
|
||
|
||
def flush():
|
||
nonlocal current, in_table, in_list
|
||
if current:
|
||
segments.append("\n".join(current))
|
||
current = []
|
||
in_table = False
|
||
in_list = False
|
||
|
||
for line in lines:
|
||
is_fence = bool(_MD_FENCE.match(line))
|
||
is_table_sep = bool(_MD_TABLE_SEP.match(line))
|
||
is_ordered = bool(_MD_ORDERED_LIST.match(line))
|
||
is_unordered = bool(_MD_UNORDERED_LIST.match(line))
|
||
is_heading = bool(_MD_HEADING.match(line))
|
||
|
||
if is_fence:
|
||
if in_fence:
|
||
current.append(line)
|
||
flush()
|
||
in_fence = False
|
||
continue
|
||
flush()
|
||
in_fence = True
|
||
current.append(line)
|
||
continue
|
||
|
||
if in_fence:
|
||
current.append(line)
|
||
continue
|
||
|
||
if is_table_sep and current and "|" in current[-1]:
|
||
flush()
|
||
in_table = True
|
||
current.append(line)
|
||
continue
|
||
|
||
if in_table and "|" in line:
|
||
current.append(line)
|
||
continue
|
||
|
||
if in_table:
|
||
flush()
|
||
|
||
if is_ordered or is_unordered:
|
||
if not in_list:
|
||
flush()
|
||
in_list = True
|
||
current.append(line)
|
||
continue
|
||
|
||
if is_heading:
|
||
flush()
|
||
current.append(line)
|
||
continue
|
||
|
||
if in_list:
|
||
flush()
|
||
|
||
if line.strip() == "":
|
||
flush()
|
||
continue
|
||
|
||
current.append(line)
|
||
|
||
if len("\n".join(current)) > self.max_length * 2:
|
||
flush()
|
||
|
||
flush()
|
||
return segments
|
||
|
||
def _chunk_overlong_md(self, text: str) -> list[str]:
|
||
lines = text.split("\n")
|
||
chunks: list[str] = []
|
||
buffer = ""
|
||
|
||
for line in lines:
|
||
candidate = buffer + ("\n" if buffer else "") + line if buffer else line
|
||
if len(candidate) <= self.max_length:
|
||
buffer = candidate
|
||
continue
|
||
|
||
if buffer:
|
||
chunks.append(buffer)
|
||
buffer = ""
|
||
|
||
if len(line) <= self.max_length:
|
||
buffer = line
|
||
else:
|
||
for sub in self._chunk_by_length(line):
|
||
if len(sub) <= self.max_length:
|
||
if buffer:
|
||
chunks.append(buffer)
|
||
buffer = sub
|
||
else:
|
||
if buffer:
|
||
chunks.append(buffer)
|
||
buffer = ""
|
||
chunks.append(sub)
|
||
|
||
if buffer:
|
||
chunks.append(buffer)
|
||
return chunks
|
||
|
||
def _find_split_point(self, text: str) -> int:
|
||
limit = self.max_length
|
||
slice_end = min(limit, len(text))
|
||
for sep in _NATURAL_BOUNDARIES:
|
||
idx = text.rfind(sep, 0, slice_end)
|
||
if idx > slice_end // 2:
|
||
return idx + len(sep.rstrip())
|
||
space_idx = text.rfind(" ", 0, slice_end)
|
||
if space_idx > slice_end // 2:
|
||
return space_idx + 1
|
||
while slice_end > 0 and (ord(text[slice_end - 1]) & 0xC0) == 0x80:
|
||
slice_end -= 1
|
||
return slice_end if slice_end > 0 else limit
|