ic-rig 0.2.0

A lean, modular library for building LLM applications. Bring your own HTTP client.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

//! [`Embed`] trait and [`TextEmbedder`] accumulator.
//!
//! `Embed` decouples your domain types from the embedding API. Rather than
//! returning a `Vec<String>` (which forces allocation at the call site),
//! your type **pushes** strings into a [`TextEmbedder`]. This lets the
//! [`EmbeddingsBuilder`](super::EmbeddingsBuilder) collect texts from many
//! documents into a single flat buffer before chunking into API requests.
//!
//! A single document can push **multiple** strings — you get one
//! [`Embedding`](super::Embedding) back per push. This is how one struct
//! can produce several vectors (e.g. embed `title` and `body` separately).
//!
//! # Example
//!
//! ```rust
//! use irig::embeddings::{Embed, TextEmbedder, EmbedError};
//!
//! struct Article { title: String, body: String }
//!
//! impl Embed for Article {
//!     fn embed(&self, e: &mut TextEmbedder) -> Result<(), EmbedError> {
//!         e.embed(self.title.clone());
//!         e.embed(self.body.clone());
//!         Ok(())
//!     }
//! }
//! ```

use thiserror::Error;

// ── EmbedError ────────────────────────────────────────────────────────────────

#[derive(Debug, Error)]
#[error("{0}")]
pub struct EmbedError(String);

impl EmbedError {
    pub fn new<E: std::error::Error>(e: E) -> Self {
        Self(e.to_string())
    }

    pub fn msg(s: impl Into<String>) -> Self {
        Self(s.into())
    }
}

// ── TextEmbedder ──────────────────────────────────────────────────────────────

/// Accumulates strings to be embedded. Passed into [`Embed::embed`].
#[derive(Default)]
pub struct TextEmbedder {
    pub(crate) texts: Vec<String>,
}

impl TextEmbedder {
    /// Queue `text` for embedding. Each call produces one output [`Embedding`](super::Embedding).
    pub fn embed(&mut self, text: String) {
        self.texts.push(text);
    }
}

// ── Embed trait ───────────────────────────────────────────────────────────────

/// Implemented by types that can be converted to text for embedding.
///
/// Use the [`derive`](irig_derive) macro or implement manually.
pub trait Embed {
    fn embed(&self, embedder: &mut TextEmbedder) -> Result<(), EmbedError>;
}

/// Extract all texts an object would embed without running the model.
pub fn to_texts(item: impl Embed) -> Result<Vec<String>, EmbedError> {
    let mut embedder = TextEmbedder::default();
    item.embed(&mut embedder)?;
    Ok(embedder.texts)
}

// ── Blanket impls for primitives ──────────────────────────────────────────────

impl Embed for String {
    fn embed(&self, e: &mut TextEmbedder) -> Result<(), EmbedError> {
        e.embed(self.clone());
        Ok(())
    }
}

impl Embed for &str {
    fn embed(&self, e: &mut TextEmbedder) -> Result<(), EmbedError> {
        e.embed(self.to_string());
        Ok(())
    }
}

impl Embed for serde_json::Value {
    fn embed(&self, e: &mut TextEmbedder) -> Result<(), EmbedError> {
        e.embed(serde_json::to_string(self).map_err(EmbedError::new)?);
        Ok(())
    }
}

impl<T: Embed> Embed for &T {
    fn embed(&self, e: &mut TextEmbedder) -> Result<(), EmbedError> {
        (*self).embed(e)
    }
}

impl<T: Embed> Embed for Vec<T> {
    fn embed(&self, e: &mut TextEmbedder) -> Result<(), EmbedError> {
        for item in self {
            item.embed(e)?;
        }
        Ok(())
    }
}

impl<T: Embed> Embed for Option<T> {
    fn embed(&self, e: &mut TextEmbedder) -> Result<(), EmbedError> {
        if let Some(inner) = self {
            inner.embed(e)?;
        }
        Ok(())
    }
}