61 lines
1.6 KiB
Python
61 lines
1.6 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import logging
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
def chunk_text(text: str, limit: int = 4000, mode: str = "markdown") -> list[str]:
|
||
|
|
if len(text) <= limit:
|
||
|
|
return [text]
|
||
|
|
|
||
|
|
if mode == "markdown":
|
||
|
|
return _chunk_markdown(text, limit)
|
||
|
|
elif mode == "newline":
|
||
|
|
return _chunk_by_newline(text, limit)
|
||
|
|
return _chunk_by_length(text, limit)
|
||
|
|
|
||
|
|
|
||
|
|
def _chunk_markdown(text: str, limit: int) -> list[str]:
|
||
|
|
chunks: list[str] = []
|
||
|
|
remainder = text
|
||
|
|
|
||
|
|
while len(remainder) > limit:
|
||
|
|
cut = remainder.rfind("\n\n", 0, limit)
|
||
|
|
if cut < limit // 2:
|
||
|
|
cut = remainder.rfind("\n", 0, limit)
|
||
|
|
if cut < limit // 4:
|
||
|
|
cut = remainder.rfind(". ", 0, limit)
|
||
|
|
if cut < 0:
|
||
|
|
cut = limit
|
||
|
|
|
||
|
|
chunk = remainder[: cut + 1].rstrip()
|
||
|
|
chunks.append(chunk)
|
||
|
|
remainder = remainder[cut + 1 :].lstrip()
|
||
|
|
|
||
|
|
if remainder.strip():
|
||
|
|
chunks.append(remainder.strip())
|
||
|
|
|
||
|
|
return chunks
|
||
|
|
|
||
|
|
|
||
|
|
def _chunk_by_newline(text: str, limit: int) -> list[str]:
|
||
|
|
chunks: list[str] = []
|
||
|
|
current = ""
|
||
|
|
for line in text.split("\n"):
|
||
|
|
if len(current) + len(line) + 1 > limit and current:
|
||
|
|
chunks.append(current.rstrip())
|
||
|
|
current = line
|
||
|
|
else:
|
||
|
|
current = f"{current}\n{line}".strip()
|
||
|
|
if current:
|
||
|
|
chunks.append(current)
|
||
|
|
return chunks
|
||
|
|
|
||
|
|
|
||
|
|
def _chunk_by_length(text: str, limit: int) -> list[str]:
|
||
|
|
chunks: list[str] = []
|
||
|
|
for i in range(0, len(text), limit):
|
||
|
|
chunks.append(text[i : i + limit])
|
||
|
|
return chunks
|