[ENH] Use an improved chroma cloud client for chroma-load

rust/load/Cargo.toml

@@ -24,7 +24,7 @@ opentelemetry_sdk = { workspace = true }
 
 # Unlikely to be used in the workspace.
 axum = "0.7"
-chromadb = { git = "https://github.com/rescrv/chromadb-rs", rev = "e364e35c34c660d4e8e862436ea600ddc2f46a1e" }
+chromadb = { git = "https://github.com/rescrv/chromadb-rs", rev = "e9a8fb2e8b8edf7acfb1accf10166720a6bbbd33" }
 guacamole = { version = "0.9", default-features = false }
 tower-http = { version = "0.6.2", features = ["trace"] }
 reqwest = { version = "0.12", features = ["json"] }

rust/load/src/bit_difference.rs

@@ -25,8 +25,8 @@
 //! Internally, guacamole provides primitives that make it easy to manage the set of seeds to get a
 //! variety of data sets out of the synthetic data.
 
-use chromadb::v2::collection::{CollectionEntries, GetOptions, QueryOptions};
-use chromadb::v2::ChromaClient;
+use chromadb::collection::{CollectionEntries, GetOptions, QueryOptions};
+use chromadb::ChromaClient;
 use guacamole::combinators::*;
 use guacamole::{FromGuacamole, Guacamole, Zipf};
 use siphasher::sip::SipHasher24;

rust/load/src/data_sets.rs

@@ -1,7 +1,7 @@
 use std::sync::Arc;
 
-use chromadb::v2::collection::{GetOptions, QueryOptions};
-use chromadb::v2::ChromaClient;
+use chromadb::collection::{GetOptions, QueryOptions};
+use chromadb::ChromaClient;
 use guacamole::combinators::*;
 use guacamole::Guacamole;
 use tracing::Instrument;

rust/load/src/lib.rs

@@ -24,8 +24,8 @@ use axum::http::StatusCode;
 use axum::response::IntoResponse;
 use axum::routing::{get, post};
 use axum::{Json, Router};
-use chromadb::v2::client::{ChromaAuthMethod, ChromaClientOptions, ChromaTokenHeader};
-use chromadb::v2::ChromaClient;
+use chromadb::client::{ChromaAuthMethod, ChromaClientOptions, ChromaTokenHeader};
+use chromadb::ChromaClient;
 use guacamole::combinators::*;
 use guacamole::{Guacamole, Zipf};
 use opentelemetry::global;
@@ -113,29 +113,33 @@ pub struct Metrics {
 
 /// Instantiate a new Chroma client.  This will use the CHROMA_HOST environment variable (or
 /// http://localhost:8000 when unset) as the argument to [client_for_url].
-pub fn client() -> ChromaClient {
+pub async fn client() -> ChromaClient {
     let url = std::env::var("CHROMA_HOST").unwrap_or_else(|_| "http://localhost:8000".into());
-    client_for_url(url)
+    client_for_url(url).await
 }
 
 /// Create a new Chroma client for the given URL.  This will use the CHROMA_TOKEN environment
 /// variable if set, or no authentication if unset.
-pub fn client_for_url(url: String) -> ChromaClient {
+pub async fn client_for_url(url: String) -> ChromaClient {
     if let Ok(auth) = std::env::var("CHROMA_TOKEN") {
         ChromaClient::new(ChromaClientOptions {
-            url,
+            url: Some(url),
             auth: ChromaAuthMethod::TokenAuth {
                 token: auth,
                 header: ChromaTokenHeader::XChromaToken,
             },
-            database: Some("hf-tiny-stories".to_string()),
+            database: "hf-tiny-stories".to_string(),
         })
+        .await
+        .unwrap()
     } else {
         ChromaClient::new(ChromaClientOptions {
-            url,
+            url: Some(url),
             auth: ChromaAuthMethod::None,
-            database: Some("hf-tiny-stories".to_string()),
+            database: "hf-tiny-stories".to_string(),
         })
+        .await
+        .unwrap()
     }
 }
 
@@ -1107,7 +1111,7 @@ impl LoadService {
         inhibit: Arc<AtomicBool>,
         spec: RunningWorkload,
     ) {
-        let client = Arc::new(client());
+        let client = Arc::new(client().await);
         let mut guac = Guacamole::new(spec.expires.timestamp_millis() as u64);
         let mut next_op = Instant::now();
         let (tx, mut rx) = tokio::sync::mpsc::channel(1000);