Struct LanceDbVectorStore

pub struct LanceDbVectorStore<M: EmbeddingModel> { /* private fields */ }

Example

use std::{env, sync::Arc};

use arrow_array::RecordBatchIterator;
use fixture::{as_record_batch, schema};
use rig::{
    embeddings::{EmbeddingModel, EmbeddingsBuilder},
    providers::openai::{Client, TEXT_EMBEDDING_ADA_002},
    vector_store::VectorStoreIndexDyn,
};
use rig_lancedb::{LanceDbVectorStore, SearchParams};
use serde::Deserialize;

#[derive(Deserialize, Debug)]
pub struct VectorSearchResult {
    pub id: String,
    pub content: String,
}

// Initialize OpenAI client. Use this to generate embeddings (and generate test data for RAG demo).
let openai_api_key = env::var("OPENAI_API_KEY").expect("OPENAI_API_KEY not set");
let openai_client = Client::new(&openai_api_key);

// Select the embedding model and generate our embeddings
let model = openai_client.embedding_model(TEXT_EMBEDDING_ADA_002);

let embeddings = EmbeddingsBuilder::new(model.clone())
    .simple_document("doc0", "Definition of *flumbrel (noun)*: a small, seemingly insignificant item that you constantly lose or misplace, such as a pen, hair tie, or remote control.")
    .simple_document("doc1", "Definition of *zindle (verb)*: to pretend to be working on something important while actually doing something completely unrelated or unproductive")
    .simple_document("doc2", "Definition of *glimber (adjective)*: describing a state of excitement mixed with nervousness, often experienced before an important event or decision.")
    .build()
    .await?;

// Define search_params params that will be used by the vector store to perform the vector search.
let search_params = SearchParams::default();

// Initialize LanceDB locally.
let db = lancedb::connect("data/lancedb-store").execute().await?;

// Create table with embeddings.
let record_batch = as_record_batch(embeddings, model.ndims());
let table = db
    .create_table(
        "definitions",
        RecordBatchIterator::new(vec![record_batch], Arc::new(schema(model.ndims()))),
    )
    .execute()
    .await?;

let vector_store = LanceDbVectorStore::new(table, model, "id", search_params).await?;

// Query the index
let results = vector_store
.top_n("My boss says I zindle too much, what does that mean?", 1)
.await?
.into_iter()
.map(|(score, id, doc)| {
    anyhow::Ok((
        score,
        id,
        serde_json::from_value::<VectorSearchResult>(doc)?,
    ))
})
.collect::<Result<Vec<_>, _>>()?;

println!("Results: {:?}", results);

Implementations

impl<M: EmbeddingModel> LanceDbVectorStore<M>

pub async fn new( table: Table, model: M, id_field: &str, search_params: SearchParams, ) -> Result<Self, Error>

Create an instance of LanceDbVectorStore with an existing table and model. Define the id field name of the table. Define search parameters that will be used to perform vector searches on the table.

Trait Implementations

impl<M: EmbeddingModel + Sync + Send> VectorStoreIndex for LanceDbVectorStore<M>

async fn top_n<T: for<'a> Deserialize<'a> + Send>( &self, query: &str, n: usize, ) -> Result<Vec<(f64, String, T)>, VectorStoreError>

Get the top n documents based on the distance to the given query. The result is a list of tuples of the form (score, id, document)

async fn top_n_ids( &self, query: &str, n: usize, ) -> Result<Vec<(f64, String)>, VectorStoreError>

Same as top_n but returns the document ids only.