Skip to main content

entelix_prompt/
example_selector.rs

1//! Example-selection strategies for [`crate::FewShotPromptTemplate`].
2//!
3//! Selectors convert the caller's input variables into a list of examples
4//! (each example is a `HashMap<String, String>` mapping example-prompt
5//! variable names to their values). Two concrete impls ship in 1.0:
6//!
7//! - [`FixedExampleSelector`] — returns the same static example list every
8//!   call.
9//! - [`LengthBasedExampleSelector`] — drops trailing examples until the
10//!   rendered length fits a configurable character cap.
11//!
12//! A semantic selector (embedding-similarity) is deferred to 1.1 alongside
13//! concrete `Embedder` impls.
14
15use std::collections::HashMap;
16use std::sync::Arc;
17
18use entelix_core::Result;
19
20use crate::template::PromptTemplate;
21
22/// One few-shot example: a map of example-prompt variable name → value.
23pub type Example = HashMap<String, String>;
24
25/// Strategy that picks which examples to inject into a few-shot prompt.
26///
27/// The trait is async-free; selection is expected to be a fast in-process
28/// computation (filter, length cap, embedding lookup once vector stores
29/// arrive). Implementors that need IO can return a precomputed list.
30pub trait ExampleSelector: Send + Sync {
31    /// Select examples for an invocation. `input_vars` are the suffix
32    /// variables the caller will pass at render time — selectors may use
33    /// them to score relevance; `FixedExampleSelector` ignores them.
34    fn select(&self, input_vars: &HashMap<String, String>) -> Result<Vec<Example>>;
35}
36
37/// Returns the configured examples verbatim, regardless of input.
38#[derive(Clone, Debug)]
39pub struct FixedExampleSelector {
40    examples: Vec<Example>,
41}
42
43impl FixedExampleSelector {
44    /// Build a selector from a static example list.
45    pub const fn new(examples: Vec<Example>) -> Self {
46        Self { examples }
47    }
48
49    /// Borrow the underlying example list.
50    pub fn examples(&self) -> &[Example] {
51        &self.examples
52    }
53}
54
55impl ExampleSelector for FixedExampleSelector {
56    fn select(&self, _input_vars: &HashMap<String, String>) -> Result<Vec<Example>> {
57        Ok(self.examples.clone())
58    }
59}
60
61/// Selector that caps the rendered example total at a character budget.
62///
63/// Drops trailing examples until the rendered character length fits.
64/// The cap measures raw UTF-8 character count over the concatenated
65/// rendered examples joined by `\n` — coarse but deterministic and
66/// dependency-free.
67#[derive(Clone, Debug)]
68pub struct LengthBasedExampleSelector {
69    examples: Vec<Example>,
70    example_prompt: PromptTemplate,
71    max_chars: usize,
72}
73
74impl LengthBasedExampleSelector {
75    /// Build a length-based selector. `example_prompt` is the same
76    /// template used by the parent `FewShotPromptTemplate`; rendered
77    /// length is what counts.
78    pub const fn new(
79        examples: Vec<Example>,
80        example_prompt: PromptTemplate,
81        max_chars: usize,
82    ) -> Self {
83        Self {
84            examples,
85            example_prompt,
86            max_chars,
87        }
88    }
89
90    /// Borrow the example pool.
91    pub fn examples(&self) -> &[Example] {
92        &self.examples
93    }
94}
95
96impl ExampleSelector for LengthBasedExampleSelector {
97    fn select(&self, _input_vars: &HashMap<String, String>) -> Result<Vec<Example>> {
98        let mut total = 0usize;
99        let mut kept: Vec<Example> = Vec::new();
100        for example in &self.examples {
101            let rendered = self.example_prompt.render(example)?;
102            let len = rendered.chars().count();
103            if total.saturating_add(len) > self.max_chars {
104                break;
105            }
106            total = total.saturating_add(len);
107            kept.push(example.clone());
108        }
109        Ok(kept)
110    }
111}
112
113/// Convenience type-alias for selectors handed around as trait objects.
114pub type SharedExampleSelector = Arc<dyn ExampleSelector>;