use std::ops::Range;
use async_trait::async_trait;
use crate::ports::{Prompt, TextModel, TextModelError};
pub fn split_text_chunks(text: &str, max_chars: usize) -> Vec<String> {
assert!(max_chars > 0, "max_chars must be at least 1, got 0");
if text.is_empty() {
return Vec::new();
}
if text.chars().all(char::is_whitespace) {
return vec![text.to_owned()];
}
let mut chunks: Vec<String> = Vec::new();
let mut current = String::new();
let mut current_chars = 0usize;
for item in scan_items(text) {
let item_text = &text[item.range];
let item_chars = item_text.chars().count();
if current_chars + item_chars <= max_chars {
current.push_str(item_text);
current_chars += item_chars;
continue;
}
if current_chars > 0 {
chunks.push(std::mem::take(&mut current));
current_chars = 0;
}
if item_chars <= max_chars {
current.push_str(item_text);
current_chars = item_chars;
} else if item.fenced {
chunks.push(item_text.to_owned());
} else {
chunks.extend(split_oversized(item_text, max_chars));
}
}
if current_chars > 0 {
chunks.push(current);
}
chunks
}
pub fn reassemble_text_chunks(chunks: impl IntoIterator<Item: AsRef<str>>) -> String {
let mut text = String::new();
for chunk in chunks {
text.push_str(chunk.as_ref());
}
text
}
#[async_trait]
pub trait TextModelChunkedExt: TextModel {
async fn complete_chunked(
&self,
text: &str,
max_chars: usize,
prompt: impl for<'a> FnMut(usize, usize, &'a str) -> Prompt + Send,
) -> Result<String, TextModelError>;
}
#[async_trait]
impl<M: TextModel + ?Sized> TextModelChunkedExt for M {
async fn complete_chunked(
&self,
text: &str,
max_chars: usize,
mut prompt: impl for<'a> FnMut(usize, usize, &'a str) -> Prompt + Send,
) -> Result<String, TextModelError> {
let chunks = split_text_chunks(text, max_chars);
let total = chunks
.iter()
.filter(|chunk| !chunk.chars().all(char::is_whitespace))
.count();
let mut answers = Vec::with_capacity(chunks.len());
let mut index = 0usize;
for chunk in &chunks {
if chunk.chars().all(char::is_whitespace) {
answers.push(chunk.clone());
continue;
}
let completion = self.complete(&prompt(index, total, chunk)).await?;
answers.push(completion.text);
index += 1;
}
Ok(reassemble_text_chunks(answers))
}
}
struct Item {
range: Range<usize>,
fenced: bool,
}
#[derive(Clone, Copy, PartialEq, Eq)]
enum TextKind {
Blank,
Paragraph,
}
fn scan_items(text: &str) -> Vec<Item> {
let mut items: Vec<Item> = Vec::new();
let mut kind: Option<TextKind> = None;
let mut start = 0;
let mut end = 0;
let mut fence: Option<(u8, usize)> = None;
let mut fence_start = 0;
let mut i = 0;
while i < text.len() {
let line_start = i;
let line_end = match text[i..].find('\n') {
Some(at) => i + at + 1,
None => text.len(),
};
let trimmed = text[line_start..line_end].trim();
if let Some((marker, count)) = fence {
if closes_fence(trimmed, marker, count) {
fence = None;
items.push(Item {
range: fence_start..line_end,
fenced: true,
});
kind = None;
}
} else if trimmed.is_empty() {
if kind != Some(TextKind::Blank) {
push_text(&mut items, kind, start, end);
kind = Some(TextKind::Blank);
start = line_start;
}
} else if let Some((marker, count)) = opens_fence(trimmed) {
push_text(&mut items, kind, start, end);
kind = None;
fence = Some((marker, count));
fence_start = line_start;
} else if kind != Some(TextKind::Paragraph) {
push_text(&mut items, kind, start, end);
kind = Some(TextKind::Paragraph);
start = line_start;
}
end = line_end;
i = line_end;
}
if fence.is_some() {
items.push(Item {
range: fence_start..text.len(),
fenced: true,
});
} else {
push_text(&mut items, kind, start, end);
}
items
}
fn push_text(items: &mut Vec<Item>, kind: Option<TextKind>, start: usize, end: usize) {
if matches!(kind, Some(TextKind::Blank | TextKind::Paragraph)) && end > start {
items.push(Item {
range: start..end,
fenced: false,
});
}
}
fn opens_fence(trimmed: &str) -> Option<(u8, usize)> {
let bytes = trimmed.as_bytes();
let marker = *bytes.first()?;
if marker != b'`' && marker != b'~' {
return None;
}
let count = bytes.iter().take_while(|&&byte| byte == marker).count();
(count >= 3).then_some((marker, count))
}
fn closes_fence(trimmed: &str, marker: u8, count: usize) -> bool {
let bytes = trimmed.as_bytes();
bytes.len() >= count && bytes.iter().all(|&byte| byte == marker)
}
fn split_oversized(text: &str, max_chars: usize) -> Vec<String> {
let mut pieces: Vec<String> = Vec::new();
let mut current = String::new();
let mut current_chars = 0usize;
let mut i = 0;
while i < text.len() {
let line_end = match text[i..].find('\n') {
Some(at) => i + at + 1,
None => text.len(),
};
let line = &text[i..line_end];
i = line_end;
let line_chars = line.chars().count();
if line_chars > max_chars {
if current_chars > 0 {
pieces.push(std::mem::take(&mut current));
current_chars = 0;
}
let mut cut = String::new();
let mut count = 0usize;
for ch in line.chars() {
cut.push(ch);
count += 1;
if count == max_chars {
pieces.push(std::mem::take(&mut cut));
count = 0;
}
}
if !cut.is_empty() {
pieces.push(cut);
}
} else if current_chars + line_chars <= max_chars {
current.push_str(line);
current_chars += line_chars;
} else {
pieces.push(std::mem::take(&mut current));
current.push_str(line);
current_chars = line_chars;
}
}
if current_chars > 0 {
pieces.push(current);
}
pieces
}
#[cfg(test)]
mod tests {
use super::*;
fn corpus() -> Vec<(&'static str, &'static str)> {
vec![
(
"plain paragraphs",
"The quick brown fox.\n\nSecond paragraph here.\n\nThird one.\n",
),
(
"crlf",
"para one\r\n\r\npara two\r\n```rust\r\nlet x = 1;\r\n```\r\n\r\nafter\r\n",
),
(
"backtick fences with info string and blank lines inside",
"intro\n\n```rust\nfn main() {}\n\nlet x = 1;\n```\n\noutro\n",
),
("tilde fence", "before\n\n~~~text\nbody\n~~~\nafter"),
("unclosed fence", "text\n\n```python\nprint(1)\nprint(2)"),
(
"long fence",
"```js\nconst a = 1;\nconst b = 2;\nconst c = 3;\nconst d = 4;\nconst e = 5;\n```\n",
),
(
"four-backtick fence around a three-backtick line",
"````\n```\nnot a close\n```\nstill inside\n````\ntail\n",
),
(
"indented fence",
"before\n\n ```js\n const a = 1;\n ```\nafter\n",
),
("blank runs at the edges", "\n\n\nlead\n\ntrail\n\n\n"),
("whitespace only", " \n\t\n \n"),
("empty", ""),
("one line", "one line only"),
(
"multibyte",
"héllo wörld: 日本語のテキストもここに。Émile résume.",
),
(
"paragraph, fence, paragraph",
"alpha\n\n```\ncode\n```\n\nomega\n",
),
(
"many short lines",
"aaa bbb\nccc ddd\neee fff\nggg hhh\niii jjj\nkkk lll\n",
),
]
}
#[test]
fn every_fixture_round_trips_byte_for_byte_at_many_bounds() {
for (name, text) in corpus() {
for max_chars in [1, 2, 3, 5, 8, 13, 40, 10_000] {
let chunks = split_text_chunks(text, max_chars);
assert_eq!(
reassemble_text_chunks(&chunks),
text,
"{name} at max_chars {max_chars}"
);
}
}
}
#[test]
fn every_chunk_respects_the_bound_except_atomic_fences() {
for (name, text) in corpus() {
if text.chars().all(char::is_whitespace) {
continue; }
for max_chars in [1, 2, 3, 5, 8, 13, 40] {
for chunk in split_text_chunks(text, max_chars) {
let chars = chunk.chars().count();
assert!(
chars <= max_chars || is_single_atomic_fence(&chunk),
"{name} at max_chars {max_chars}: {chars}-char chunk {chunk:?}"
);
}
}
}
}
#[test]
fn every_chunk_of_a_fence_heavy_fixture_has_balanced_fences() {
for (name, text) in corpus() {
for max_chars in [1, 5, 8, 13, 40] {
let chunks = split_text_chunks(text, max_chars);
let open_at_end: Vec<bool> = chunks
.iter()
.map(|chunk| open_fence_at_end(chunk).is_some())
.collect();
let open_count = open_at_end.iter().filter(|open| **open).count();
assert!(
open_count == 0 || (open_count == 1 && open_at_end[open_at_end.len() - 1]),
"{name} at max_chars {max_chars}: a fence crosses a chunk: {chunks:?}"
);
}
}
}
fn is_single_atomic_fence(chunk: &str) -> bool {
let lines: Vec<&str> = chunk.lines().collect();
let Some((marker, count)) = lines
.first()
.copied()
.and_then(|line| opens_fence(line.trim_start()))
else {
return false;
};
match lines[1..]
.iter()
.position(|line| closes_fence(line.trim(), marker, count))
{
Some(at) => at + 1 == lines.len() - 1,
None => true,
}
}
fn open_fence_at_end(chunk: &str) -> Option<(u8, usize)> {
let mut open: Option<(u8, usize)> = None;
for line in chunk.lines() {
if let Some((marker, count)) = open {
if closes_fence(line.trim(), marker, count) {
open = None;
}
} else if let Some((marker, count)) = opens_fence(line.trim_start()) {
open = Some((marker, count));
}
}
open
}
#[test]
fn breaks_land_on_blank_line_runs_not_inside_paragraphs() {
let text = "alpha bravo\n\ncharlie delta\n\necho foxtrot";
let chunks = split_text_chunks(text, 30);
assert_eq!(
chunks,
vec!["alpha bravo\n\ncharlie delta\n\n", "echo foxtrot"],
"the break is the blank run between the second and third paragraph"
);
}
#[test]
fn a_crlf_blank_line_counts_as_blank_and_travels_with_its_paragraph() {
let text = "para one\r\n\r\npara two is longer\r\n";
let chunks = split_text_chunks(text, 14);
assert_eq!(
chunks[0], "para one\r\n\r\n",
"the CRLF blank run is a boundary"
);
assert_eq!(
chunks[1..],
["para two is lo", "nger\r\n"],
"the oversized paragraph falls back to line and char splits"
);
assert_eq!(reassemble_text_chunks(&chunks), text);
}
#[test]
fn blank_runs_at_the_edges_stay_with_their_neighbours() {
let text = "\n\n\nlead\n\ntrail\n\n\n";
let chunks = split_text_chunks(text, 9);
assert_eq!(
chunks,
vec!["\n\n\nlead\n\n", "trail\n\n\n"],
"the leading run rides with the first block, the trailing run with the last"
);
}
#[test]
fn a_leading_blank_run_that_cannot_fit_stands_alone_rather_than_bursting_the_bound() {
let chunks = split_text_chunks("\n\n\nfirst paragraph body", 20);
assert_eq!(chunks, vec!["\n\n\n", "first paragraph body"]);
assert_eq!(
reassemble_text_chunks(&chunks),
"\n\n\nfirst paragraph body"
);
}
#[test]
fn a_fence_is_never_split_even_when_it_exceeds_the_bound() {
let fence = "```js\nconst a = 1;\nconst b = 2;\nconst c = 3;\nconst d = 4;\n```\n";
let text = format!("before\n\n{fence}after\n");
let chunks = split_text_chunks(&text, 10);
assert!(
chunks.contains(&fence.to_owned()),
"the oversized fence stands alone as one chunk: {chunks:?}"
);
assert_eq!(reassemble_text_chunks(&chunks), text);
}
#[test]
fn a_four_backtick_fence_does_not_close_at_an_inner_three_backtick_line() {
let text = "````\n```\ninner\n```\nmore\n````\ntail\n";
for max_chars in [4, 6, 10, 28] {
let chunks = split_text_chunks(text, max_chars);
let fence = chunks
.iter()
.find(|chunk| chunk.starts_with("````"))
.unwrap_or_else(|| {
panic!("the ```` block is one chunk at {max_chars}: {chunks:?}")
});
assert!(
fence.contains("```\ninner") && fence.ends_with("````\n"),
"the inner ``` did not close the ```` fence at {max_chars}: {fence:?}"
);
assert_eq!(reassemble_text_chunks(&chunks), text);
}
assert_eq!(split_text_chunks(text, 40), vec![text]);
}
#[test]
fn an_unclosed_fence_runs_to_the_end_of_the_input() {
let text = "intro\n\n```python\nprint(1)\nprint(2)";
let chunks = split_text_chunks(text, 8);
assert_eq!(reassemble_text_chunks(&chunks), text);
let last = chunks
.last()
.unwrap_or_else(|| panic!("chunks: {chunks:?}"));
assert!(last.starts_with("```python\n"), "{last:?}");
assert!(
is_single_atomic_fence(last),
"the unclosed fence is one atomic chunk: {last:?}"
);
}
#[test]
fn blank_lines_inside_a_fence_stay_inside_it() {
let text = "```rust\nfn main() {}\n\n\nlet x = 1;\n```\n";
let chunks = split_text_chunks(text, 7);
assert_eq!(
chunks,
vec![text],
"the blank lines did not end the fence, so the whole block is one chunk"
);
}
#[test]
fn an_oversized_paragraph_splits_at_line_boundaries() {
let chunks = split_text_chunks("aaa bbb\nccc ddd\neee fff", 8);
assert_eq!(chunks, vec!["aaa bbb\n", "ccc ddd\n", "eee fff"]);
}
#[test]
fn an_oversized_line_is_hard_cut_on_char_boundaries() {
let text = "héllo 日本語 wörld";
let chunks = split_text_chunks(text, 3);
assert_eq!(reassemble_text_chunks(&chunks), text);
assert_eq!(
chunks,
vec!["hél", "lo ", "日本語", " wö", "rld"],
"every cut is every third char, never mid-character"
);
}
#[test]
fn empty_input_has_no_chunks() {
assert_eq!(split_text_chunks("", 100), Vec::<String>::new());
}
#[test]
fn whitespace_only_input_is_one_chunk() {
let text = " \n\t\n ";
assert_eq!(split_text_chunks(text, 2), vec![text]);
}
#[test]
fn text_with_no_split_points_is_one_chunk() {
assert_eq!(
split_text_chunks("one line only", 100),
vec!["one line only"]
);
}
#[test]
fn a_document_that_fits_is_one_chunk() {
let text = "para\n\n```js\ncode\n```\n\nfin\n";
assert_eq!(split_text_chunks(text, 1_000), vec![text]);
}
#[test]
#[should_panic(expected = "max_chars")]
fn a_zero_bound_panics() {
split_text_chunks("text", 0);
}
#[test]
fn reassembly_is_plain_concatenation_in_order() {
assert_eq!(reassemble_text_chunks(["a", "b", "c"]), "abc");
assert_eq!(
reassemble_text_chunks(vec![String::new(), "x".to_owned()]),
"x"
);
assert_eq!(reassemble_text_chunks([""]), "");
assert_eq!(reassemble_text_chunks(Vec::<&str>::new()), "");
}
}