use std::collections::HashMap;
use synwire_core::BoxFuture;
use synwire_core::error::SynwireError;
pub trait ExampleSelector: Send + Sync {
fn select_examples<'a>(
&'a self,
input_variables: &'a HashMap<String, String>,
) -> BoxFuture<'a, Result<Vec<HashMap<String, String>>, SynwireError>>;
fn add_example(
&self,
example: HashMap<String, String>,
) -> BoxFuture<'_, Result<(), SynwireError>>;
}
pub struct SemanticSimilarityExampleSelector {
examples: std::sync::Mutex<Vec<HashMap<String, String>>>,
}
impl SemanticSimilarityExampleSelector {
pub const fn new() -> Self {
Self {
examples: std::sync::Mutex::new(Vec::new()),
}
}
}
impl Default for SemanticSimilarityExampleSelector {
fn default() -> Self {
Self::new()
}
}
fn lock_err(e: impl std::fmt::Display) -> SynwireError {
SynwireError::Other(Box::new(std::io::Error::other(e.to_string())))
}
impl ExampleSelector for SemanticSimilarityExampleSelector {
fn select_examples<'a>(
&'a self,
_input_variables: &'a HashMap<String, String>,
) -> BoxFuture<'a, Result<Vec<HashMap<String, String>>, SynwireError>> {
Box::pin(async move {
let guard = self.examples.lock().map_err(lock_err)?;
Ok(guard.clone())
})
}
fn add_example(
&self,
example: HashMap<String, String>,
) -> BoxFuture<'_, Result<(), SynwireError>> {
Box::pin(async move {
self.examples.lock().map_err(lock_err)?.push(example);
Ok(())
})
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
use super::*;
#[tokio::test]
async fn selector_returns_all_examples() {
let selector = SemanticSimilarityExampleSelector::new();
let mut ex1 = HashMap::new();
let _ = ex1.insert("input".into(), "hello".into());
let _ = ex1.insert("output".into(), "world".into());
selector.add_example(ex1).await.unwrap();
let mut ex2 = HashMap::new();
let _ = ex2.insert("input".into(), "foo".into());
let _ = ex2.insert("output".into(), "bar".into());
selector.add_example(ex2).await.unwrap();
let input = HashMap::new();
let examples = selector.select_examples(&input).await.unwrap();
assert_eq!(examples.len(), 2);
}
#[tokio::test]
async fn empty_selector_returns_empty() {
let selector = SemanticSimilarityExampleSelector::new();
let input = HashMap::new();
let examples = selector.select_examples(&input).await.unwrap();
assert!(examples.is_empty());
}
}