langchain_rust/prompt/
example_selector.rs

1//! Example selectors pick which examples a [`FewShotPromptTemplate`](super::FewShotPromptTemplate)
2//! shows for a given input.
3use std::{
4    cmp::Ordering,
5    fmt,
6    sync::{Arc, RwLock},
7};
8
9use async_trait::async_trait;
10
11use super::{PromptArgs, PromptError, PromptFromatter, PromptTemplate};
12use crate::{embedding::Embedder, schemas::messages::value_to_text};
13
14/// Chooses the examples to put in a few-shot prompt for a given input.
15///
16/// Examples are [`PromptArgs`] maps, formatted with the prompt's example template.
17#[async_trait]
18pub trait ExampleSelector: Send + Sync {
19    /// Adds an example to the pool the selector chooses from.
20    async fn add_example(&self, example: PromptArgs) -> Result<(), PromptError>;
21
22    /// The examples to use for `input_variables` (the prompt's inputs), in the order
23    /// they should appear.
24    async fn select_examples(
25        &self,
26        input_variables: &PromptArgs,
27    ) -> Result<Vec<PromptArgs>, PromptError>;
28}
29
30type LengthFn = Arc<dyn Fn(&str) -> usize + Send + Sync>;
31
32/// Selects examples in order until they would no longer fit in `max_length`, counting
33/// the length of the input too. Useful to keep a prompt within the context window.
34///
35/// Length is measured in words by default (see [`with_length_fn`](Self::with_length_fn)).
36///
37/// ```rust,ignore
38/// let selector = LengthBasedExampleSelector::new(
39///     template_fstring!("Input: {input}\nOutput: {output}", "input", "output"),
40///     vec![prompt_args! {"input" => "happy", "output" => "sad"}],
41/// )?
42/// .with_max_length(25);
43/// ```
44pub struct LengthBasedExampleSelector {
45    example_prompt: PromptTemplate,
46    max_length: usize,
47    length_fn: LengthFn,
48    /// Each example with the length of its formatted text.
49    examples: RwLock<Vec<(PromptArgs, usize)>>,
50}
51
52impl LengthBasedExampleSelector {
53    /// Default `max_length`, in words.
54    pub const DEFAULT_MAX_LENGTH: usize = 2048;
55
56    /// A selector over `examples`, formatted with `example_prompt` to measure them.
57    /// Fails if an example is missing a variable of `example_prompt`.
58    pub fn new(
59        example_prompt: PromptTemplate,
60        examples: Vec<PromptArgs>,
61    ) -> Result<Self, PromptError> {
62        let selector = Self {
63            example_prompt,
64            max_length: Self::DEFAULT_MAX_LENGTH,
65            length_fn: Arc::new(word_count),
66            examples: RwLock::new(Vec::new()),
67        };
68        for example in examples {
69            selector.push(example)?;
70        }
71        Ok(selector)
72    }
73
74    /// The most the selected examples plus the input may add up to.
75    pub fn with_max_length(mut self, max_length: usize) -> Self {
76        self.max_length = max_length;
77        self
78    }
79
80    /// How to measure text, e.g. a tokenizer's token count. Example lengths are
81    /// recomputed with it.
82    pub fn with_length_fn<F>(mut self, length_fn: F) -> Self
83    where
84        F: Fn(&str) -> usize + Send + Sync + 'static,
85    {
86        self.length_fn = Arc::new(length_fn);
87        let examples = std::mem::take(self.examples.get_mut().unwrap_or_else(|e| e.into_inner()));
88        for (example, _) in examples {
89            // Every example already formatted once, so this can't fail.
90            let _ = self.push(example);
91        }
92        self
93    }
94
95    fn push(&self, example: PromptArgs) -> Result<(), PromptError> {
96        let length = (self.length_fn)(&self.example_prompt.format(example.clone())?);
97        self.examples
98            .write()
99            .unwrap_or_else(|e| e.into_inner())
100            .push((example, length));
101        Ok(())
102    }
103}
104
105impl fmt::Debug for LengthBasedExampleSelector {
106    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
107        f.debug_struct("LengthBasedExampleSelector")
108            .field("max_length", &self.max_length)
109            .field(
110                "examples",
111                &self
112                    .examples
113                    .read()
114                    .unwrap_or_else(|e| e.into_inner())
115                    .len(),
116            )
117            .finish()
118    }
119}
120
121#[async_trait]
122impl ExampleSelector for LengthBasedExampleSelector {
123    async fn add_example(&self, example: PromptArgs) -> Result<(), PromptError> {
124        self.push(example)
125    }
126
127    async fn select_examples(
128        &self,
129        input_variables: &PromptArgs,
130    ) -> Result<Vec<PromptArgs>, PromptError> {
131        let input_length: usize = input_variables
132            .values()
133            .map(|v| (self.length_fn)(&value_to_text(v)))
134            .sum();
135        let mut remaining = self.max_length.saturating_sub(input_length);
136        let mut selected = Vec::new();
137        for (example, length) in self
138            .examples
139            .read()
140            .unwrap_or_else(|e| e.into_inner())
141            .iter()
142        {
143            if *length > remaining {
144                break;
145            }
146            remaining -= length;
147            selected.push(example.clone());
148        }
149        Ok(selected)
150    }
151}
152
153/// Words in `text`, the default length for [`LengthBasedExampleSelector`].
154fn word_count(text: &str) -> usize {
155    text.split_whitespace().count()
156}
157
158/// Selects the `k` examples most similar to the input, by embedding cosine similarity,
159/// most similar first.
160///
161/// An example's text is its values joined by spaces in key order. With
162/// [`with_input_keys`](Self::with_input_keys) only those keys are used, both for the
163/// examples and for the input (so the expected output doesn't count).
164///
165/// ```rust,ignore
166/// let selector = SemanticSimilarityExampleSelector::from_examples(
167///     examples,
168///     OllamaEmbedder::default(),
169///     2,
170/// )
171/// .await?
172/// .with_input_keys(["input"]);
173/// ```
174pub struct SemanticSimilarityExampleSelector {
175    embedder: Arc<dyn Embedder>,
176    k: usize,
177    input_keys: Option<Vec<String>>,
178    examples: RwLock<Vec<(PromptArgs, Vec<f64>)>>,
179}
180
181impl SemanticSimilarityExampleSelector {
182    /// Default number of examples to select.
183    pub const DEFAULT_K: usize = 4;
184
185    /// An empty selector; add examples with [`ExampleSelector::add_example`].
186    pub fn new<E: Embedder + 'static>(embedder: E) -> Self {
187        Self::with_embedder(Arc::new(embedder))
188    }
189
190    /// Like [`new`](Self::new), with an embedder shared with other components.
191    pub fn with_embedder(embedder: Arc<dyn Embedder>) -> Self {
192        Self {
193            embedder,
194            k: Self::DEFAULT_K,
195            input_keys: None,
196            examples: RwLock::new(Vec::new()),
197        }
198    }
199
200    /// A selector over `examples` (embedded in one batch) that picks `k` of them.
201    pub async fn from_examples<E: Embedder + 'static>(
202        examples: Vec<PromptArgs>,
203        embedder: E,
204        k: usize,
205    ) -> Result<Self, PromptError> {
206        let selector = Self::new(embedder).with_k(k);
207        selector.add_examples(examples).await?;
208        Ok(selector)
209    }
210
211    /// How many examples to select.
212    pub fn with_k(mut self, k: usize) -> Self {
213        self.k = k;
214        self
215    }
216
217    /// Only these keys are embedded, for examples and input alike. Examples are
218    /// re-embedded when added, so set this before adding them.
219    pub fn with_input_keys<I, S>(mut self, keys: I) -> Self
220    where
221        I: IntoIterator<Item = S>,
222        S: Into<String>,
223    {
224        self.input_keys = Some(keys.into_iter().map(Into::into).collect());
225        self
226    }
227
228    /// Adds several examples, embedded in one batch.
229    pub async fn add_examples(&self, examples: Vec<PromptArgs>) -> Result<(), PromptError> {
230        if examples.is_empty() {
231            return Ok(());
232        }
233        let texts: Vec<String> = examples.iter().map(|e| self.text_of(e)).collect();
234        let vectors = self
235            .embedder
236            .embed_documents(&texts)
237            .await
238            .map_err(|e| PromptError::OtherError(format!("embedding examples: {e}")))?;
239        if vectors.len() != examples.len() {
240            return Err(PromptError::OtherError(format!(
241                "embedder returned {} vectors for {} examples",
242                vectors.len(),
243                examples.len()
244            )));
245        }
246        self.examples
247            .write()
248            .unwrap_or_else(|e| e.into_inner())
249            .extend(examples.into_iter().zip(vectors));
250        Ok(())
251    }
252
253    #[cfg(test)]
254    pub(crate) fn examples_for_test(&self, examples: Vec<(PromptArgs, Vec<f64>)>) {
255        self.examples
256            .write()
257            .unwrap_or_else(|e| e.into_inner())
258            .extend(examples);
259    }
260
261    /// The text embedded for `args`: the values of the input keys (or all keys),
262    /// in key order, joined by spaces.
263    fn text_of(&self, args: &PromptArgs) -> String {
264        let mut keys: Vec<&String> = match &self.input_keys {
265            Some(keys) => keys.iter().filter(|k| args.contains_key(*k)).collect(),
266            None => args.keys().collect(),
267        };
268        keys.sort();
269        keys.iter()
270            .map(|k| value_to_text(&args[*k]))
271            .collect::<Vec<_>>()
272            .join(" ")
273    }
274}
275
276impl fmt::Debug for SemanticSimilarityExampleSelector {
277    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
278        f.debug_struct("SemanticSimilarityExampleSelector")
279            .field("k", &self.k)
280            .field("input_keys", &self.input_keys)
281            .field(
282                "examples",
283                &self
284                    .examples
285                    .read()
286                    .unwrap_or_else(|e| e.into_inner())
287                    .len(),
288            )
289            .finish()
290    }
291}
292
293#[async_trait]
294impl ExampleSelector for SemanticSimilarityExampleSelector {
295    async fn add_example(&self, example: PromptArgs) -> Result<(), PromptError> {
296        self.add_examples(vec![example]).await
297    }
298
299    async fn select_examples(
300        &self,
301        input_variables: &PromptArgs,
302    ) -> Result<Vec<PromptArgs>, PromptError> {
303        if self.k == 0
304            || self
305                .examples
306                .read()
307                .unwrap_or_else(|e| e.into_inner())
308                .is_empty()
309        {
310            return Ok(Vec::new());
311        }
312        let query = self
313            .embedder
314            .embed_query(&self.text_of(input_variables))
315            .await
316            .map_err(|e| PromptError::OtherError(format!("embedding the input: {e}")))?;
317        let examples = self.examples.read().unwrap_or_else(|e| e.into_inner());
318        let mut scored: Vec<(f64, &PromptArgs)> = examples
319            .iter()
320            .map(|(example, vector)| (cosine_similarity(&query, vector), example))
321            .collect();
322        // Stable sort: ties keep insertion order.
323        scored.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap_or(Ordering::Equal));
324        Ok(scored
325            .into_iter()
326            .take(self.k)
327            .map(|(_, example)| example.clone())
328            .collect())
329    }
330}
331
332/// Cosine similarity; 0 when either vector is all zeros.
333fn cosine_similarity(a: &[f64], b: &[f64]) -> f64 {
334    let dot: f64 = a.iter().zip(b).map(|(x, y)| x * y).sum();
335    let norm = |v: &[f64]| v.iter().map(|x| x * x).sum::<f64>().sqrt();
336    let denominator = norm(a) * norm(b);
337    if denominator == 0.0 {
338        0.0
339    } else {
340        dot / denominator
341    }
342}
langchain_rust/prompt/example_selector.rs

langchain_rust/prompt/
example_selector.rs
