dspy_rs/core/
module.rs

1use anyhow::Result;
2use futures::stream::{self, StreamExt};
3use indexmap::IndexMap;
4use kdam::{BarExt, tqdm};
5
6use crate::{Example, Prediction, core::MetaSignature};
7
8#[allow(async_fn_in_trait)]
9pub trait Module: Send + Sync {
10    async fn forward(&self, inputs: Example) -> Result<Prediction>;
11
12    async fn batch(
13        &self,
14        inputs: Vec<Example>,
15        max_concurrency: usize,
16        display_progress: bool,
17    ) -> Result<Vec<Prediction>> {
18        let total = inputs.len();
19        let mut pb = if display_progress {
20            Some(tqdm!(total = total, desc = "Processing"))
21        } else {
22            None
23        };
24
25        // Pair each input with its index to maintain order
26        let indexed_results: Vec<(usize, Result<Prediction>)> =
27            stream::iter(inputs.into_iter().enumerate())
28                .map(|(idx, example)| async move {
29                    let result = self.forward(example).await;
30                    (idx, result)
31                })
32                .buffer_unordered(max_concurrency)
33                .inspect(|_| {
34                    if let Some(ref mut progress) = pb {
35                        let _ = progress.update(1);
36                    }
37                })
38                .collect()
39                .await;
40
41        // Sort results back to original order
42        let mut indexed_results = indexed_results;
43        indexed_results.sort_by_key(|(idx, _)| *idx);
44
45        // Collect predictions and handle errors
46        let mut predictions = Vec::with_capacity(total);
47        for (_, result) in indexed_results {
48            predictions.push(result?);
49        }
50
51        Ok(predictions)
52    }
53}
54
55#[allow(unused_variables)]
56pub trait Optimizable {
57    fn get_signature(&self) -> &dyn MetaSignature {
58        todo!()
59    }
60
61    fn parameters(&mut self) -> IndexMap<String, &mut dyn Optimizable>;
62
63    fn update_signature_instruction(&mut self, instruction: String) -> anyhow::Result<()> {
64        todo!()
65    }
66}