Struct SqliteVectorIndex

pub struct SqliteVectorIndex<E: EmbeddingModel + 'static, T: SqliteVectorStoreTable + 'static> { /* private fields */ }

SQLite vector store implementation for Rig.

This crate provides a SQLite-based vector store implementation that can be used with Rig. It uses the sqlite-vec extension to enable vector similarity search capabilities.

Example

use rig::{
    embeddings::EmbeddingsBuilder,
    providers::openai::{Client, TEXT_EMBEDDING_ADA_002},
    vector_store::VectorStoreIndex,
    Embed,
};
use rig_sqlite::{Column, ColumnValue, SqliteVectorStore, SqliteVectorStoreTable};
use serde::Deserialize;
use tokio_rusqlite::Connection;

#[derive(Embed, Clone, Debug, Deserialize)]
struct Document {
    id: String,
    #[embed]
    content: String,
}

impl SqliteVectorStoreTable for Document {
    fn name() -> &'static str {
        "documents"
    }

    fn schema() -> Vec<Column> {
        vec![
            Column::new("id", "TEXT PRIMARY KEY"),
            Column::new("content", "TEXT"),
        ]
    }

    fn id(&self) -> String {
        self.id.clone()
    }

    fn column_values(&self) -> Vec<(&'static str, Box<dyn ColumnValue>)> {
        vec![
            ("id", Box::new(self.id.clone())),
            ("content", Box::new(self.content.clone())),
        ]
    }
}

let conn = Connection::open("vector_store.db").await?;
let openai_client = Client::new("YOUR_API_KEY");
let model = openai_client.embedding_model(TEXT_EMBEDDING_ADA_002);

// Initialize vector store
let vector_store = SqliteVectorStore::new(conn, &model).await?;

// Create documents
let documents = vec![
    Document {
        id: "doc1".to_string(),
        content: "Example document 1".to_string(),
    },
    Document {
        id: "doc2".to_string(),
        content: "Example document 2".to_string(),
    },
];

// Generate embeddings
let embeddings = EmbeddingsBuilder::new(model.clone())
    .documents(documents)?
    .build()
    .await?;

// Add to vector store
vector_store.add_rows(embeddings).await?;

// Create index and search
let index = vector_store.index(model);
let results = index
    .top_n::<Document>("Example query", 2)
    .await?;

Implementations

impl<E: EmbeddingModel + 'static, T: SqliteVectorStoreTable> SqliteVectorIndex<E, T>

pub fn new(embedding_model: E, store: SqliteVectorStore<E, T>) -> Self

Trait Implementations

impl<E: EmbeddingModel + Sync, T: SqliteVectorStoreTable> VectorStoreIndex for SqliteVectorIndex<E, T>

async fn top_n<D: for<'a> Deserialize<'a>>( &self, query: &str, n: usize, ) -> Result<Vec<(f64, String, D)>, 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.