ic_rig/embeddings/
embed.rs

1//! [`Embed`] trait and [`TextEmbedder`] accumulator.
2//!
3//! `Embed` decouples your domain types from the embedding API. Rather than
4//! returning a `Vec<String>` (which forces allocation at the call site),
5//! your type **pushes** strings into a [`TextEmbedder`]. This lets the
6//! [`EmbeddingsBuilder`](super::EmbeddingsBuilder) collect texts from many
7//! documents into a single flat buffer before chunking into API requests.
8//!
9//! A single document can push **multiple** strings — you get one
10//! [`Embedding`](super::Embedding) back per push. This is how one struct
11//! can produce several vectors (e.g. embed `title` and `body` separately).
12//!
13//! # Example
14//!
15//! ```rust
16//! use irig::embeddings::{Embed, TextEmbedder, EmbedError};
17//!
18//! struct Article { title: String, body: String }
19//!
20//! impl Embed for Article {
21//!     fn embed(&self, e: &mut TextEmbedder) -> Result<(), EmbedError> {
22//!         e.embed(self.title.clone());
23//!         e.embed(self.body.clone());
24//!         Ok(())
25//!     }
26//! }
27//! ```
28
29use thiserror::Error;
30
31// ── EmbedError ────────────────────────────────────────────────────────────────
32
33#[derive(Debug, Error)]
34#[error("{0}")]
35pub struct EmbedError(String);
36
37impl EmbedError {
38    pub fn new<E: std::error::Error>(e: E) -> Self {
39        Self(e.to_string())
40    }
41
42    pub fn msg(s: impl Into<String>) -> Self {
43        Self(s.into())
44    }
45}
46
47// ── TextEmbedder ──────────────────────────────────────────────────────────────
48
49/// Accumulates strings to be embedded. Passed into [`Embed::embed`].
50#[derive(Default)]
51pub struct TextEmbedder {
52    pub(crate) texts: Vec<String>,
53}
54
55impl TextEmbedder {
56    /// Queue `text` for embedding. Each call produces one output [`Embedding`](super::Embedding).
57    pub fn embed(&mut self, text: String) {
58        self.texts.push(text);
59    }
60}
61
62// ── Embed trait ───────────────────────────────────────────────────────────────
63
64/// Implemented by types that can be converted to text for embedding.
65///
66/// Use the [`derive`](irig_derive) macro or implement manually.
67pub trait Embed {
68    fn embed(&self, embedder: &mut TextEmbedder) -> Result<(), EmbedError>;
69}
70
71/// Extract all texts an object would embed without running the model.
72pub fn to_texts(item: impl Embed) -> Result<Vec<String>, EmbedError> {
73    let mut embedder = TextEmbedder::default();
74    item.embed(&mut embedder)?;
75    Ok(embedder.texts)
76}
77
78// ── Blanket impls for primitives ──────────────────────────────────────────────
79
80impl Embed for String {
81    fn embed(&self, e: &mut TextEmbedder) -> Result<(), EmbedError> {
82        e.embed(self.clone());
83        Ok(())
84    }
85}
86
87impl Embed for &str {
88    fn embed(&self, e: &mut TextEmbedder) -> Result<(), EmbedError> {
89        e.embed(self.to_string());
90        Ok(())
91    }
92}
93
94impl Embed for serde_json::Value {
95    fn embed(&self, e: &mut TextEmbedder) -> Result<(), EmbedError> {
96        e.embed(serde_json::to_string(self).map_err(EmbedError::new)?);
97        Ok(())
98    }
99}
100
101impl<T: Embed> Embed for &T {
102    fn embed(&self, e: &mut TextEmbedder) -> Result<(), EmbedError> {
103        (*self).embed(e)
104    }
105}
106
107impl<T: Embed> Embed for Vec<T> {
108    fn embed(&self, e: &mut TextEmbedder) -> Result<(), EmbedError> {
109        for item in self {
110            item.embed(e)?;
111        }
112        Ok(())
113    }
114}
115
116impl<T: Embed> Embed for Option<T> {
117    fn embed(&self, e: &mut TextEmbedder) -> Result<(), EmbedError> {
118        if let Some(inner) = self {
119            inner.embed(e)?;
120        }
121        Ok(())
122    }
123}
ic_rig/embeddings/embed.rs

ic_rig/embeddings/
embed.rs
