Feat(embeddings): Embeddable derive macro

Cargo.toml

@@ -1,6 +1,7 @@
 [workspace]
 resolver = "2"
 members = [
-    "rig-core", "rig-lancedb",
+    "rig-core", "rig-core/rig-core-derive",
     "rig-mongodb",
+    "rig-lancedb"
 ]

rig-core/Cargo.toml

@@ -23,8 +23,16 @@ futures = "0.3.29"
 ordered-float = "4.2.0"
 schemars = "0.8.16"
 thiserror = "1.0.61"
+rig-derive = { path = "./rig-core-derive", optional = true }
 
 [dev-dependencies]
 anyhow = "1.0.75"
 tokio = { version = "1.34.0", features = ["full"] }
 tracing-subscriber = "0.3.18"
+
+[features]
+rig_derive = ["dep:rig-derive"]
+
+[[test]]
+name = "embeddable_macro"
+required-features = ["rig_derive"]

rig-core/examples/calculator_chatbot.rs

@@ -2,10 +2,11 @@ use anyhow::Result;
 use rig::{
     cli_chatbot::cli_chatbot,
     completion::ToolDefinition,
-    embeddings::{DocumentEmbeddings, EmbeddingsBuilder},
+    embeddings::builder::DocumentEmbeddings,
     providers::openai::{Client, TEXT_EMBEDDING_ADA_002},
     tool::{Tool, ToolEmbedding, ToolSet},
     vector_store::in_memory_store::InMemoryVectorStore,
+    EmbeddingsBuilder,
 };
 use serde::{Deserialize, Serialize};
 use serde_json::json;

rig-core/examples/rag.rs

@@ -2,9 +2,10 @@ use std::env;
 
 use rig::{
     completion::Prompt,
-    embeddings::{DocumentEmbeddings, EmbeddingsBuilder},
+    embeddings::builder::DocumentEmbeddings,
     providers::openai::{Client, TEXT_EMBEDDING_ADA_002},
     vector_store::in_memory_store::InMemoryVectorStore,
+    EmbeddingsBuilder,
 };
 
 #[tokio::main]

rig-core/examples/rag_dynamic_tools.rs

@@ -1,10 +1,11 @@
 use anyhow::Result;
 use rig::{
     completion::{Prompt, ToolDefinition},
-    embeddings::{DocumentEmbeddings, EmbeddingsBuilder},
+    embeddings::builder::DocumentEmbeddings,
     providers::openai::{Client, TEXT_EMBEDDING_ADA_002},
     tool::{Tool, ToolEmbedding, ToolSet},
     vector_store::in_memory_store::InMemoryVectorStore,
+    EmbeddingsBuilder,
 };
 use serde::{Deserialize, Serialize};
 use serde_json::json;

rig-core/examples/vector_search.rs

@@ -1,9 +1,10 @@
 use std::env;
 
 use rig::{
-    embeddings::{DocumentEmbeddings, EmbeddingsBuilder},
+    embeddings::builder::DocumentEmbeddings,
     providers::openai::{Client, TEXT_EMBEDDING_ADA_002},
     vector_store::{in_memory_store::InMemoryVectorStore, VectorStoreIndex},
+    EmbeddingsBuilder,
 };
 
 #[tokio::main]

rig-core/examples/vector_search_cohere.rs

@@ -1,9 +1,10 @@
 use std::env;
 
 use rig::{
-    embeddings::{DocumentEmbeddings, EmbeddingsBuilder},
+    embeddings::builder::DocumentEmbeddings,
     providers::cohere::{Client, EMBED_ENGLISH_V3},
     vector_store::{in_memory_store::InMemoryVectorStore, VectorStoreIndex},
+    EmbeddingsBuilder,
 };
 
 #[tokio::main]

rig-core/rig-core-derive/Cargo.toml

@@ -0,0 +1,13 @@
+[package]
+name = "rig-derive"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+indoc = "2.0.5"
+proc-macro2 = { version = "1.0.87", features = ["proc-macro"] }
+quote = "1.0.37"
+syn = { version = "2.0.79", features = ["full"]}
+
+[lib]
+proc-macro = true

rig-core/rig-core-derive/src/basic.rs

