73 lines
1.9 KiB
Python
73 lines
1.9 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
|
||
|
|
def chunk_text(text: str, limit: int, mode: str = "paragraph") -> list[str]:
|
||
|
|
if len(text) <= limit:
|
||
|
|
return [text]
|
||
|
|
|
||
|
|
if mode == "length":
|
||
|
|
return _chunk_by_length(text, limit)
|
||
|
|
if mode == "newline":
|
||
|
|
return _chunk_by_newline(text, limit)
|
||
|
|
return _chunk_by_paragraph(text, limit)
|
||
|
|
|
||
|
|
|
||
|
|
def _chunk_by_paragraph(text: str, limit: int) -> list[str]:
|
||
|
|
chunks = []
|
||
|
|
paragraphs = text.split("\n\n")
|
||
|
|
|
||
|
|
for para in paragraphs:
|
||
|
|
if len(para) <= limit:
|
||
|
|
chunks.append(para)
|
||
|
|
continue
|
||
|
|
|
||
|
|
sentences = para.replace("\n", " ").split(". ")
|
||
|
|
current = ""
|
||
|
|
|
||
|
|
for sentence in sentences:
|
||
|
|
sep = ". " if current else ""
|
||
|
|
candidate = f"{current}{sep}{sentence}"
|
||
|
|
|
||
|
|
if current and len(candidate) + 1 > limit:
|
||
|
|
chunks.append(current + ".")
|
||
|
|
current = sentence
|
||
|
|
else:
|
||
|
|
current = candidate
|
||
|
|
|
||
|
|
if current:
|
||
|
|
if not current.endswith("."):
|
||
|
|
current += "."
|
||
|
|
chunks.append(current)
|
||
|
|
|
||
|
|
merged = []
|
||
|
|
for chunk in chunks:
|
||
|
|
if merged and len(merged[-1]) + len(chunk) + 2 <= limit:
|
||
|
|
merged[-1] = f"{merged[-1]}\n\n{chunk}"
|
||
|
|
else:
|
||
|
|
merged.append(chunk)
|
||
|
|
|
||
|
|
return merged
|
||
|
|
|
||
|
|
|
||
|
|
def _chunk_by_length(text: str, limit: int) -> list[str]:
|
||
|
|
chunks = []
|
||
|
|
for i in range(0, len(text), limit):
|
||
|
|
chunks.append(text[i : i + limit])
|
||
|
|
return chunks
|
||
|
|
|
||
|
|
|
||
|
|
def _chunk_by_newline(text: str, limit: int) -> list[str]:
|
||
|
|
lines = text.split("\n")
|
||
|
|
chunks = []
|
||
|
|
current = ""
|
||
|
|
for line in lines:
|
||
|
|
candidate = f"{current}\n{line}" if current else line
|
||
|
|
if len(candidate) > limit and current:
|
||
|
|
chunks.append(current)
|
||
|
|
current = line
|
||
|
|
else:
|
||
|
|
current = candidate
|
||
|
|
if current:
|
||
|
|
chunks.append(current)
|
||
|
|
return chunks
|