feat(ingestion_pipeline): default concurrency is the number of cpus (#6)

bosun-ai · Jun 13, 2024 · f10bc30 · f10bc30
1 parent 1036d56
commit f10bc30
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 2 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/swiftide/Cargo.toml b/swiftide/Cargo.toml
@@ -25,6 +25,7 @@ tokio = { version = "1.38.0", features = ["full"] }
 tracing = { version = "0.1.40", features = ["log"] }
 strum = "0.26.2"
 strum_macros = "0.26.4"
+num_cpus = "1.16.0"
 
 # Integrations
 async-openai = { version = "0.23.2", optional = true }

diff --git a/swiftide/src/ingestion/ingestion_pipeline.rs b/swiftide/src/ingestion/ingestion_pipeline.rs
@@ -12,14 +12,23 @@ pub struct IngestionPipeline {
     concurrency: usize,
 }
 
+impl Default for IngestionPipeline {
+    fn default() -> Self {
+        Self {
+            stream: Box::pin(futures_util::stream::empty()),
+            storage: None,
+            concurrency: num_cpus::get(),
+        }
+    }
+}
+
 // A lazy pipeline for ingesting files, adding metadata, chunking, transforming, embedding and then storing them.
 impl IngestionPipeline {
     pub fn from_loader(loader: impl Loader + 'static) -> Self {
         let stream = loader.into_stream();
         Self {
             stream: stream.boxed(),
-            storage: None,
-            concurrency: 10,
+            ..Default::default()
         }
     }