85 lines
2.4 KiB
Python
85 lines
2.4 KiB
Python
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
from typing import Literal
|
|||
|
|
|
|||
|
|
NATURAL_BOUNDARIES = [
|
|||
|
|
"\n\n",
|
|||
|
|
"\n",
|
|||
|
|
"。",
|
|||
|
|
"！",
|
|||
|
|
"？",
|
|||
|
|
"；",
|
|||
|
|
"：",
|
|||
|
|
". ",
|
|||
|
|
"! ",
|
|||
|
|
"? ",
|
|||
|
|
]
|
|||
|
|
|
|||
|
|
|
|||
|
|
class WeChatChunker:
|
|||
|
|
def chunk_text(self, text: str, limit: int = 2048) -> list[str]:
|
|||
|
|
if len(text) <= limit:
|
|||
|
|
return [text]
|
|||
|
|
|
|||
|
|
chunks: list[str] = []
|
|||
|
|
remaining = text
|
|||
|
|
while len(remaining) > limit:
|
|||
|
|
cut = self._find_boundary(remaining, limit)
|
|||
|
|
if cut <= 0:
|
|||
|
|
cut = limit
|
|||
|
|
chunks.append(remaining[:cut])
|
|||
|
|
remaining = remaining[cut:]
|
|||
|
|
if remaining:
|
|||
|
|
chunks.append(remaining)
|
|||
|
|
return chunks
|
|||
|
|
|
|||
|
|
def chunk_markdown(self, text: str, limit: int = 2048) -> list[str]:
|
|||
|
|
if len(text) <= limit:
|
|||
|
|
return [text]
|
|||
|
|
|
|||
|
|
chunks: list[str] = []
|
|||
|
|
remaining = text
|
|||
|
|
code_block = False
|
|||
|
|
|
|||
|
|
while len(remaining) > limit:
|
|||
|
|
if "```" in remaining[:limit]:
|
|||
|
|
fence_idx = remaining.find("```")
|
|||
|
|
if fence_idx >= 0:
|
|||
|
|
code_block = not code_block
|
|||
|
|
if code_block:
|
|||
|
|
close_idx = remaining.find("```", fence_idx + 3)
|
|||
|
|
block_end = close_idx + 3 if close_idx >= 0 else limit
|
|||
|
|
if block_end <= limit:
|
|||
|
|
cut = block_end
|
|||
|
|
else:
|
|||
|
|
cut = limit
|
|||
|
|
else:
|
|||
|
|
cut = self._find_boundary(remaining, limit)
|
|||
|
|
chunks.append(remaining[:cut])
|
|||
|
|
remaining = remaining[cut:]
|
|||
|
|
continue
|
|||
|
|
|
|||
|
|
cut = self._find_boundary(remaining, limit)
|
|||
|
|
if cut <= 0:
|
|||
|
|
cut = limit
|
|||
|
|
chunks.append(remaining[:cut])
|
|||
|
|
remaining = remaining[cut:]
|
|||
|
|
|
|||
|
|
if remaining:
|
|||
|
|
chunks.append(remaining)
|
|||
|
|
return chunks
|
|||
|
|
|
|||
|
|
def chunk(self, text: str, limit: int, mode: Literal["text", "markdown"] = "text") -> list[str]:
|
|||
|
|
if mode == "markdown":
|
|||
|
|
return self.chunk_markdown(text, limit)
|
|||
|
|
return self.chunk_text(text, limit)
|
|||
|
|
|
|||
|
|
@staticmethod
|
|||
|
|
def _find_boundary(text: str, limit: int) -> int:
|
|||
|
|
best = -1
|
|||
|
|
for boundary in NATURAL_BOUNDARIES:
|
|||
|
|
idx = text.rfind(boundary, 0, limit)
|
|||
|
|
if idx > best:
|
|||
|
|
best = idx
|
|||
|
|
return best + 1 if best >= 0 else 0
|