rig_core/embeddings/
embed.rs

1//! Text extraction for embedding builders, including implementations for common types.
2//!
3//! ```
4//! use rig_core::embeddings::to_texts;
5//!
6//! assert_eq!(to_texts("document")?, vec!["document"]);
7//! # Ok::<(), rig_core::embeddings::EmbedError>(())
8//! ```
9
10/// Failure extracting text for embedding.
11#[derive(Debug, thiserror::Error)]
12#[error("{0}")]
13pub struct EmbedError(#[from] Box<dyn std::error::Error + Send + Sync>);
14
15impl EmbedError {
16    pub fn new<E: std::error::Error + Send + Sync + 'static>(error: E) -> Self {
17        EmbedError(Box::new(error))
18    }
19}
20
21/// Extracts text fragments for vector embedding in append order.
22/// Implementations report extraction failures through [`EmbedError`].
23///
24/// ```
25/// use rig_core::{Embed, embeddings::{EmbedError, TextEmbedder, to_texts}};
26///
27/// struct Definitions(String);
28/// impl Embed for Definitions {
29///     fn embed(&self, out: &mut TextEmbedder) -> Result<(), EmbedError> {
30///         for definition in self.0.split(',') {
31///             out.embed(definition.trim().to_owned());
32///         }
33///         Ok(())
34///     }
35/// }
36/// assert_eq!(to_texts(Definitions("fruit, company".into()))?, vec!["fruit", "company"]);
37/// # Ok::<(), EmbedError>(())
38/// ```
39pub trait Embed {
40    /// Append all text fragments that should be embedded for this value.
41    fn embed(&self, embedder: &mut TextEmbedder) -> Result<(), EmbedError>;
42}
43
44/// Accumulates string values that need to be embedded.
45/// Used by the [Embed] trait.
46#[derive(Default)]
47pub struct TextEmbedder {
48    pub(crate) texts: Vec<String>,
49}
50
51impl TextEmbedder {
52    /// Appends a text fragment without filtering or deduplication.
53    pub fn embed(&mut self, text: String) {
54        self.texts.push(text);
55    }
56}
57
58/// Extracts text fragments in append order, propagating extraction errors.
59pub fn to_texts(item: impl Embed) -> Result<Vec<String>, EmbedError> {
60    let mut embedder = TextEmbedder::default();
61    item.embed(&mut embedder)?;
62    Ok(embedder.texts)
63}
64
65impl Embed for String {
66    fn embed(&self, embedder: &mut TextEmbedder) -> Result<(), EmbedError> {
67        embedder.embed(self.clone());
68        Ok(())
69    }
70}
71
72/// Embeds each type as its `Display` text.
73macro_rules! embed_as_text {
74    ($($type:ty),*) => {$(
75        impl Embed for $type {
76            fn embed(&self, embedder: &mut TextEmbedder) -> Result<(), EmbedError> {
77                embedder.embed(self.to_string());
78                Ok(())
79            }
80        }
81    )*};
82}
83
84embed_as_text!(&str, i8, i16, i32, i64, i128, f32, f64, bool, char);
85
86impl Embed for serde_json::Value {
87    fn embed(&self, embedder: &mut TextEmbedder) -> Result<(), EmbedError> {
88        embedder.embed(serde_json::to_string(self).map_err(EmbedError::new)?);
89        Ok(())
90    }
91}
92
93impl<T: Embed> Embed for &T {
94    fn embed(&self, embedder: &mut TextEmbedder) -> Result<(), EmbedError> {
95        (*self).embed(embedder)
96    }
97}
98
99impl<T: Embed> Embed for Vec<T> {
100    fn embed(&self, embedder: &mut TextEmbedder) -> Result<(), EmbedError> {
101        for item in self {
102            item.embed(embedder).map_err(EmbedError::new)?;
103        }
104        Ok(())
105    }
106}
rig_core/embeddings/embed.rs

rig_core/embeddings/
embed.rs
