50 lines
1.2 KiB
Rust
50 lines
1.2 KiB
Rust
|
/// Boundary mode - where chunks can be split.
|
||
|
|
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
|
||
|
|
pub enum Boundary {
|
||
|
|
/// Never split inside a Record
|
||
|
|
Record,
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Trigger for flushing a chunk.
|
||
|
|
#[derive(Clone, Debug)]
|
||
|
|
pub enum FlushTrigger {
|
||
|
|
/// Flush when this many tokens is reached
|
||
|
|
Tokens(usize),
|
||
|
|
// OrIdle(Duration) will land with the first streaming source.
|
||
|
|
// Carrying the enum now means that change is one variant, not a signature change
|
||
|
|
// threaded through the loop.
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Chunking policy.
|
||
|
|
#[derive(Clone, Debug)]
|
||
|
|
pub struct ChunkPolicy {
|
||
|
|
/// Maximum tokens per chunk (default 5000 - GRU-Mem paper default)
|
||
|
|
pub max_tokens: usize,
|
||
|
|
/// Boundary mode - never split inside a Record
|
||
|
|
pub split_on: Boundary,
|
||
|
|
/// Flush trigger
|
||
|
|
pub flush: FlushTrigger,
|
||
|
|
}
|
||
|
|
|
||
|
|
impl Default for ChunkPolicy {
|
||
|
|
fn default() -> Self {
|
||
|
|
ChunkPolicy {
|
||
|
|
max_tokens: 5000,
|
||
|
|
split_on: Boundary::Record,
|
||
|
|
flush: FlushTrigger::Tokens(5000),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
#[cfg(test)]
|
||
|
|
mod tests {
|
||
|
|
use super::*;
|
||
|
|
|
||
|
|
#[test]
|
||
|
|
fn test_default_chunk_policy() {
|
||
|
|
let policy = ChunkPolicy::default();
|
||
|
|
assert_eq!(policy.max_tokens, 5000);
|
||
|
|
assert_eq!(policy.split_on, Boundary::Record);
|
||
|
|
}
|
||
|
|
}
|