@@ -0,0 +1,29 @@
+use syn::{parse_quote, Attribute, DataStruct, Meta};
+
+use crate::EMBED;
+
+/// Finds and returns fields with simple #[embed] attribute tags only.
+pub(crate) fn basic_embed_fields(data_struct: &DataStruct) -> impl Iterator<Item = syn::Field> {
+    data_struct.fields.clone().into_iter().filter(|field| {
+        field
+            .attrs
+            .clone()
+            .into_iter()
+            .any(|attribute| match attribute {
+                Attribute {
+                    meta: Meta::Path(path),
+                    ..
+                } => path.is_ident(EMBED),
+                _ => false,
+            })
+    })
+}
+
+// Adds bounds to where clause that force all fields tagged with #[embed] to implement the Embeddable trait.
+pub(crate) fn add_struct_bounds(generics: &mut syn::Generics, field_type: &syn::Type) {
+    let where_clause = generics.make_where_clause();
+
+    where_clause.predicates.push(parse_quote! {
+        #field_type: Embeddable
+    });
+}

rig-core/rig-core-derive/src/custom.rs

@@ -0,0 +1,125 @@
+use quote::ToTokens;
+use syn::{meta::ParseNestedMeta, ExprPath};
+
+use crate::EMBED;
+
+const EMBED_WITH: &str = "embed_with";
+
+/// Finds and returns fields with #[embed(embed_with = "...")] attribute tags only.
+/// Also returns the "..." part of the tag (ie. the custom function).
+pub(crate) fn custom_embed_fields(
+    data_struct: &syn::DataStruct,
+) -> syn::Result<Vec<(syn::Field, syn::ExprPath)>> {
+    data_struct
+        .fields
+        .clone()
+        .into_iter()
+        .filter_map(|field| {
+            field
+                .attrs
+                .clone()
+                .into_iter()
+                .filter_map(|attribute| match attribute.is_custom() {
+                    Ok(true) => match attribute.expand_tag() {
+                        Ok(path) => Some(Ok((field.clone(), path))),
+                        Err(e) => Some(Err(e)),
+                    },
+                    Ok(false) => None,
+                    Err(e) => Some(Err(e)),
+                })
+                .next()
+        })
+        .collect::<Result<Vec<_>, _>>()
+}
+
+trait CustomAttributeParser {
+    // Determine if field is tagged with an #[embed(embed_with = "...")] attribute.
+    fn is_custom(&self) -> syn::Result<bool>;
+
+    // Get the "..." part of the #[embed(embed_with = "...")] attribute.
+    // Ex: If attribute is tagged with #[embed(embed_with = "my_embed")], returns "my_embed".
+    fn expand_tag(&self) -> syn::Result<syn::ExprPath>;
+}
+
+impl CustomAttributeParser for syn::Attribute {
+    fn is_custom(&self) -> syn::Result<bool> {
+        // Check that the attribute is a list.
+        match &self.meta {
+            syn::Meta::List(meta) => {
+                if meta.tokens.is_empty() {
+                    return Ok(false);
+                }
+            }
+            _ => return Ok(false),
+        };
+
+        // Check the first attribute tag (the first "embed")
+        if !self.path().is_ident(EMBED) {
+            return Ok(false);
+        }
+
+        self.parse_nested_meta(|meta| {
+            // Parse the meta attribute as an expression. Need this to compile.
+            meta.value()?.parse::<syn::Expr>()?;
+
+            if meta.path.is_ident(EMBED_WITH) {
+                Ok(())
+            } else {
+                let path = meta.path.to_token_stream().to_string().replace(' ', "");
+                Err(syn::Error::new_spanned(
+                    meta.path,
+                    format_args!("unknown embedding field attribute `{}`", path),
+                ))
+            }
+        })?;
+
+        Ok(true)
+    }
+
+    fn expand_tag(&self) -> syn::Result<syn::ExprPath> {
+        fn function_path(meta: &ParseNestedMeta<'_>) -> syn::Result<ExprPath> {
+            // #[embed(embed_with = "...")]
+            let expr = meta.value()?.parse::<syn::Expr>().unwrap();
+            let mut value = &expr;
+            while let syn::Expr::Group(e) = value {
+                value = &e.expr;
+            }
+            let string = if let syn::Expr::Lit(syn::ExprLit {
+                lit: syn::Lit::Str(lit_str),
+                ..
+            }) = value
+            {
+                let suffix = lit_str.suffix();
+                if !suffix.is_empty() {
+                    return Err(syn::Error::new_spanned(
+                        lit_str,
+                        format!("unexpected suffix `{}` on string literal", suffix),
+                    ));
+                }
+                lit_str.clone()
+            } else {
+                return Err(syn::Error::new_spanned(
+                    value,
+                    format!(
+                        "expected {} attribute to be a string: `{} = \"...\"`",
+                        EMBED_WITH, EMBED_WITH
+                    ),
+                ));
+            };
+
+            string.parse()
+        }
+
+        let mut custom_func_path = None;
+
+        self.parse_nested_meta(|meta| match function_path(&meta) {
+            Ok(path) => {
+                custom_func_path = Some(path);
+                Ok(())
+            }
+            Err(e) => Err(e),
+        })?;
+
+        Ok(custom_func_path.unwrap())
+    }
+}

