diff --git a/Cargo.lock b/Cargo.lock index b2a26b24..974dd666 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -587,6 +587,7 @@ dependencies = [ "ciborium", "clap", "criterion-plot", + "futures", "is-terminal", "itertools 0.10.5", "num-traits", @@ -599,6 +600,7 @@ dependencies = [ "serde_derive", "serde_json", "tinytemplate", + "tokio", "walkdir", ] diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index ba2d47b9..a1f2e0f1 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -9,7 +9,7 @@ tokio = { version = "1.0", features = ["full"] } swiftide = { path = "../swiftide/", features = ["all"] } tracing-subscriber = "0.3" serde_json = "1.0" -criterion = { version = "0.5.1", features = ["html_reports"] } +criterion = { version = "0.5.1", features = ["html_reports", "async_tokio"] } anyhow = "1.0" futures-util = "0.3" @@ -17,3 +17,8 @@ futures-util = "0.3" name = "fileloader" path = "fileloader.rs" harness = false + +[[bench]] +name = "local_pipeline" +path = "local_pipeline.rs" +harness = false diff --git a/benchmarks/local_pipeline.rs b/benchmarks/local_pipeline.rs new file mode 100644 index 00000000..6090fcc6 --- /dev/null +++ b/benchmarks/local_pipeline.rs @@ -0,0 +1,30 @@ +use anyhow::Result; +use criterion::{criterion_group, criterion_main, Criterion}; +use swiftide::{ + ingestion::IngestionPipeline, + integrations::{fastembed::FastEmbed, qdrant::Qdrant}, + loaders::FileLoader, + transformers::Embed, +}; + +async fn run_pipeline() -> Result<()> { + let qdrant_url = "http://localhost:6334"; + IngestionPipeline::from_loader(FileLoader::new(".").with_extensions(&["rs"])) + .then_in_batch(10, Embed::new(FastEmbed::builder().batch_size(10).build()?)) + .then_store_with( + Qdrant::try_from_url(qdrant_url)? + .batch_size(50) + .vector_size(384) + .collection_name("swiftide-examples-fastembed".to_string()) + .build()?, + ) + .run() + .await +} + +fn criterion_benchmark(c: &mut Criterion) { + c.bench_function("run_local_pipeline", |b| b.iter(run_pipeline)); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches);