feat: Add support for Chroma's Database feature. (#16)

* Add support for Chroma's Database feature.

Different collections can be in different databases.  Passing the
optional database argument at client creation time binds the client to
the appropriate database.

* Fix doctest example

* chore: Review updates

Signed-off-by: Anush008 <[email protected]>

* docs: nit update README.md

Signed-off-by: Anush008 <[email protected]>

---------

Co-authored-by: Anush008 <[email protected]>

Author: rescrv

Cargo.lock

@@ -1,59 +1,57 @@
 <div align="center">
   <h1><a href="https://crates.io/crates/chromadb">ChromaDB-rs</a></h1>
-  <h3>A Rust based client library for the Chroma vector database.</h3>
+  <h3>A Rust client library for the Chroma vector database.</h3>
   <a href="https://crates.io/crates/chromadb"><img src="https://img.shields.io/crates/v/chromadb.svg" alt="Crates.io"></a>
   <a href="https://github.com/Anush008/chromadb-rs/blob/master/LICENSE"><img src="https://img.shields.io/badge/license-mit-blue.svg" alt="MIT Licensed"></a>
   <a href="https://github.com/Anush008/chromadb-rs/actions/workflows/cargo-test.yml"><img src="https://github.com/Anush008/chromadb-rs/actions/workflows/release.yml/badge.svg?branch=master" alt="Tests"></a>
 </div>
 
