74 lines
2.6 KiB
Rust
74 lines
2.6 KiB
Rust
use mem_ingest::PiSessionSource;
|
|
use mem_chunk::{RecordSource, chunks, ChunkPolicy};
|
|
use futures::stream::StreamExt;
|
|
use std::path::PathBuf;
|
|
|
|
#[tokio::main]
|
|
async fn main() -> anyhow::Result<()> {
|
|
let fixture = PathBuf::from("fixtures/pi-session-small.jsonl");
|
|
println!("Testing E2E pipeline with: {}", fixture.display());
|
|
|
|
// Step 1: Create source and extract project
|
|
println!("\n1. Creating pi session source...");
|
|
let source = PiSessionSource::new(fixture.clone());
|
|
let project = source.read_project_key().await?;
|
|
println!(" Project: {}", project);
|
|
|
|
// Step 2: Stream records
|
|
println!("\n2. Streaming records...");
|
|
let source = PiSessionSource::new(fixture.clone());
|
|
let mut records_stream = source.records();
|
|
let mut record_count = 0;
|
|
while let Some(result) = records_stream.next().await {
|
|
match result {
|
|
Ok(record) => {
|
|
record_count += 1;
|
|
println!(" Record {}: role={{:?}}, text_len={}", record_count, record.text.len());
|
|
}
|
|
Err(e) => eprintln!(" Error: {}", e),
|
|
}
|
|
}
|
|
println!(" Total records: {}", record_count);
|
|
|
|
// Step 3: Chunk them
|
|
println!("\n3. Chunking with 5000 token budget...");
|
|
let source = PiSessionSource::new(fixture);
|
|
let policy = ChunkPolicy::default();
|
|
let mut chunk_stream = chunks(source, policy);
|
|
let mut chunk_count = 0;
|
|
let mut total_records_in_chunks = 0;
|
|
while let Some(result) = chunk_stream.next().await {
|
|
match result {
|
|
Ok(chunk) => {
|
|
chunk_count += 1;
|
|
total_records_in_chunks += chunk.records.len();
|
|
println!(" Chunk {}: t={}, records={}, tokens={}",
|
|
chunk_count, chunk.t, chunk.records.len(), chunk.tokens);
|
|
}
|
|
Err(e) => eprintln!(" Error: {}", e),
|
|
}
|
|
}
|
|
|
|
println!("\n=== E2E PIPELINE RESULTS ===");
|
|
println!("Records streamed: {}", record_count);
|
|
println!("Chunks produced: {}", chunk_count);
|
|
println!("Records in chunks: {}", total_records_in_chunks);
|
|
println!("Lossless: {}", record_count == total_records_in_chunks);
|
|
|
|
if record_count == 0 {
|
|
eprintln!("\n❌ FAILED: No records parsed!");
|
|
std::process::exit(1);
|
|
}
|
|
if chunk_count == 0 {
|
|
eprintln!("\n❌ FAILED: No chunks produced!");
|
|
std::process::exit(1);
|
|
}
|
|
if record_count != total_records_in_chunks {
|
|
eprintln!("\n❌ FAILED: Records lost in chunking!");
|
|
std::process::exit(1);
|
|
}
|
|
|
|
println!("\n✅ SUCCESS: Full pipeline works!");
|
|
Ok(())
|
|
}
|