rig-core/rig-core-derive/src/embeddable.rs

@@ -0,0 +1,134 @@
+use proc_macro2::TokenStream;
+use quote::quote;
+use syn::DataStruct;
+
+use crate::{
+    basic::{add_struct_bounds, basic_embed_fields},
+    custom::custom_embed_fields,
+};
+
+pub(crate) fn expand_derive_embedding(input: &mut syn::DeriveInput) -> syn::Result<TokenStream> {
+    let name = &input.ident;
+    let data = &input.data;
+    let generics = &mut input.generics;
+
+    let target_stream = match data {
+        syn::Data::Struct(data_struct) => {
+            let (basic_targets, basic_target_size) = data_struct.basic(generics);
+            let (custom_targets, custom_target_size) = data_struct.custom()?;
+
+            // If there are no fields tagged with #[embed] or #[embed(embed_with = "...")], return an empty TokenStream.
+            // ie. do not implement Embeddable trait for the struct.
+            if basic_target_size + custom_target_size == 0 {
+                return Err(syn::Error::new_spanned(
+                    name,
+                    "Add at least one field tagged with #[embed] or #[embed(embed_with = \"...\")].",
+                ));
+            }
+
+            // Determine whether the Embeddable::Kind should be SingleEmbedding or ManyEmbedding
+            quote! {
+                let mut embed_targets = #basic_targets;
+                embed_targets.extend(#custom_targets)
+            }
+        }
+        _ => {
+            return Err(syn::Error::new_spanned(
+                input,
+                "Embeddable derive macro should only be used on structs",
+            ))
+        }
+    };
+
+    let (impl_generics, ty_generics, where_clause) = input.generics.split_for_impl();
+
+    let gen = quote! {
+        // Note: Embeddable trait is imported with the macro.
+
+        impl #impl_generics Embeddable for #name #ty_generics #where_clause {
+            type Error = rig::embeddings::embeddable::EmbeddableError;
+
+            fn embeddable(&self) -> Result<rig::embeddings::embeddable::OneOrMany<String>, Self::Error> {
+                #target_stream;
+
+                Ok(rig::embeddings::embeddable::OneOrMany::from(
+                    embed_targets.into_iter()
+                        .collect::<Result<Vec<_>, _>>()?
+                ))
+            }
+        }
+    };
+
+    Ok(gen)
+}
+
+trait StructParser {
+    // Handles fields tagged with #[embed]
+    fn basic(&self, generics: &mut syn::Generics) -> (TokenStream, usize);
+
+    // Handles fields tagged with #[embed(embed_with = "...")]
+    fn custom(&self) -> syn::Result<(TokenStream, usize)>;
+}
+
+impl StructParser for DataStruct {
+    fn basic(&self, generics: &mut syn::Generics) -> (TokenStream, usize) {
+        let embed_targets = basic_embed_fields(self)
+            // Iterate over every field tagged with #[embed]
+            .map(|field| {
+                add_struct_bounds(generics, &field.ty);
+
+                let field_name = field.ident;
+
+                quote! {
+                    self.#field_name
+                }
+            })
+            .collect::<Vec<_>>();
+
+        if !embed_targets.is_empty() {
+            (
+                quote! {
+                    vec![#(#embed_targets.embeddable()),*]
+                },
+                embed_targets.len(),
+            )
+        } else {
+            (
+                quote! {
+                    vec![]
+                },
+                0,
+            )
+        }
+    }
+
+    fn custom(&self) -> syn::Result<(TokenStream, usize)> {
+        let embed_targets = custom_embed_fields(self)?
+            // Iterate over every field tagged with #[embed(embed_with = "...")]
+            .into_iter()
+            .map(|(field, custom_func_path)| {
+                let field_name = field.ident;
+
+                quote! {
+                    #custom_func_path(self.#field_name.clone())
+                }
+            })
+            .collect::<Vec<_>>();
+
+        Ok(if !embed_targets.is_empty() {
+            (
+                quote! {
+                    vec![#(#embed_targets),*]
+                },
+                embed_targets.len(),
+            )
+        } else {
+            (
+                quote! {
+                    vec![]
+                },
+                0,
+            )
+        })
+    }
+}