-## ⚙️ Running ChromaDB
-> ℹ Chroma can be run in-memory in Python (without Docker), but this feature is not yet available in other languages.
-> To use this library you either need a hosted or local version of ChromaDB running.
-
-If you can run `docker-compose up -d --build` you can run Chroma.
-
-```shell
-git clone https://github.com/chroma-core/chroma.git
-cd chroma
-# Run a ChromaDB instance at localhost:8000
-docker-compose up -d --build
-```
-
-More information about deploying Chroma to production can be found [here](https://docs.trychroma.com/deployment).
-
 ## 🚀 Installing the library
+
 ```shell
 cargo add chromadb
 ```
-The library crate can be found at [crates.io](https://crates.io/crates/chromadb).
+
+The crate can be found at [crates.io](https://crates.io/crates/chromadb).
 
 ## 📖 Documentation
+
 The library reference can be found [here](https://docs.rs/chromadb).
 
 ## 🔍 Overview
 
-#### The library provides 2 modules to interact with the ChromaDB server via API V1:
- * `client` - To interface with the ChromaDB server.
- * `collection` - To interface with an associated ChromaDB collection.
+#### The library provides 2 modules to interact with the ChromaDB server via API V1
+
+* `client` - To interface with the ChromaDB server.
+* `collection` - To interface with an associated ChromaDB collection.
 
 #### You can connect to ChromaDB by instantiating a [ChromaClient](https://docs.rs/chromadb/latest/chromadb/v1/client/struct.ChromaClient.html)
- 
+
  ```rust
-use chromadb::v2::ChromaClient;
+use chromadb::v2::client::{ChromaAuthMethod, ChromaClient, ChromaClientOptions, ChromaTokenHeader};
 use chromadb::v2::collection::{ChromaCollection, GetQuery, GetResult, CollectionEntries};
-use serde_json::json;
 
 // With default ChromaClientOptions
 // Defaults to http://localhost:8000
 let client: ChromaClient = ChromaClient::new(Default::default());
 
 // With custom ChromaClientOptions
-let client: ChromaClient = ChromaClient::new(ChromaClientOptions { url: "<CHROMADB_URL>".into() });
+let auth = ChromaAuthMethod::TokenAuth {
+    token: "<TOKEN>".to_string(),
+    header: ChromaTokenHeader::Authorization
+};
+let client: ChromaClient = ChromaClient::new(ChromaClientOptions {
+    url: "<CHROMADB_URL>".into(),
+    database: Some("<DATABASE>".into()),
+    auth
+});
 ```
 
-#### Now that a client is instantiated, we can interface with the ChromaDB server.
+#### Now that a client is instantiated, we can interface with the ChromaDB server
 
  ```rust
+use serde_json::json;
+
 // Get or create a collection with the given name and no metadata.
 let collection: ChromaCollection = client.get_or_create_collection("my_collection", None).await?;
 
@@ -62,7 +60,7 @@ let collection_uuid = collection.id();
 println!("Collection UUID: {}", collection_uuid);
 ```
 
-###  With a collection instance, we can perform queries on the database
+### With a collection instance, we can perform queries on the database
 
 ```rust
 // Upsert some embeddings with documents and no metadata.
@@ -81,7 +79,7 @@ let result: bool = 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.
@@ -97,10 +95,11 @@ let get_result: GetResult = collection.get(get_query).await?;
 println!("Get result: {:?}", get_result);
 
 ```
-Find more information about the available filters and options in the [get()](https://docs.rs/chromadb/latest/chromadb/v1/collection/struct.ChromaCollection.html#method.get) documentation.
 
+Find more information about the available filters and options in the [get()](https://docs.rs/chromadb/latest/chromadb/v1/collection/struct.ChromaCollection.html#method.get) documentation.
 
 ### Performing a similarity search
+
 ```rust
 //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
@@ -117,7 +116,8 @@ let query_result: QueryResult = collection.query(query, None).await?;
 println!("Query result: {:?}", query_result);
 ```
 
- ### Support for Embedding providers
+### Support for Embedding providers
+
  This crate has built-in support for OpenAI and SBERT embeddings.
 
  To use [OpenAI](https://platform.openai.com/docs/guides/embeddings) embeddings, enable the `openai` feature in your Cargo.toml.

README.md

@@ -6,7 +6,7 @@
 //!
 //! ### Instantiating [ChromaClient](crate::v2::ChromaClient)
 //! ```
-//! use chromadb::v2::client::{ChromaAuthMethod, ChromaClient, ChromaClientOptions};
+//! use chromadb::v2::client::{ChromaAuthMethod, ChromaClient, ChromaClientOptions, ChromaTokenHeader};
 //! use chromadb::v2::collection::{ChromaCollection, GetResult, GetOptions};
 //! use serde_json::json;
 //!
@@ -16,13 +16,14 @@
 //! let client: ChromaClient = ChromaClient::new(Default::default());
 //!
 //! // With custom ChromaClientOptions
-//! let auth = ChromaAuthMethod::TokenAuth { 
-//!     token: "<TOKEN>".to_string(), 
-//!     header: chromadb::v2::client::ChromaTokenHeader::Authorization 
+//! let auth = ChromaAuthMethod::TokenAuth {
+//!     token: "<TOKEN>".to_string(),
+//!     header: ChromaTokenHeader::Authorization
 //! };
-//! let client: ChromaClient = ChromaClient::new(ChromaClientOptions { 
+//! let client: ChromaClient = ChromaClient::new(ChromaClientOptions {
 //!     url: "<CHROMADB_URL>".into(),
-//!     auth 
+//!     database: Some("<DATABASE>".into()),
+//!     auth
 //! });
 //!
 //! # Ok(())

src/lib.rs

@@ -1,6 +1,6 @@
 use super::commons::Result;
 use base64::prelude::*;
-use reqwest::{Client, Response, Method};
+use reqwest::{Client, Method, Response};
 use serde_json::Value;
 
 #[derive(Clone, Debug)]
@@ -12,7 +12,10 @@ pub enum ChromaTokenHeader {
 #[derive(Clone, Debug)]
 pub enum ChromaAuthMethod {
     None,
-    BasicAuth { username: String, password: String },
+    BasicAuth {
+        username: String,
+        password: String,
+    },
     TokenAuth {
         token: String,
         header: ChromaTokenHeader,
@@ -29,14 +32,16 @@ impl Default for ChromaAuthMethod {
 pub(super) struct APIClientAsync {
     pub(super) api_endpoint: String,
     pub(super) auth_method: ChromaAuthMethod,
+    pub(super) database: Option<String>,
     client: Client,
 }
 
 impl APIClientAsync {
-    pub fn new(endpoint: String, auth_method: ChromaAuthMethod) -> Self {
+    pub fn new(endpoint: String, auth_method: ChromaAuthMethod, database: Option<String>) -> Self {
         Self {
             api_endpoint: format!("{}/api/v1", endpoint),
             auth_method,
+            database,
             client: Client::new(),
         }
     }
@@ -57,29 +62,35 @@ impl APIClientAsync {
         self.send_request(Method::DELETE, path, None).await
     }
 
-    async fn send_request(&self, method: Method, path: &str, json_body: Option<Value>) -> Result<Response> {
-        let url = format!("{}{}", self.api_endpoint, path);
-        
-        let mut request = self.client
-            .request(method, &url);
+    async fn send_request(
+        &self,
+        method: Method,
+        path: &str,
+        json_body: Option<Value>,
+    ) -> Result<Response> {
+        let url = if let Some(database) = &self.database {
+            format!("{}{}?database={}", self.api_endpoint, path, database)
+        } else {
+            format!("{}{}", self.api_endpoint, path,)
+        };
+
+        let mut request = self.client.request(method, &url);
 
         // Add auth headers if needed
         match &self.auth_method {
-            ChromaAuthMethod::None => {},
+            ChromaAuthMethod::None => {}
             ChromaAuthMethod::BasicAuth { username, password } => {
                 let credentials = BASE64_STANDARD.encode(format!("{username}:{password}"));
                 request = request.header("Authorization", format!("Basic {credentials}"));
             }
-            ChromaAuthMethod::TokenAuth { token, header } => {
-                match header {
-                    ChromaTokenHeader::Authorization => {
-                        request = request.header("Authorization", format!("Bearer {token}"));
-                    }
-                    ChromaTokenHeader::XChromaToken => {
-                        request = request.header("X-Chroma-Token", token);
-                    }
+            ChromaAuthMethod::TokenAuth { token, header } => match header {
+                ChromaTokenHeader::Authorization => {
+                    request = request.header("Authorization", format!("Bearer {token}"));
                 }
-            }
+                ChromaTokenHeader::XChromaToken => {
+                    request = request.header("X-Chroma-Token", token);
+                }
+            },
         }
 
         // Add JSON body if present
@@ -97,7 +108,7 @@ impl APIClientAsync {
         } else {
             let error_text = response.text().await?;
             anyhow::bail!(
-                "{} {}: {}", 
+                "{} {}: {}",
                 status.as_u16(),
                 status.canonical_reason().unwrap_or("Unknown"),
                 error_text

src/v2/api.rs

@@ -20,22 +20,32 @@ pub struct ChromaClient {
 /// The options for instantiating ChromaClient.
 #[derive(Debug, Default)]
 pub struct ChromaClientOptions {
+    /// The URL of the Chroma Server.
     pub url: String,
+    /// Authentication to use to connect to the Chroma Server.
     pub auth: ChromaAuthMethod,
+    /// Optional database name to scope all queries to.
+    pub database: Option<String>,
 }
 
 impl ChromaClient {
     /// Create a new Chroma client with the given options.
     /// * Defaults to `url`: http://localhost:8000
-    pub fn new(ChromaClientOptions { url, auth }: ChromaClientOptions) -> ChromaClient {
+    pub fn new(
+        ChromaClientOptions {
+            url,
+            auth,
+            database,
+        }: ChromaClientOptions,
+    ) -> ChromaClient {
         let endpoint = if url.is_empty() {
             std::env::var("CHROMA_URL").unwrap_or(DEFAULT_ENDPOINT.to_string())
         } else {
             url
         };
 
         ChromaClient {
-            api: Arc::new(APIClientAsync::new(endpoint, auth)),
+            api: Arc::new(APIClientAsync::new(endpoint, auth, database)),
         }
     }
 
@@ -165,7 +175,6 @@ mod tests {
     use super::*;
     use tokio;
 
-
     const TEST_COLLECTION: &str = "8-recipies-for-octopus";
 
     #[tokio::test]
@@ -208,7 +217,7 @@ mod tests {
     #[tokio::test]
     async fn test_get_collection() {
         let client: ChromaClient = ChromaClient::new(Default::default());
-        
+
         const GET_TEST_COLLECTION: &str = "100-recipes-for-octopus";
 
         client