feat: Add Op::batch_call and TryOp::try_batch_call

rig-core/examples/multi_extract.rs

@@ -50,6 +50,9 @@ async fn main() -> anyhow::Result<()> {
         )
         .build();
 
+    // Create a chain that extracts names, topics, and sentiment from a given text
+    // using three different GPT-4 based extractors.
+    // The chain will output a formatted string containing the extracted information.
     let chain = pipeline::new()
         .chain(try_parallel!(
             agent_ops::extract(names_extractor),
@@ -65,9 +68,21 @@ async fn main() -> anyhow::Result<()> {
             )
         });
 
-    let response = chain.try_call("Screw you Putin!").await?;
+    // Batch call the chain with up to 4 inputs concurrently
+    let response = chain
+        .try_batch_call(
+            4,
+            vec![
+                "Screw you Putin!",
+                "I love my dog, but I hate my cat.",
+                "I'm going to the store to buy some milk.",
+            ],
+        )
+        .await?;
 
-    println!("Text analysis:\n{response}");
+    for response in response {
+        println!("Text analysis:\n{response}");
+    }
 
     Ok(())
 }

rig-core/src/pipeline/mod.rs

@@ -233,10 +233,7 @@ mod tests {
             .await
             .expect("Failed to run chain");
 
-        assert_eq!(
-            result,
-            "Top documents:\nbar"
-        );
+        assert_eq!(result, "Top documents:\nbar");
     }
 
     #[tokio::test]

rig-core/src/pipeline/op.rs

@@ -2,6 +2,7 @@ use std::future::Future;
 
 #[allow(unused_imports)] // Needed since this is used in a macro rule
 use futures::join;
+use futures::{stream, StreamExt};
 
 // ================================================================
 // Core Op trait
@@ -12,6 +13,23 @@ pub trait Op: Send + Sync {
 
     fn call(&self, input: Self::Input) -> impl Future<Output = Self::Output> + Send;
 
+    /// Execute the current pipeline with the given inputs. `n` is the number of concurrent
+    /// inputs that will be processed concurrently.
+    fn batch_call<I>(&self, n: usize, input: I) -> impl Future<Output = Vec<Self::Output>> + Send
+    where
+        I: IntoIterator<Item = Self::Input> + Send,
+        I::IntoIter: Send,
+        Self: Sized,
+    {
+        async move {
+            stream::iter(input)
+                .map(|input| self.call(input))
+                .buffered(n)
+                .collect()
+                .await
+        }
+    }
+
     /// Chain a function to the current pipeline
     ///
     /// # Example

rig-core/src/pipeline/try_op.rs

@@ -2,6 +2,7 @@ use std::future::Future;
 
 #[allow(unused_imports)] // Needed since this is used in a macro rule
 use futures::try_join;
+use futures::{stream, StreamExt, TryStreamExt};
 
 use super::op::{self, map, then};
 
@@ -18,6 +19,29 @@ pub trait TryOp: Send + Sync {
         input: Self::Input,
     ) -> impl Future<Output = Result<Self::Output, Self::Error>> + Send;
 
+    /// Execute the current pipeline with the given inputs. `n` is the number of concurrent
+    /// inputs that will be processed concurrently.
+    /// If one of the inputs fails, the entire operation will fail and the error will
+    /// be returned.
+    fn try_batch_call<I>(
+        &self,
+        n: usize,
+        input: I,
+    ) -> impl Future<Output = Result<Vec<Self::Output>, Self::Error>> + Send
+    where
+        I: IntoIterator<Item = Self::Input> + Send,
+        I::IntoIter: Send,
+        Self: Sized,
+    {
+        async move {
+            stream::iter(input)
+                .map(|input| self.try_call(input))
+                .buffered(n)
+                .try_collect()
+                .await
+        }
+    }
+
     fn map_ok<F, T>(self, f: F) -> impl op::Op<Input = Self::Input, Output = Result<T, Self::Error>>
     where
         F: Fn(Self::Output) -> T + Send + Sync,