[ENH]: (Rust client): export more types, config -> schema on create_collection() (#5699)

Author: codetheweb

clients/new-js/packages/chromadb/src/api/types.gen.ts

@@ -284,6 +284,7 @@ export type RemoveTaskResponse = {
 
 /**
  * Schema representation for collection index configurations
+ *
  * This represents the server-side schema structure used for index management
  */
 export type Schema = {

rust/chroma/src/client/chroma_http_client.rs

@@ -3,8 +3,8 @@ use backon::Retryable;
 use chroma_api_types::ErrorResponse;
 use chroma_error::ChromaValidationError;
 use chroma_types::Collection;
-use chroma_types::CollectionConfiguration;
 use chroma_types::Metadata;
+use chroma_types::Schema;
 use chroma_types::WhereError;
 use parking_lot::Mutex;
 use reqwest::Method;
@@ -437,10 +437,10 @@ impl ChromaHttpClient {
     pub async fn get_or_create_collection(
         &self,
         name: impl AsRef<str>,
-        configuration: Option<CollectionConfiguration>,
+        schema: Option<Schema>,
         metadata: Option<Metadata>,
     ) -> Result<ChromaCollection, ChromaHttpClientError> {
-        self.common_create_collection(name, configuration, metadata, true)
+        self.common_create_collection(name, schema, metadata, true)
             .await
     }
 
@@ -473,10 +473,10 @@ impl ChromaHttpClient {
     pub async fn create_collection(
         &self,
         name: impl AsRef<str>,
-        configuration: Option<CollectionConfiguration>,
+        schema: Option<Schema>,
         metadata: Option<Metadata>,
     ) -> Result<ChromaCollection, ChromaHttpClientError> {
-        self.common_create_collection(name, configuration, metadata, false)
+        self.common_create_collection(name, schema, metadata, false)
             .await
     }
 
@@ -620,7 +620,7 @@ impl ChromaHttpClient {
     async fn common_create_collection(
         &self,
         name: impl AsRef<str>,
-        configuration: Option<CollectionConfiguration>,
+        schema: Option<Schema>,
         metadata: Option<Metadata>,
         get_or_create: bool,
     ) -> Result<ChromaCollection, ChromaHttpClientError> {
@@ -637,7 +637,7 @@ impl ChromaHttpClient {
                 ),
                 Some(serde_json::json!({
                     "name": name.as_ref(),
-                    "configuration": configuration,
+                    "schema": schema,
                     "metadata": metadata,
                     "get_or_create": get_or_create,
                 })),
@@ -877,6 +877,7 @@ mod tests {
     use super::*;
     use crate::client::ChromaRetryOptions;
     use crate::tests::with_client;
+    use chroma_types::{EmbeddingFunctionConfiguration, EmbeddingFunctionNewConfiguration};
     use httpmock::{HttpMockResponse, MockServer};
     use std::sync::atomic::AtomicBool;
     use std::time::Duration;
@@ -1059,13 +1060,23 @@ mod tests {
     #[test_log::test]
     async fn test_live_cloud_create_collection() {
         with_client(|client| async move {
-            let collection = client.create_collection("foo", None, None).await.unwrap();
+            let schema = Schema::default_with_embedding_function(
+                EmbeddingFunctionConfiguration::Known(EmbeddingFunctionNewConfiguration {
+                    name: "bar".to_string(),
+                    config: serde_json::json!({}),
+                }),
+            );
+            let collection = client
+                .create_collection("foo", Some(schema.clone()), None)
+                .await
+                .unwrap();
             assert_eq!(collection.collection.name, "foo");
 
-            client
+            let collection = client
                 .get_or_create_collection("foo", None, None)
                 .await
                 .unwrap();
+            assert_eq!(collection.schema().clone().unwrap(), schema);
         })
         .await;
     }

rust/chroma/src/lib.rs

@@ -101,12 +101,13 @@
 #![deny(missing_docs)]
 
 pub mod client;
-pub mod collection;
+mod collection;
 pub mod embed;
 pub mod types;
 
 pub use client::ChromaHttpClient;
 pub use client::ChromaHttpClientOptions;
+pub use collection::ChromaCollection;
 
 #[cfg(test)]
 mod tests {

rust/chroma/src/types.rs

@@ -3,9 +3,15 @@
 pub use chroma_api_types::{GetUserIdentityResponse, HeartbeatResponse};
 
 pub use chroma_types::{
-    plan::SearchPayload, AddCollectionRecordsRequest, AddCollectionRecordsResponse, Collection,
-    DeleteCollectionRecordsRequest, DeleteCollectionRecordsResponse, ForkCollectionRequest,
-    GetRequest, GetResponse, IncludeList, Metadata, QueryRequest, QueryResponse, Schema,
-    SearchRequest, SearchResponse, UpdateCollectionRecordsRequest, UpdateCollectionRecordsResponse,
-    UpdateMetadata, UpsertCollectionRecordsRequest, UpsertCollectionRecordsResponse, Where,
+    plan::SearchPayload, AddCollectionRecordsRequest, AddCollectionRecordsResponse,
+    BooleanOperator, Collection, CompositeExpression, DeleteCollectionRecordsRequest,
+    DeleteCollectionRecordsResponse, DocumentExpression, DocumentOperator,
+    EmbeddingFunctionConfiguration, EmbeddingFunctionNewConfiguration, ForkCollectionRequest,
+    GetRequest, GetResponse, Include, IncludeList, Metadata, MetadataComparison,
+    MetadataExpression, MetadataSetValue, MetadataValue, PrimitiveOperator, QueryRequest,
+    QueryResponse, Schema, SearchRequest, SearchResponse, SetOperator,
+    UpdateCollectionRecordsRequest, UpdateCollectionRecordsResponse, UpdateMetadata,
+    UpdateMetadataValue, UpsertCollectionRecordsRequest, UpsertCollectionRecordsResponse, Where,
 };
+
+pub use chroma_types::operator::Key;

rust/types/src/collection_schema.rs

@@ -88,6 +88,7 @@ pub const EMBEDDING_KEY: &str = "#embedding";
 // ============================================================================
 
 /// Schema representation for collection index configurations
+///
 /// This represents the server-side schema structure used for index management
 
 #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
@@ -1114,6 +1115,25 @@ impl Schema {
         }
     }
 
+    pub fn default_with_embedding_function(
+        embedding_function: EmbeddingFunctionConfiguration,
+    ) -> Schema {
+        let mut schema = Schema::new_default(KnnIndex::Spann);
+        if let Some(float_list) = &mut schema.defaults.float_list {
+            if let Some(vector_index) = &mut float_list.vector_index {
+                vector_index.config.embedding_function = Some(embedding_function.clone());
+            }
+        }
+        if let Some(embedding_types) = schema.keys.get_mut(EMBEDDING_KEY) {
+            if let Some(float_list) = &mut embedding_types.float_list {
+                if let Some(vector_index) = &mut float_list.vector_index {
+                    vector_index.config.embedding_function = Some(embedding_function);
+                }
+            }
+        }
+        schema
+    }
+
     /// Check if schema is default by comparing it word-by-word with new_default
     fn is_schema_default(schema: &Schema) -> bool {
         // Compare with both possible default schemas (HNSW and SPANN)