Feat: small improvements + fixes + tests

rig-core/Cargo.toml

@@ -35,8 +35,8 @@ tracing-subscriber = "0.3.18"
 tokio-test = "0.4.4"
 
 [features]
+all = ["derive", "pdf"]
 derive = ["dep:rig-derive"]
-all = ["pdf"]
 pdf = ["dep:lopdf"]
 
 [[test]]

rig-core/examples/rag.rs

@@ -10,17 +10,24 @@ use rig::{
 use serde::Serialize;
 
 // Data to be RAGged.
-// A vector search needs to be performed on the `definitions` field, so we derive the `Embed` trait for `FakeDefinition`
+// A vector search needs to be performed on the `definitions` field, so we derive the `Embed` trait for `WordDefinition`
 // and tag that field with `#[embed]`.
 #[derive(Embed, Serialize, Clone, Debug, Eq, PartialEq, Default)]
-struct FakeDefinition {
+struct WordDefinition {
     id: String,
+    word: String,
     #[embed]
     definitions: Vec<String>,
 }
 
 #[tokio::main]
 async fn main() -> Result<(), anyhow::Error> {
+    // Initialize tracing
+    tracing_subscriber::fmt()
+        .with_max_level(tracing::Level::INFO)
+        .with_target(false)
+        .init();
+
     // Create OpenAI client
     let openai_api_key = env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY not set");
     let openai_client = Client::new(&openai_api_key);
@@ -30,25 +37,28 @@ async fn main() -> Result<(), anyhow::Error> {
     // Generate embeddings for the definitions of all the documents using the specified embedding model.
     let embeddings = EmbeddingsBuilder::new(embedding_model.clone())
         .documents(vec![
-            FakeDefinition {
+            WordDefinition {
                 id: "doc0".to_string(),
+                word: "flurbo".to_string(),
                 definitions: vec![
-                    "Definition of a *flurbo*: A flurbo is a green alien that lives on cold planets.".to_string(),
-                    "Definition of a *flurbo*: A fictional digital currency that originated in the animated series Rick and Morty.".to_string()
+                    "1. *flurbo* (name): A flurbo is a green alien that lives on cold planets.".to_string(),
+                    "2. *flurbo* (name): A fictional digital currency that originated in the animated series Rick and Morty.".to_string()
                 ]
             },
-            FakeDefinition {
+            WordDefinition {
                 id: "doc1".to_string(),
+                word: "glarb-glarb".to_string(),
                 definitions: vec![
-                    "Definition of a *glarb-glarb*: A glarb-glarb is a ancient tool used by the ancestors of the inhabitants of planet Jiro to farm the land.".to_string(),
-                    "Definition of a *glarb-glarb*: A fictional creature found in the distant, swampy marshlands of the planet Glibbo in the Andromeda galaxy.".to_string()
+                    "1. *glarb-glarb* (noun): A glarb-glarb is a ancient tool used by the ancestors of the inhabitants of planet Jiro to farm the land.".to_string(),
+                    "2. *glarb-glarb* (noun): A fictional creature found in the distant, swampy marshlands of the planet Glibbo in the Andromeda galaxy.".to_string()
                 ]
             },
-            FakeDefinition {
+            WordDefinition {
                 id: "doc2".to_string(),
+                word: "linglingdong".to_string(),
                 definitions: vec![
-                    "Definition of a *linglingdong*: A term used by inhabitants of the far side of the moon to describe humans.".to_string(),
-                    "Definition of a *linglingdong*: A rare, mystical instrument crafted by the ancient monks of the Nebulon Mountain Ranges on the planet Quarm.".to_string()
+                    "1. *linglingdong* (noun): A term used by inhabitants of the far side of the moon to describe humans.".to_string(),
+                    "2. *linglingdong* (noun): A rare, mystical instrument crafted by the ancient monks of the Nebulon Mountain Ranges on the planet Quarm.".to_string()
                 ]
             },
         ])?

rig-core/examples/rag_dynamic_tools.rs

@@ -137,11 +137,6 @@ async fn main() -> Result<(), anyhow::Error> {
         .with_max_level(tracing::Level::INFO)
         // disable printing the name of the module in every log line.
         .with_target(false)
-        // this needs to be set to false, otherwise ANSI color codes will
-        // show up in a confusing manner in CloudWatch logs.
-        .with_ansi(false)
-        // disabling time is handy because CloudWatch will add the ingestion time.
-        .without_time()
         .init();
 
     // Create OpenAI client

rig-core/examples/vector_search.rs

@@ -11,7 +11,7 @@ use serde::{Deserialize, Serialize};
 // Shape of data that needs to be RAG'ed.
 // The definition field will be used to generate embeddings.
 #[derive(Embed, Clone, Deserialize, Debug, Serialize, Eq, PartialEq, Default)]
-struct FakeDefinition {
+struct WordDefinition {
     id: String,
     word: String,
     #[embed]
@@ -28,23 +28,23 @@ async fn main() -> Result<(), anyhow::Error> {
 
     let embeddings = EmbeddingsBuilder::new(model.clone())
         .documents(vec![
-            FakeDefinition {
+            WordDefinition {
                 id: "doc0".to_string(),
                 word: "flurbo".to_string(),
                 definitions: vec![
                     "A green alien that lives on cold planets.".to_string(),
                     "A fictional digital currency that originated in the animated series Rick and Morty.".to_string()
                 ]
             },
-            FakeDefinition {
+            WordDefinition {
                 id: "doc1".to_string(),
                 word: "glarb-glarb".to_string(),
                 definitions: vec![
                     "An ancient tool used by the ancestors of the inhabitants of planet Jiro to farm the land.".to_string(),
                     "A fictional creature found in the distant, swampy marshlands of the planet Glibbo in the Andromeda galaxy.".to_string()
                 ]
             },
-            FakeDefinition {
+            WordDefinition {
                 id: "doc2".to_string(),
                 word: "linglingdong".to_string(),
                 definitions: vec![
@@ -61,7 +61,7 @@ async fn main() -> Result<(), anyhow::Error> {
         .index(model);
 
     let results = index
-        .top_n::<FakeDefinition>("I need to buy something in a fictional universe. What type of money can I use for this?", 1)
+        .top_n::<WordDefinition>("I need to buy something in a fictional universe. What type of money can I use for this?", 1)
         .await?
         .into_iter()
         .map(|(score, id, doc)| (score, id, doc.word))

rig-core/examples/vector_search_cohere.rs

@@ -11,7 +11,7 @@ use serde::{Deserialize, Serialize};
 // Shape of data that needs to be RAG'ed.
 // The definition field will be used to generate embeddings.
 #[derive(Embed, Clone, Deserialize, Debug, Serialize, Eq, PartialEq, Default)]
-struct FakeDefinition {
+struct WordDefinition {
     id: String,
     word: String,
     #[embed]
@@ -29,23 +29,23 @@ async fn main() -> Result<(), anyhow::Error> {
 
     let embeddings = EmbeddingsBuilder::new(document_model.clone())
         .documents(vec![
-            FakeDefinition {
+            WordDefinition {
                 id: "doc0".to_string(),
                 word: "flurbo".to_string(),
                 definitions: vec![
                     "A green alien that lives on cold planets.".to_string(),
                     "A fictional digital currency that originated in the animated series Rick and Morty.".to_string()
                 ]
             },
-            FakeDefinition {
+            WordDefinition {
                 id: "doc1".to_string(),
                 word: "glarb-glarb".to_string(),
                 definitions: vec![
                     "An ancient tool used by the ancestors of the inhabitants of planet Jiro to farm the land.".to_string(),
                     "A fictional creature found in the distant, swampy marshlands of the planet Glibbo in the Andromeda galaxy.".to_string()
                 ]
             },
-            FakeDefinition {
+            WordDefinition {
                 id: "doc2".to_string(),
                 word: "linglingdong".to_string(),
                 definitions: vec![
@@ -62,7 +62,7 @@ async fn main() -> Result<(), anyhow::Error> {
         .index(search_model);
 
     let results = index
-        .top_n::<FakeDefinition>(
+        .top_n::<WordDefinition>(
             "Which instrument is found in the Nebulon Mountain Ranges?",
             1,
         )