[ENH] Implement collection apis for ChromaClient (#5653)

## Description of changes

_Summarize the changes made by this PR._

- Improvements & Bug fixes
    - N/A
- New functionality
    - Implement collection level apis

## Test plan

_How are these changes tested?_

- [ ] Tests pass locally with `pytest` for python, `yarn test` for js,
`cargo test` for rust

## Migration plan

_Are there any migrations, or any forwards/backwards compatibility
changes needed in order to make sure this change deploys reliably?_

## Observability plan

_What is the plan to instrument and monitor this change?_

## Documentation Changes

_Are all docstrings for user-facing APIs updated if required? Do we need
to make documentation changes in the_ [_docs
section](https://github.com/chroma-core/chroma/tree/main/docs/docs.trychroma.com)?_

Author: Sicheng-Pan

Cargo.lock

@@ -19,6 +19,7 @@ tracing.workspace = true
 parking_lot.workspace = true
 
 chroma-api-types = { workspace = true }
+chroma-error = { workspace = true }
 chroma-types = { workspace = true }
 
 [features]

rust/chroma/Cargo.toml

@@ -1,5 +1,6 @@
 use backon::ExponentialBuilder;
 use backon::Retryable;
+use chroma_error::ChromaValidationError;
 use parking_lot::Mutex;
 use reqwest::Method;
 use reqwest::StatusCode;
@@ -24,6 +25,8 @@ pub enum ChromaClientError {
     CouldNotResolveDatabaseId(String),
     #[error("Serialization/Deserialization error: {0}")]
     SerdeError(#[from] serde_json::Error),
+    #[error("Validation error: {0}")]
+    ValidationError(#[from] ChromaValidationError),
 }
 
 #[derive(Debug)]

rust/chroma/src/client/chroma_client.rs

@@ -1,6 +1,12 @@
 use std::sync::Arc;
 
-use chroma_types::{Collection, InternalSchema, Metadata};
+use chroma_types::{
+    plan::SearchPayload, AddCollectionRecordsRequest, AddCollectionRecordsResponse, Collection,
+    DeleteCollectionRecordsRequest, DeleteCollectionRecordsResponse, GetRequest, GetResponse,
+    IncludeList, InternalSchema, Metadata, QueryRequest, QueryResponse, SearchRequest,
+    SearchResponse, UpdateCollectionRecordsRequest, UpdateCollectionRecordsResponse,
+    UpdateMetadata, UpsertCollectionRecordsRequest, UpsertCollectionRecordsResponse, Where,
+};
 use reqwest::Method;
 use serde::{de::DeserializeOwned, Serialize};
 
@@ -33,6 +39,146 @@ impl ChromaCollection {
         self.send::<(), u32>("count", Method::GET, None).await
     }
 
+    pub async fn get(
+        &self,
+        ids: Option<Vec<String>>,
+        r#where: Option<Where>,
+        limit: Option<u32>,
+        offset: Option<u32>,
+        include: Option<IncludeList>,
+    ) -> Result<GetResponse, ChromaClientError> {
+        let request = GetRequest::try_new(
+            self.collection.tenant.clone(),
+            self.collection.database.clone(),
+            self.collection.collection_id,
+            ids,
+            r#where,
+            limit,
+            offset.unwrap_or_default(),
+            include.unwrap_or_else(IncludeList::default_get),
+        )?;
+
+        self.send("get", Method::POST, Some(request)).await
+    }
+
+    pub async fn query(
+        &self,
+        query_embeddings: Vec<Vec<f32>>,
+        n_results: Option<u32>,
+        r#where: Option<Where>,
+        ids: Option<Vec<String>>,
+        include: Option<IncludeList>,
+    ) -> Result<QueryResponse, ChromaClientError> {
+        let request = QueryRequest::try_new(
+            self.collection.tenant.clone(),
+            self.collection.database.clone(),
+            self.collection.collection_id,
+            ids,
+            r#where,
+            query_embeddings,
+            n_results.unwrap_or(10),
+            include.unwrap_or_else(IncludeList::default_query),
+        )?;
+
+        self.send("query", Method::POST, Some(request)).await
+    }
+
+    pub async fn search(
+        &self,
+        searches: Vec<SearchPayload>,
+    ) -> Result<SearchResponse, ChromaClientError> {
+        let request = SearchRequest::try_new(
+            self.collection.tenant.clone(),
+            self.collection.database.clone(),
+            self.collection.collection_id,
+            searches,
+        )?;
+
+        self.send("search", Method::POST, Some(request)).await
+    }
+
+    pub async fn add(
+        &self,
+        ids: Vec<String>,
+        embeddings: Vec<Vec<f32>>,
+        documents: Option<Vec<Option<String>>>,
+        uris: Option<Vec<Option<String>>>,
+        metadatas: Option<Vec<Option<Metadata>>>,
+    ) -> Result<AddCollectionRecordsResponse, ChromaClientError> {
+        let request = AddCollectionRecordsRequest::try_new(
+            self.collection.tenant.clone(),
+            self.collection.database.clone(),
+            self.collection.collection_id,
+            ids,
+            embeddings,
+            documents,
+            uris,
+            metadatas,
+        )?;
+
+        self.send("add", Method::POST, Some(request)).await
+    }
+
+    pub async fn update(
+        &self,
+        ids: Vec<String>,
+        embeddings: Option<Vec<Option<Vec<f32>>>>,
+        documents: Option<Vec<Option<String>>>,
+        uris: Option<Vec<Option<String>>>,
+        metadatas: Option<Vec<Option<UpdateMetadata>>>,
+    ) -> Result<UpdateCollectionRecordsResponse, ChromaClientError> {
+        let request = UpdateCollectionRecordsRequest::try_new(
+            self.collection.tenant.clone(),
+            self.collection.database.clone(),
+            self.collection.collection_id,
+            ids,
+            embeddings,
+            documents,
+            uris,
+            metadatas,
+        )?;
+
+        self.send("update", Method::POST, Some(request)).await
+    }
+
+    pub async fn upsert(
+        &self,
+        ids: Vec<String>,
+        embeddings: Vec<Vec<f32>>,
+        documents: Option<Vec<Option<String>>>,
+        uris: Option<Vec<Option<String>>>,
+        metadatas: Option<Vec<Option<UpdateMetadata>>>,
+    ) -> Result<UpsertCollectionRecordsResponse, ChromaClientError> {
+        let request = UpsertCollectionRecordsRequest::try_new(
+            self.collection.tenant.clone(),
+            self.collection.database.clone(),
+            self.collection.collection_id,
+            ids,
+            embeddings,
+            documents,
+            uris,
+            metadatas,
+        )?;
+
+        self.send("upsert", Method::POST, Some(request)).await
+    }
+
+    pub async fn delete(
+        &self,
+        ids: Option<Vec<String>>,
+        r#where: Option<Where>,
+    ) -> Result<DeleteCollectionRecordsResponse, ChromaClientError> {
+        let request = DeleteCollectionRecordsRequest::try_new(
+            self.collection.tenant.clone(),
+            self.collection.database.clone(),
+            self.collection.collection_id,
+            ids,
+            r#where,
+        )?;
+
+        self.send("delete", Method::POST, Some(request)).await
+    }
+
     async fn send<Body: Serialize, Response: DeserializeOwned>(
         &self,
         operation: &str,

rust/chroma/src/collection.rs

@@ -1202,7 +1202,7 @@ impl UpdateCollectionRecordsRequest {
     }
 }
 
-#[derive(Serialize, ToSchema)]
+#[derive(Serialize, Deserialize, ToSchema)]
 pub struct UpdateCollectionRecordsResponse {}
 
 #[derive(Error, Debug)]
@@ -1266,7 +1266,7 @@ impl UpsertCollectionRecordsRequest {
     }
 }
 
-#[derive(Serialize, ToSchema)]
+#[derive(Serialize, Deserialize, ToSchema)]
 pub struct UpsertCollectionRecordsResponse {}
 
 #[derive(Error, Debug)]
@@ -1326,7 +1326,7 @@ impl DeleteCollectionRecordsRequest {
     }
 }
 
-#[derive(Serialize, ToSchema)]
+#[derive(Serialize, Deserialize, ToSchema)]
 pub struct DeleteCollectionRecordsResponse {}
 
 #[derive(Error, Debug)]