rig 0.36.0

An opinionated library for building LLM powered applications.
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

//! Demonstrates vector search against a local Ollama embedding model.
//! Requires a local Ollama server and the `derive` feature.
//! Run it to compare semantic search results and matching document ids.

use rig::client::Nothing;
use rig::prelude::*;
use rig::providers::ollama;
use rig::vector_store::request::VectorSearchRequest;
use rig::{
    Embed,
    embeddings::EmbeddingsBuilder,
    vector_store::{VectorStoreIndex, in_memory_store::InMemoryVectorStore},
};

use serde::{Deserialize, Serialize};

type SearchMatch = (f64, String, String);

// 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 WordDefinition {
    id: String,
    word: String,
    #[embed]
    definitions: Vec<String>,
}

fn sample_documents() -> Vec<WordDefinition> {
    vec![
        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(),
            ],
        },
        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(),
            ],
        },
        WordDefinition {
            id: "doc2".to_string(),
            word: "linglingdong".to_string(),
            definitions: vec![
                "A term used by inhabitants of the sombrero galaxy to describe humans.".to_string(),
                "A rare, mystical instrument crafted by the ancient monks of the Nebulon Mountain Ranges on the planet Quarm.".to_string(),
            ],
        },
    ]
}

fn print_matches(label: &str, matches: &[SearchMatch]) {
    println!("{label}:");
    for (score, id, word) in matches {
        println!("  score={score:.4} id={id} word={word}");
    }
}

fn print_id_matches(label: &str, matches: &[(f64, String)]) {
    println!("{label}:");
    for (score, id) in matches {
        println!("  score={score:.4} id={id}");
    }
}

#[tokio::main]
async fn main() -> Result<(), anyhow::Error> {
    let client = ollama::Client::builder()
        .api_key(Nothing)
        .base_url("http://localhost:11434")
        .build()?;

    let embedding_model = client.embedding_model("nomic-embed-text");

    let embeddings = EmbeddingsBuilder::new(embedding_model.clone())
        .documents(sample_documents())?
        .build()
        .await?;

    let vector_store =
        InMemoryVectorStore::from_documents_with_id_f(embeddings, |doc| doc.id.clone());

    let index = vector_store.index(embedding_model);

    let query =
        "I need to buy something in a fictional universe. What type of money can I use for this?";
    let req = VectorSearchRequest::builder()
        .query(query)
        .samples(1)
        .build();

    let results = index
        .top_n::<WordDefinition>(req.clone())
        .await?
        .into_iter()
        .map(|(score, id, doc)| (score, id, doc.word))
        .collect::<Vec<SearchMatch>>();

    let id_results = index.top_n_ids(req).await?.into_iter().collect::<Vec<_>>();

    print_matches("Top document matches", &results);
    print_id_matches("Top document ids", &id_results);

    Ok(())
}