use std::collections::BTreeMap;
use std::ops::Range;
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use super::chunk::{self, chunk_text, ChunkBoundary, ChunkPolicy};
use super::log_normalize::mask_volatile_prefix;
use super::model::ContextFragment;
use crate::error::MemoryError;
use crate::limits::{MAX_FRAGMENTS, MAX_FRAGMENT_BYTES, MAX_TRANSCRIPT_BYTES};
const MIN_LOG_RUN_LINES: usize = 8;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(rename_all = "lowercase")]
#[non_exhaustive] pub enum SegmentFormat {
Auto,
Plain,
Jsonl,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, JsonSchema)]
#[serde(rename_all = "lowercase")]
#[non_exhaustive] pub enum SegmentKind {
Body,
Code,
Log,
}
impl SegmentKind {
fn fragment_kind(self) -> Option<&'static str> {
match self {
Self::Body => None,
Self::Code => Some("code"),
Self::Log => Some("log"),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
#[serde(default)]
#[schemars(transform = crate::schema::strip_int_formats)]
pub struct SegmentationPolicy {
pub format: SegmentFormat,
pub min_segment_bytes: usize,
pub cache_system_turn: bool,
}
impl Default for SegmentationPolicy {
fn default() -> Self {
Self {
format: SegmentFormat::Auto,
min_segment_bytes: 256,
cache_system_turn: true,
}
}
}
#[derive(Debug, Clone)]
pub struct TranscriptSegment {
pub fragment: ContextFragment,
pub turn: usize,
pub role: Option<String>,
pub kind: SegmentKind,
pub byte_start: usize,
pub byte_end: usize,
}
#[derive(Debug, Clone)]
pub struct SegmentationOutcome {
pub format_detected: SegmentFormat,
pub segments: Vec<TranscriptSegment>,
pub merged_segments: usize,
}
pub fn segment_transcript(
text: &str,
policy: &SegmentationPolicy,
) -> Result<SegmentationOutcome, MemoryError> {
if text.len() > MAX_TRANSCRIPT_BYTES {
return Err(MemoryError::ContextOverLimit(format!(
"transcript of {} bytes exceeds the cap of {MAX_TRANSCRIPT_BYTES} bytes",
text.len()
)));
}
let (format_detected, pieces) = detect_and_segment(text, policy.format)?;
reject_oversized_fences(&pieces)?;
let pieces = resplit_oversized_bodies(text, pieces);
let pieces_before_merge = pieces.len();
let merged = merge_tiny(pieces, policy.min_segment_bytes);
if merged.len() > MAX_FRAGMENTS {
return Err(MemoryError::ContextOverLimit(format!(
"transcript segmented into {} fragments, exceeding the cap of {MAX_FRAGMENTS} — \
raise segmentation.min_segment_bytes to merge more small segments",
merged.len()
)));
}
let merged_segments = pieces_before_merge - merged.len();
let segments = merged
.into_iter()
.map(|piece| build_segment(text, piece, policy))
.collect();
Ok(SegmentationOutcome {
format_detected,
segments,
merged_segments,
})
}
struct RawPiece {
kind: SegmentKind,
range: Range<usize>,
turn: usize,
role: Option<String>,
content_override: Option<String>,
}
fn detect_and_segment(
text: &str,
requested: SegmentFormat,
) -> Result<(SegmentFormat, Vec<RawPiece>), MemoryError> {
match requested {
SegmentFormat::Plain => Ok((SegmentFormat::Plain, plain_pieces(text))),
SegmentFormat::Jsonl => {
let pieces = jsonl_pieces(text).map_err(MemoryError::SegmentationError)?;
Ok((SegmentFormat::Jsonl, pieces))
}
SegmentFormat::Auto => {
if !text.is_empty() {
if let Ok(pieces) = jsonl_pieces(text) {
return Ok((SegmentFormat::Jsonl, pieces));
}
}
Ok((SegmentFormat::Plain, plain_pieces(text)))
}
}
}
#[derive(Deserialize)]
struct JsonlLine {
role: String,
content: String,
}
fn jsonl_pieces(text: &str) -> Result<Vec<RawPiece>, String> {
let mut pieces: Vec<RawPiece> = Vec::new();
let mut pending_prefix_start: Option<usize> = None;
let mut turn = 0_usize;
let mut cursor = 0_usize;
for (line_index, line) in text.split_inclusive('\n').enumerate() {
let start = cursor;
cursor += line.len();
let trimmed = line.trim_end_matches(['\r', '\n']);
if trimmed.is_empty() {
if let Some(last) = pieces.last_mut() {
last.range.end = cursor;
} else {
pending_prefix_start.get_or_insert(start);
}
continue;
}
let parsed: JsonlLine = serde_json::from_str(trimmed).map_err(|err| {
format!(
"jsonl line {}: not a valid {{role, content}} object: {err}",
line_index + 1
)
})?;
let piece_start = pending_prefix_start.take().unwrap_or(start);
pieces.push(RawPiece {
kind: SegmentKind::Body,
range: piece_start..cursor,
turn,
role: Some(parsed.role),
content_override: Some(parsed.content),
});
turn += 1;
}
if pieces.is_empty() {
return Err("no non-blank jsonl line found".to_owned());
}
Ok(pieces)
}
const PLAIN_MARKERS: &[&str] = &[
"System:",
"User:",
"Human:",
"Assistant:",
"AI:",
"Tool:",
"### User",
"### Assistant",
];
fn match_marker(line: &str) -> Option<&'static str> {
PLAIN_MARKERS
.iter()
.find(|marker| line.starts_with(*marker))
.copied()
}
fn marker_role(marker: &str) -> String {
marker
.strip_prefix("### ")
.unwrap_or(marker)
.trim_end_matches(':')
.to_owned()
}
fn plain_turns(text: &str) -> Vec<(Range<usize>, Option<String>)> {
let mut turns = Vec::new();
let mut turn_start = 0_usize;
let mut pending_role: Option<String> = None;
let mut cursor = 0_usize;
for line in text.split_inclusive('\n') {
let line_start = cursor;
if let Some(marker) = match_marker(line) {
if line_start > turn_start {
turns.push((turn_start..line_start, pending_role.take()));
}
pending_role = Some(marker_role(marker));
turn_start = line_start;
}
cursor += line.len();
}
turns.push((turn_start..text.len(), pending_role));
turns
}
fn plain_pieces(text: &str) -> Vec<RawPiece> {
let mut pieces = Vec::new();
for (turn, (range, role)) in plain_turns(text).into_iter().enumerate() {
if range.is_empty() {
continue;
}
for segment in chunk::fence_segments(&text[range.clone()]) {
match segment {
chunk::Segment::Fence(relative) => pieces.push(RawPiece {
kind: SegmentKind::Code,
range: (range.start + relative.start)..(range.start + relative.end),
turn,
role: role.clone(),
content_override: None,
}),
chunk::Segment::Plain(relative) => {
let absolute = (range.start + relative.start)..(range.start + relative.end);
for (kind, sub_range) in log_split(text, absolute) {
pieces.push(RawPiece {
kind,
range: sub_range,
turn,
role: role.clone(),
content_override: None,
});
}
}
}
}
}
pieces
}
fn log_split(text: &str, range: Range<usize>) -> Vec<(SegmentKind, Range<usize>)> {
if range.is_empty() {
return Vec::new();
}
let slice = &text[range.clone()];
let mut lines: Vec<(Range<usize>, &str)> = Vec::new();
let mut cursor = range.start;
for line in slice.split_inclusive('\n') {
let end = cursor + line.len();
lines.push((cursor..end, line));
cursor = end;
}
if lines.is_empty() {
return Vec::new();
}
let trimmed: Vec<&str> = lines
.iter()
.map(|(_, line)| line.trim_end_matches(['\r', '\n']))
.collect();
let mut repeat_counts: BTreeMap<&str, usize> = BTreeMap::new();
for line in &trimmed {
*repeat_counts.entry(line).or_insert(0) += 1;
}
let candidate: Vec<bool> = trimmed
.iter()
.map(|line| {
!line.is_empty() && (mask_volatile_prefix(line).is_some() || repeat_counts[line] > 1)
})
.collect();
let mut pieces = Vec::new();
let mut body_start: Option<usize> = None;
let mut index = 0_usize;
while index < lines.len() {
if candidate[index] {
let run_start = index;
while index < lines.len() && candidate[index] {
index += 1;
}
if index - run_start >= MIN_LOG_RUN_LINES {
if let Some(start) = body_start.take() {
pieces.push((
SegmentKind::Body,
lines[start].0.start..lines[run_start - 1].0.end,
));
}
pieces.push((
SegmentKind::Log,
lines[run_start].0.start..lines[index - 1].0.end,
));
} else if body_start.is_none() {
body_start = Some(run_start);
}
} else {
if body_start.is_none() {
body_start = Some(index);
}
index += 1;
}
}
if let Some(start) = body_start {
pieces.push((
SegmentKind::Body,
lines[start].0.start..lines[lines.len() - 1].0.end,
));
}
pieces
}
fn reject_oversized_fences(pieces: &[RawPiece]) -> Result<(), MemoryError> {
if let Some(piece) = pieces
.iter()
.find(|piece| piece.kind == SegmentKind::Code && piece.range.len() > MAX_FRAGMENT_BYTES)
{
return Err(MemoryError::ContextOverLimit(format!(
"an unsplittable fenced code block of {} bytes exceeds the cap of {MAX_FRAGMENT_BYTES} bytes",
piece.range.len()
)));
}
Ok(())
}
fn resplit_oversized_bodies(text: &str, pieces: Vec<RawPiece>) -> Vec<RawPiece> {
let chunk_policy = ChunkPolicy {
max_chunk_bytes: MAX_FRAGMENT_BYTES,
overlap_bytes: 0,
boundary: ChunkBoundary::Paragraph,
};
pieces
.into_iter()
.flat_map(|piece| resplit_one(text, piece, &chunk_policy))
.collect()
}
fn resplit_one(text: &str, piece: RawPiece, chunk_policy: &ChunkPolicy) -> Vec<RawPiece> {
match piece.kind {
SegmentKind::Body => resplit_body(text, piece, chunk_policy),
SegmentKind::Log => resplit_log(text, piece),
SegmentKind::Code => vec![piece],
}
}
fn resplit_body(text: &str, piece: RawPiece, chunk_policy: &ChunkPolicy) -> Vec<RawPiece> {
let effective_len = piece
.content_override
.as_ref()
.map_or(piece.range.len(), String::len);
if effective_len <= MAX_FRAGMENT_BYTES {
return vec![piece];
}
match &piece.content_override {
Some(content) => {
let chunks = chunk_text(content, chunk_policy);
let weights: Vec<usize> = chunks.iter().map(|chunk| chunk.text.len()).collect();
let ranges = partition_range_by_weight(&piece.range, &weights);
chunks
.into_iter()
.zip(ranges)
.map(|(chunk, range)| RawPiece {
kind: SegmentKind::Body,
range,
turn: piece.turn,
role: piece.role.clone(),
content_override: Some(chunk.text),
})
.collect()
}
None => chunk_text(&text[piece.range.clone()], chunk_policy)
.into_iter()
.map(|chunk| RawPiece {
kind: SegmentKind::Body,
range: (piece.range.start + chunk.byte_range.start)
..(piece.range.start + chunk.byte_range.end),
turn: piece.turn,
role: piece.role.clone(),
content_override: None,
})
.collect(),
}
}
fn partition_range_by_weight(range: &Range<usize>, weights: &[usize]) -> Vec<Range<usize>> {
debug_assert!(!weights.is_empty(), "must have at least one child");
let total: usize = weights.iter().sum::<usize>().max(1);
let span = range.len();
let mut start = range.start;
let mut cumulative = 0_usize;
let last_index = weights.len() - 1;
weights
.iter()
.enumerate()
.map(|(index, weight)| {
cumulative += weight;
let end = if index == last_index {
range.end
} else {
range.start + (span * cumulative) / total
};
let sub_range = start..end;
start = end;
sub_range
})
.collect()
}
fn resplit_log(text: &str, piece: RawPiece) -> Vec<RawPiece> {
if piece.range.len() <= MAX_FRAGMENT_BYTES {
return vec![piece];
}
let hard_split_policy = ChunkPolicy {
max_chunk_bytes: MAX_FRAGMENT_BYTES,
overlap_bytes: 0,
boundary: ChunkBoundary::Fixed,
};
let mut result = Vec::new();
let mut chunk_start = piece.range.start;
let mut cursor = piece.range.start;
for line in text[piece.range.clone()].split_inclusive('\n') {
let line_start = cursor;
let line_end = line_start + line.len();
cursor = line_end;
if line_end - line_start > MAX_FRAGMENT_BYTES {
if chunk_start < line_start {
result.push(log_piece(&piece, chunk_start..line_start));
}
for hard in chunk_text(&text[line_start..line_end], &hard_split_policy) {
result.push(log_piece(
&piece,
(line_start + hard.byte_range.start)..(line_start + hard.byte_range.end),
));
}
chunk_start = line_end;
continue;
}
if line_end - chunk_start > MAX_FRAGMENT_BYTES {
result.push(log_piece(&piece, chunk_start..line_start));
chunk_start = line_start;
}
}
if chunk_start < piece.range.end {
result.push(log_piece(&piece, chunk_start..piece.range.end));
}
result
}
fn log_piece(source: &RawPiece, range: Range<usize>) -> RawPiece {
RawPiece {
kind: SegmentKind::Log,
range,
turn: source.turn,
role: source.role.clone(),
content_override: None,
}
}
fn merge_tiny(pieces: Vec<RawPiece>, min_bytes: usize) -> Vec<RawPiece> {
let mut merged: Vec<RawPiece> = Vec::new();
for piece in pieces {
let mergeable = merged
.last()
.is_some_and(|last| can_absorb(last, &piece, min_bytes));
match merged.last_mut() {
Some(last) if mergeable => last.range.end = piece.range.end,
_ => merged.push(piece),
}
}
merged
}
fn can_absorb(last: &RawPiece, next: &RawPiece, min_bytes: usize) -> bool {
last.turn == next.turn
&& last.kind == next.kind
&& last.content_override.is_none()
&& next.content_override.is_none()
&& last.range.end == next.range.start
&& last.range.len() + next.range.len() <= MAX_FRAGMENT_BYTES
&& (last.range.len() < min_bytes || next.range.len() < min_bytes)
}
fn build_segment(text: &str, piece: RawPiece, policy: &SegmentationPolicy) -> TranscriptSegment {
let content = piece
.content_override
.clone()
.unwrap_or_else(|| text[piece.range.clone()].to_owned());
let mut metadata = Map::new();
metadata.insert(
"role".to_owned(),
piece.role.clone().map_or(Value::Null, Value::String),
);
metadata.insert("turn".to_owned(), Value::Number(piece.turn.into()));
let is_first_turn_system = piece.turn == 0
&& piece
.role
.as_deref()
.is_some_and(|role| role.eq_ignore_ascii_case("system"));
if policy.cache_system_turn && is_first_turn_system {
metadata.insert("cache".to_owned(), Value::Bool(true));
}
let fragment = ContextFragment {
id: None,
content,
path: None,
kind: piece.kind.fragment_kind().map(str::to_owned),
priority: None,
metadata: Some(metadata),
media: None,
};
TranscriptSegment {
fragment,
turn: piece.turn,
role: piece.role,
kind: piece.kind,
byte_start: piece.range.start,
byte_end: piece.range.end,
}
}
#[cfg(test)]
#[path = "segment_tests.rs"]
mod tests;