poimen-memory/crates/mem-chunk/src/chunker.rs

use crate::record_source::RecordSource;
use crate::chunk_policy::ChunkPolicy;
use crate::token_counter::{TokenCounter, CharsOverFourCounter};
use mem_core::{Chunk, Record};
use futures::stream::Stream;

/// Create a stream of chunks from a record source.
pub fn chunks<S: RecordSource + 'static>(
    src: S,
    policy: ChunkPolicy,
) -> impl Stream<Item = Result<Chunk, String>> + Unpin {
    ChunkingAdapter {
        records: src.records(),
        policy,
        counter: CharsOverFourCounter,
        current_records: Vec::new(),
        current_tokens: 0,
        turn_index: 0,
    }
}

struct ChunkingAdapter {
    records: Box<dyn Stream<Item = Result<Record, String>> + Unpin>,
    policy: ChunkPolicy,
    counter: CharsOverFourCounter,
    current_records: Vec<Record>,
    current_tokens: usize,
    turn_index: u32,
}

impl Stream for ChunkingAdapter {
    type Item = Result<Chunk, String>;

    fn poll_next(
        mut self: std::pin::Pin<&mut Self>,
        cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Option<Self::Item>> {
        use std::pin::Pin;
        use std::task::Poll;

        loop {
            // Try to get the next record
            match Pin::new(&mut self.records).poll_next(cx) {
                Poll::Pending => {
                    // No record available right now
                    return Poll::Pending;
                }
                Poll::Ready(Some(Ok(record))) => {
                    let tokens = self.counter.count(&record);

                    // Check if adding this record would exceed the budget
                    if !self.current_records.is_empty()
                        && self.current_tokens + tokens > self.policy.max_tokens
                    {
                        // Flush the current chunk before adding this record
                        self.turn_index += 1;
                        let chunk = Chunk::new(
                            self.turn_index,
                            std::mem::take(&mut self.current_records),
                            self.current_tokens,
                        );
                        self.current_tokens = tokens;
                        self.current_records.push(record);
                        return Poll::Ready(Some(Ok(chunk)));
                    }

                    // Add record to current chunk
                    self.current_records.push(record);
                    self.current_tokens += tokens;
                    // Continue the loop to try getting the next record
                }
                Poll::Ready(Some(Err(e))) => {
                    return Poll::Ready(Some(Err(e)));
                }
                Poll::Ready(None) => {
                    // Stream exhausted
                    if !self.current_records.is_empty() {
                        self.turn_index += 1;
                        let chunk = Chunk::new(
                            self.turn_index,
                            std::mem::take(&mut self.current_records),
                            self.current_tokens,
                        );
                        self.current_tokens = 0;
                        return Poll::Ready(Some(Ok(chunk)));
                    }
                    return Poll::Ready(None);
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::record_source::VecSource;
    use mem_core::{Provenance, Role};
    use time::macros::datetime;
    use futures::stream::StreamExt;

    #[tokio::test]
    async fn test_basic_chunking() {
        let records = vec![
            Record {
                role: Role::User,
                text: "Hello world".to_string(),
                timestamp: datetime!(2024-08-20 12:00:00 UTC),
                provenance: Provenance {
                    source_id: "session1".to_string(),
                    offset: 0,
                },
            },
        ];

        let source = VecSource(records);
        let policy = ChunkPolicy::default();
        let mut chunk_stream = chunks(source, policy);

        let chunk = chunk_stream.next().await;
        assert!(chunk.is_some());
        let chunk = chunk.unwrap().unwrap();
        assert_eq!(chunk.t, 1);
        assert_eq!(chunk.records.len(), 1);
    }

    #[tokio::test]
    async fn test_empty_source() {
        let source = VecSource(vec![]);
        let policy = ChunkPolicy::default();
        let mut chunk_stream = chunks(source, policy);

        let result = chunk_stream.next().await;
        assert!(result.is_none());
    }

    #[tokio::test]
    async fn test_multiple_chunks() {
        let records = vec![
            Record {
                role: Role::User,
                text: "a".repeat(2000).to_string(),  // ~500 tokens
                timestamp: datetime!(2024-08-20 12:00:00 UTC),
                provenance: Provenance {
                    source_id: "s1".to_string(),
                    offset: 0,
                },
            },
            Record {
                role: Role::Assistant,
                text: "b".repeat(2000).to_string(),  // ~500 tokens
                timestamp: datetime!(2024-08-20 12:00:01 UTC),
                provenance: Provenance {
                    source_id: "s1".to_string(),
                    offset: 1,
                },
            },
            Record {
                role: Role::User,
                text: "c".repeat(2000).to_string(),  // ~500 tokens
                timestamp: datetime!(2024-08-20 12:00:02 UTC),
                provenance: Provenance {
                    source_id: "s1".to_string(),
                    offset: 2,
                },
            },
        ];

        let source = VecSource(records);
        let policy = ChunkPolicy {
            max_tokens: 800,
            split_on: crate::chunk_policy::Boundary::Record,
            flush: crate::chunk_policy::FlushTrigger::Tokens(800),
        };
        let mut chunk_stream = chunks(source, policy);

        // First chunk should have first two records (~1000 tokens, over budget)
        // Actually, since 500 + 500 = 1000 > 800, the second should cause a flush
        let chunk1 = chunk_stream.next().await.unwrap().unwrap();
        assert_eq!(chunk1.t, 1);
        assert_eq!(chunk1.records.len(), 1);

        let chunk2 = chunk_stream.next().await.unwrap().unwrap();
        assert_eq!(chunk2.t, 2);
    }
}