from __future__ import annotations import re from typing import Any, Literal ChunkMode = Literal["length", "newline", "paragraph"] TELEGRAM_MAX_TEXT_LENGTH = 4096 _SENTENCE_BOUNDARY = re.compile(r"[。!?.!?\n]") def _chunk_by_length(text: str, limit: int) -> list[str]: chunks: list[str] = [] start = 0 while start < len(text): end = min(start + limit, len(text)) if end < len(text): split_at = _find_split_point(text, limit, start, end) chunks.append(text[start:split_at].strip()) start = split_at else: chunks.append(text[start:end].strip()) start = end return chunks or [text] def _chunk_by_newline(text: str, limit: int) -> list[str]: chunks: list[str] = [] lines = text.split("\n") current = "" for line in lines: if len(line) > limit: if current: chunks.append(current) sub_chunks = _chunk_by_length(line, limit) if sub_chunks: current = sub_chunks.pop() chunks.extend(sub_chunks) else: current = "" elif len(current) + len(line) + (1 if current else 0) <= limit: current = f"{current}\n{line}" if current else line else: chunks.append(current) current = line if current: chunks.append(current) return chunks or [text] def _chunk_by_paragraph(text: str, limit: int) -> list[str]: if len(text) <= limit: return [text] chunks: list[str] = [] paragraphs = text.split("\n\n") current = "" for para in paragraphs: if len(current) + len(para) + 2 <= limit: current = f"{current}\n\n{para}" if current else para else: if current: chunks.append(current) if len(para) > limit: sub_chunks = _chunk_long_paragraph(para, limit) if sub_chunks: if sub_chunks[-1]: current = sub_chunks.pop() else: sub_chunks.pop() current = "" chunks.extend(sub_chunks) else: current = "" else: current = para if current: chunks.append(current) return chunks or [text] def chunk_text( text: str, limit: int = TELEGRAM_MAX_TEXT_LENGTH, mode: ChunkMode = "paragraph", ) -> list[str]: modes: dict[ChunkMode, callable] = { "length": _chunk_by_length, "newline": _chunk_by_newline, "paragraph": _chunk_by_paragraph, } chunker = modes.get(mode, _chunk_by_paragraph) return chunker(text, limit) def chunk_mode_from_config(config: dict[str, Any] | None = None) -> ChunkMode: cfg = config or {} streaming_cfg = cfg.get("streaming", {}) if isinstance(streaming_cfg, dict): return streaming_cfg.get("chunk_mode", "paragraph") return cfg.get("chunking_mode", "paragraph") def _chunk_long_paragraph(text: str, limit: int) -> list[str]: chunks: list[str] = [] while len(text) > limit: split_at = _find_split_point(text, limit, 0, limit) chunks.append(text[:split_at].rstrip()) text = text[split_at:].lstrip() chunks.append(text) return chunks def _find_split_point(text: str, limit: int, lo: int = 0, hi: int | None = None) -> int: look_start = max(lo, lo or limit // 2) look_end = hi or limit candidates = [m.start() for m in _SENTENCE_BOUNDARY.finditer(text, look_start, look_end)] if candidates: return candidates[-1] + 1 newline = text.rfind("\n", look_start, look_end) if newline != -1: return newline + 1 space = text.rfind(" ", look_start, look_end) if space != -1: return space + 1 return look_end