Crate chromadb

ChromaDB client library for Rust.

The library provides 2 modules to interact with the ChromaDB server via the API:

• client - To interface with the ChromaDB server.
• collection - To interface with an associated ChromaDB collection.

Instantiating ChromaClient

use chromadb::client::{ChromaAuthMethod, ChromaClient, ChromaClientOptions, ChromaTokenHeader};
use chromadb::collection::{ChromaCollection, GetResult, GetOptions};
use serde_json::json;

// With default ChromaClientOptions
// Defaults to http://localhost:8000
let client: ChromaClient = ChromaClient::new(Default::default()).await.unwrap();

// With custom ChromaClientOptions
let auth = ChromaAuthMethod::TokenAuth {
    token: "<TOKEN>".to_string(),
    header: ChromaTokenHeader::Authorization
};
let client: ChromaClient = ChromaClient::new(ChromaClientOptions {
    url: Some("<CHROMADB_URL>".to_string()),
    database: "<DATABASE>".to_string(),
    auth
}).await.unwrap();

Now that a client is instantiated, we can interface with the ChromaDB server and execute queries.

Collection Queries

// Get or create a collection with the given name and no metadata.
let collection: ChromaCollection = client.get_or_create_collection("my_collection", None).await?;

// Get the UUID of the collection
let collection_uuid = collection.id();
println!("Collection UUID: {}", collection_uuid);

// Upsert some embeddings with documents and no metadata.
let collection_entries = CollectionEntries {
   ids: vec!["demo-id-1", "demo-id-2"],
   embeddings: Some(vec![vec![0.0_f32; 768], vec![0.0_f32; 768]]),
   metadatas: None,
   documents: Some(vec![
       "Some document about 9 octopus recipies",
       "Some other document about DCEU Superman Vs CW Superman"
   ])
};

let result  = collection.upsert(collection_entries, None).await?;

// Create a filter object to filter by document content.
let where_document = json!({
   "$contains": "Superman"
    });

// Get embeddings from a collection with filters and limit set to 1.
// An empty IDs vec will return all embeddings.

let get_query = GetOptions {
    ids: vec![],
    where_metadata: None,
    limit: Some(1),
    offset: None,
    where_document: Some(where_document),
    include: Some(vec!["documents".into(),"embeddings".into()])
};

let get_result: GetResult = collection.get(get_query).await?;
println!("Get result: {:?}", get_result);

Find more information about on the available filters and options in the get() documentation.

Perform a similarity search.

//Instantiate QueryOptions to perform a similarity search on the collection
//Alternatively, an embedding_function can also be provided with query_texts to perform the search
let query = QueryOptions {
    query_texts: None,
    query_embeddings: Some(vec![vec![0.0_f32; 768], vec![0.0_f32; 768]]),
    where_metadata: None,
    where_document: None,
    n_results: Some(5),
    include: None,
};

let query_result: QueryResult = collection.query(query, None).await?;
println!("Query result: {:?}", query_result);

Support for Embedding providers

This crate has built-in support for OpenAI and SBERT embeddings.

To use OpenAI embeddings, enable the openai feature in your Cargo.toml.

let collection: ChromaCollection = client.get_or_create_collection("openai_collection",
None).await?;

let collection_entries = CollectionEntries {
  ids: vec!["demo-id-1", "demo-id-2"],
  embeddings: None,
  metadatas: None,
  documents: Some(vec![
           "Some document about 9 octopus recipies",
           "Some other document about DCEU Superman Vs CW Superman"])
};

// Use OpenAI embeddings
let openai_embeddings = OpenAIEmbeddings::new(Default::default());
collection.upsert(collection_entries, Some(Box::new(openai_embeddings))).await?;
Ok(())

To use SBERT embeddings, enable the bert feature in your Cargo.toml.

let collection: ChromaCollection = client.get_or_create_collection("sbert_collection",
None).await?;

let collection_entries = CollectionEntries {
  ids: vec!["demo-id-1", "demo-id-2"],
  embeddings: None,
  metadatas: None,
  documents: Some(vec![
           "Some document about 9 octopus recipies",
           "Some other document about DCEU Superman Vs CW Superman"])
};

// Use SBERT embeddings
let sbert_embeddings = SentenceEmbeddingsBuilder::remote(
                        SentenceEmbeddingsModelType::AllMiniLmL6V2
                       ).create_model()?;

collection.upsert(collection_entries, Some(Box::new(sbert_embeddings))).await?;

Re-exports

pub use client::ChromaClient;

pub use collection::ChromaCollection;

Modules

client

collection

embeddings