chroma 0.15.0

//! Collection operations for managing and querying vector embeddings.
//!
//! This module provides the [`ChromaCollection`] type, which represents a handle to a specific
//! collection within a Chroma database. Collections are the primary storage unit for vector
//! embeddings and their associated metadata, documents, and URIs.
//!
//! # Operations
//!
//! Collections support the following categories of operations:
//!
//! - **Metadata access**: [`database()`](ChromaCollection::database), [`metadata()`](ChromaCollection::metadata), [`schema()`](ChromaCollection::schema), [`tenant()`](ChromaCollection::tenant)
//! - **Read operations**: [`count()`](ChromaCollection::count), [`get()`](ChromaCollection::get), [`get_indexing_status()`](ChromaCollection::get_indexing_status), [`query()`](ChromaCollection::query), [`search()`](ChromaCollection::search)
//! - **Write operations**: [`add()`](ChromaCollection::add), [`update()`](ChromaCollection::update), [`upsert()`](ChromaCollection::upsert), [`delete()`](ChromaCollection::delete), [`modify()`](ChromaCollection::modify)
//! - **Attached functions**: [`attach_function()`](ChromaCollection::attach_function), [`get_attached_function()`](ChromaCollection::get_attached_function), [`detach_function()`](ChromaCollection::detach_function)

use std::{error::Error, fmt::Display, sync::Arc};

use chroma_api_types::ForkCollectionPayload;
use chroma_types::{
    plan::{ReadLevel, SearchPayload},
    AddCollectionRecordsRequest, AddCollectionRecordsResponse, AttachFunctionResponse, Collection,
    CollectionUuid, DeleteCollectionRecordsRequest, DeleteCollectionRecordsResponse,
    DetachFunctionResponse, EmbeddingFunctionConfiguration, GetRequest, GetResponse, IncludeList,
    IndexStatusResponse, Metadata, QueryRequest, QueryResponse, Schema, SearchRequest,
    SearchResponse, UpdateCollectionRecordsRequest, UpdateCollectionRecordsResponse,
    UpdateMetadata, UpsertCollectionRecordsRequest, UpsertCollectionRecordsResponse, Where,
    EMBEDDING_KEY,
};
use reqwest::Method;
use serde::{de::DeserializeOwned, Deserialize, Serialize};

use crate::embed::{chroma_cloud::ChromaCloudQwenEmbeddingFunction, EmbeddingFunction};
use crate::{client::ChromaHttpClientError, ChromaAttachedFunction, ChromaHttpClient};

#[derive(Deserialize)]
struct ForkCountResponse {
    count: usize,
}

/// A handle to a specific collection within a Chroma database.
///
/// Collections are the primary storage unit in Chroma for vector embeddings and associated metadata.
/// Each collection belongs to a specific tenant and database, and contains records that can be
/// queried by similarity, filtered by metadata, and managed through CRUD operations.
///
/// # Architecture
///
/// A `ChromaCollection` is a lightweight reference to collection metadata that shares the underlying
/// HTTP client. Operations on the collection are executed immediately against the server.
///
/// # Examples
///
/// ```
/// # use chroma::ChromaCollection;
/// # use chroma::client::ChromaHttpClientError;
/// # async fn example(collection: ChromaCollection) -> Result<(), ChromaHttpClientError> {
/// let count = collection.count().await?;
/// println!("Collection contains {} records", count);
///
/// let embeddings = vec![vec![0.1, 0.2, 0.3]];
/// let results = collection.query(embeddings, Some(10), None, None, None).await?;
/// println!("Found {} similar records", results.ids.len());
/// # Ok(())
/// # }
/// ```
#[derive(Clone)]
pub struct ChromaCollection {
    pub(crate) client: ChromaHttpClient,
    pub(crate) collection: Arc<Collection>,
    pub(crate) embedding_function: Option<Arc<ErasedEmbeddingFunction>>,
}

type ErasedEmbeddingFunction =
    dyn EmbeddingFunction<Embedding = Vec<f32>, Error = BoxedEmbeddingError>;

#[derive(Debug)]
pub(crate) struct BoxedEmbeddingError(Box<dyn Error + Send + Sync>);

impl Display for BoxedEmbeddingError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        Display::fmt(&self.0, f)
    }
}

impl Error for BoxedEmbeddingError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        Some(self.0.as_ref())
    }
}

struct ErrorErasedEmbeddingFunction<E> {
    inner: E,
}

#[async_trait::async_trait]
impl<E> EmbeddingFunction for ErrorErasedEmbeddingFunction<E>
where
    E: EmbeddingFunction<Embedding = Vec<f32>>,
    E::Error: Send + Sync + 'static,
{
    type Embedding = Vec<f32>;
    type Error = BoxedEmbeddingError;

    async fn embed_strs(&self, batches: &[&str]) -> Result<Vec<Self::Embedding>, Self::Error> {
        self.inner
            .embed_strs(batches)
            .await
            .map_err(|err| BoxedEmbeddingError(Box::new(err)))
    }

    async fn embed_query_strs(
        &self,
        batches: &[&str],
    ) -> Result<Vec<Self::Embedding>, Self::Error> {
        self.inner
            .embed_query_strs(batches)
            .await
            .map_err(|err| BoxedEmbeddingError(Box::new(err)))
    }
}

fn erase_embedding_function<E>(embedding_function: E) -> Arc<ErasedEmbeddingFunction>
where
    E: EmbeddingFunction<Embedding = Vec<f32>>,
    E::Error: Send + Sync + 'static,
{
    Arc::new(ErrorErasedEmbeddingFunction {
        inner: embedding_function,
    }) as Arc<ErasedEmbeddingFunction>
}

fn default_embedding_function(
    client: &ChromaHttpClient,
    collection: &Collection,
) -> Option<Arc<ErasedEmbeddingFunction>> {
    let config = dense_embedding_function_configuration(collection)?;
    match config {
        EmbeddingFunctionConfiguration::Known(config)
            if config.name == ChromaCloudQwenEmbeddingFunction::name() =>
        {
            ChromaCloudQwenEmbeddingFunction::try_from_config(config, client.chroma_cloud_api_key())
                .map(erase_embedding_function)
                .ok()
        }
        EmbeddingFunctionConfiguration::Legacy
        | EmbeddingFunctionConfiguration::Known(_)
        | EmbeddingFunctionConfiguration::Unknown => None,
    }
}

fn dense_embedding_function_configuration(
    collection: &Collection,
) -> Option<&EmbeddingFunctionConfiguration> {
    let schema = collection.schema.as_ref()?;
    schema
        .keys
        .get(EMBEDDING_KEY)
        .and_then(|value_types| value_types.float_list.as_ref())
        .and_then(|float_list| float_list.vector_index.as_ref())
        .and_then(|vector_index| vector_index.config.embedding_function.as_ref())
        .or_else(|| {
            schema
                .defaults
                .float_list
                .as_ref()
                .and_then(|float_list| float_list.vector_index.as_ref())
                .and_then(|vector_index| vector_index.config.embedding_function.as_ref())
        })
}

/// Converts an embeddings argument into the optional dense embedding form used by writes.
///
/// This keeps the write methods to a single public method while accepting either explicit
/// embeddings or `None` to request embedding-function fallback.
pub trait IntoOptionalEmbeddings {
    /// Converts this value into optional dense embeddings.
    fn into_optional_embeddings(self) -> Option<Vec<Vec<f32>>>;
}

impl IntoOptionalEmbeddings for Vec<Vec<f32>> {
    fn into_optional_embeddings(self) -> Option<Vec<Vec<f32>>> {
        Some(self)
    }
}

impl IntoOptionalEmbeddings for Option<Vec<Vec<f32>>> {
    fn into_optional_embeddings(self) -> Option<Vec<Vec<f32>>> {
        self
    }
}

impl std::fmt::Debug for ChromaCollection {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ChromaCollection")
            .field("database", &self.collection.database)
            .field("tenant", &self.collection.tenant)
            .field("name", &self.collection.name)
            .field("collection_id", &self.collection.collection_id)
            .field("version", &self.collection.version)
            .finish()
    }
}

impl ChromaCollection {
    pub(crate) fn new(client: ChromaHttpClient, collection: Collection) -> Self {
        let embedding_function = default_embedding_function(&client, &collection);
        Self {
            client,
            collection: Arc::new(collection),
            embedding_function,
        }
    }

    /// Builds a collection handle from an already-resolved [`Collection`]
    /// model without performing a network lookup.
    ///
    /// This is useful when the collection identity (id, schema, metadata) has
    /// been cached out of band and only needs to be re-bound to a fresh client
    /// (for example, one carrying a per-request authentication token). The
    /// dense embedding function is re-derived from the collection schema,
    /// mirroring [`ChromaHttpClient::get_collection`](crate::ChromaHttpClient::get_collection).
    pub fn from_collection_model(client: ChromaHttpClient, collection: Collection) -> Self {
        Self::new(client, collection)
    }

    /// Returns a clone of the underlying [`Collection`] model (id, name,
    /// schema, metadata, ...).
    ///
    /// Useful for caching the collection identity so a handle can later be
    /// rebuilt with [`ChromaCollection::from_collection_model`] without another
    /// network round trip.
    pub fn to_collection_model(&self) -> Collection {
        (*self.collection).clone()
    }

    /// Sets the embedding function used when record embeddings are omitted.
    ///
    /// Passing `None` clears the callback. When set, [`add`](Self::add),
    /// [`update`](Self::update), and [`upsert`](Self::upsert) can compute dense
    /// embeddings from provided documents.
    pub fn set_embedding_function<E>(&mut self, embedding_function: Option<E>)
    where
        E: EmbeddingFunction<Embedding = Vec<f32>>,
        E::Error: Send + Sync + 'static,
    {
        self.embedding_function = embedding_function.map(erase_embedding_function);
    }

    /// Returns the database ID that contains this collection.
    pub fn database(&self) -> &str {
        &self.collection.database
    }

    /// Returns the collection's metadata, if any was specified during creation.
    pub fn metadata(&self) -> &Option<Metadata> {
        &self.collection.metadata
    }

    /// Returns the collection's internal schema definition, if available.
    pub fn schema(&self) -> &Option<Schema> {
        &self.collection.schema
    }

    /// Returns the tenant ID that owns this collection.
    pub fn tenant(&self) -> &str {
        &self.collection.tenant
    }

    /// Returns the user-assigned name of this collection.
    ///
    /// # Examples
    ///
    /// ```
    /// # use chroma::ChromaCollection;
    /// # fn example(collection: ChromaCollection) {
    /// assert_eq!(collection.name(), "my_collection");
    /// # }
    /// ```
    pub fn name(&self) -> &str {
        &self.collection.name
    }

    /// Returns the unique identifier assigned to this collection by Chroma.
    ///
    /// # Examples
    ///
    /// ```
    /// # use chroma::ChromaCollection;
    /// # fn example(collection: ChromaCollection) {
    /// let id = collection.id();
    /// println!("Collection ID: {}", id);
    /// # }
    /// ```
    pub fn id(&self) -> CollectionUuid {
        self.collection.collection_id
    }

    /// Gets an attached function by name for this collection.
    ///
    /// Returns a handle scoped to this collection's participation in the attached
    /// function. For multi-input async attached functions, the returned handle can
    /// be used to add more input collections via [`ChromaAttachedFunction::add_input`].
    pub async fn get_attached_function(
        &self,
        name: impl AsRef<str>,
    ) -> Result<ChromaAttachedFunction, ChromaHttpClientError> {
        self.client
            .get_attached_function(self.collection.collection_id, name)
            .await
    }

    /// Returns the version number of this collection's schema or configuration.
    ///
    /// The version increments when the collection's structure or settings change.
    ///
    /// # Examples
    ///
    /// ```
    /// # use chroma::ChromaCollection;
    /// # fn example(collection: ChromaCollection) {
    /// let version = collection.version();
    /// assert!(version >= 0);
    /// # }
    /// ```
    pub fn version(&self) -> i32 {
        self.collection.version
    }

    /// Computes the total number of records stored in this collection.
    ///
    /// # Errors
    ///
    /// Returns an error if network communication fails.
    ///
    /// # Examples
    ///
    /// ```
    /// # use chroma::ChromaCollection;
    /// # async fn example(collection: ChromaCollection) -> Result<(), Box<dyn std::error::Error>> {
    /// let count = collection.count().await?;
    /// assert!(count >= 0);
    /// # Ok(())
    /// # }
    /// ```
    pub async fn count(&self) -> Result<u32, ChromaHttpClientError> {
        self.count_with_options(ReadLevel::IndexAndWal).await
    }

    /// Count with custom read level controlling whether to read from the write-ahead log.
    ///
    /// By default, counts read from both the compacted index and the write-ahead log (WAL),
    /// ensuring all committed writes are visible. For higher throughput at the cost of
    /// potentially missing recent uncommitted writes, use `ReadLevel::IndexOnly` to skip
    /// the WAL and read only from the compacted index.
    ///
    /// # Arguments
    ///
    /// * `read_level` - Controls data sources for the query:
    ///   - [`ReadLevel::IndexAndWal`] - Read from both the compacted index and WAL (default).
    ///     All committed writes will be visible.
    ///   - [`ReadLevel::IndexOnly`] - Read only from the compacted index, skipping the WAL.
    ///     Faster, but recent writes that haven't been compacted may not be visible.
    ///   - [`ReadLevel::IndexAndBoundedWal`] - Read from the index and up to a server-configured
    ///     number of WAL entries for bounded query latency.
    ///
    /// # Example
    ///
    /// ```
    /// use chroma::types::ReadLevel;
    ///
    /// # async fn example(collection: &chroma::ChromaCollection) -> Result<(), Box<dyn std::error::Error>> {
    /// // Skip WAL for faster count (may miss recent writes)
    /// let count = collection.count_with_options(ReadLevel::IndexOnly).await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn count_with_options(
        &self,
        read_level: ReadLevel,
    ) -> Result<u32, ChromaHttpClientError> {
        #[derive(Serialize)]
        struct CountQueryParams {
            read_level: ReadLevel,
        }
        self.send_with_query::<(), CountQueryParams, u32>(
            true,
            "count",
            "count",
            Method::GET,
            None,
            Some(CountQueryParams { read_level }),
        )
        .await
    }

    /// Gets the indexing status of this collection.
    ///
    /// Returns information about how many operations have been indexed versus how many are
    /// pending. This is useful for monitoring the progress of bulk data ingestion.
    ///
    /// # Errors
    ///
    /// Returns an error if network communication fails, or if unauthenticated.
    ///
    /// # Examples
    ///
    /// ```
    /// # use chroma::ChromaCollection;
    /// # async fn example(collection: ChromaCollection) -> Result<(), Box<dyn std::error::Error>> {
    /// let status = collection.get_indexing_status().await?;
    /// println!("Indexing progress: {:.1}%", status.op_indexing_progress * 100.0);
    /// println!("Indexed: {} / Total: {}", status.num_indexed_ops, status.total_ops);
    /// # Ok(())
    /// # }
    /// ```
    pub async fn get_indexing_status(&self) -> Result<IndexStatusResponse, ChromaHttpClientError> {
        self.send::<(), IndexStatusResponse>(
            true,
            "indexing_status",
            "indexing_status",
            Method::GET,
            None,
        )
        .await
    }

    /// Modifies the collection's name or metadata.
    ///
    /// Updates the collection's configuration on the server and synchronizes the local
    /// collection handle with the new values. At least one of `new_name` or `new_metadata`
    /// should be specified; passing `None` for both parameters is valid but has no effect.
    ///
    /// Note that this method takes a mutable reference to self because it updates the
    /// local collection metadata after the server-side modification succeeds.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - A collection with the new name already exists in the same database
    /// - Network communication fails
    /// - The authenticated user lacks sufficient permissions
    ///
    /// # Examples
    ///
    /// ```
    /// # use chroma::ChromaCollection;
    /// # use chroma_types::Metadata;
    /// # async fn example(mut collection: ChromaCollection) -> Result<(), Box<dyn std::error::Error>> {
    /// let mut metadata = Metadata::new();
    /// metadata.insert("version".to_string(), "2.0".into());
    ///
    /// collection.modify(
    ///     Some("my_renamed_collection"),
    ///     Some(metadata)
    /// ).await?;
    ///
    /// assert_eq!(collection.name(), "my_renamed_collection");
    /// # Ok(())
    /// # }
    /// ```
    pub async fn modify(
        &mut self,
        new_name: Option<impl AsRef<str>>,
        new_metadata: Option<Metadata>,
    ) -> Result<(), ChromaHttpClientError> {
        // Returns empty map ({})
        self.send::<_, serde_json::Value>(
            false,
            "modify",
            "",
            Method::PUT,
            Some(serde_json::json!({
                "new_name": new_name.as_ref().map(|s| s.as_ref()),
                "new_metadata": new_metadata,
            })),
        )
        .await?;

        let mut updated_collection = (*self.collection).clone();
        if let Some(name) = new_name {
            updated_collection.name = name.as_ref().to_string();
        }
        if let Some(metadata) = new_metadata {
            updated_collection.metadata = Some(metadata);
        }

        self.collection = Arc::new(updated_collection);

        Ok(())
    }

    /// Retrieves records from the collection by ID or metadata filter.
    ///
    /// At least one of `ids` or `where` must be specified. The `include` parameter controls
    /// which fields are returned (embeddings, documents, metadata, URIs).
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - Network communication fails
    /// - Request validation fails
    ///
    /// # Examples
    ///
    /// ```
    /// # use chroma::ChromaCollection;
    /// # use chroma_types::IncludeList;
    /// # async fn example(collection: ChromaCollection) -> Result<(), Box<dyn std::error::Error>> {
    /// let response = collection.get(
    ///     Some(vec!["id1".to_string(), "id2".to_string()]),
    ///     None,
    ///     Some(10),
    ///     Some(0),
    ///     Some(IncludeList::default_get())
    /// ).await?;
    /// println!("Retrieved {} records", response.ids.len());
    /// # Ok(())
    /// # }
    /// ```
    pub async fn get(
        &self,
        ids: Option<Vec<String>>,
        r#where: Option<Where>,
        limit: Option<u32>,
        offset: Option<u32>,
        include: Option<IncludeList>,
    ) -> Result<GetResponse, ChromaHttpClientError> {
        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),
        )?;
        let request = request.into_payload()?;
        self.send(true, "get", "get", Method::POST, Some(request))
            .await
    }

    /// Performs vector similarity search against the collection.
    ///
    /// Finds the `n_results` nearest neighbors for each query embedding using the collection's
    /// configured distance metric. Results can be filtered by metadata conditions via `where`
    /// or restricted to a subset of IDs.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - Embedding dimensions don't match the collection's schema
    /// - Network communication fails
    /// - Request validation fails
    ///
    /// # Examples
    ///
    /// ```
    /// # use chroma::ChromaCollection;
    /// # async fn example(collection: ChromaCollection) -> Result<(), Box<dyn std::error::Error>> {
    /// let query = vec![vec![0.1, 0.2, 0.3]];
    /// let results = collection.query(
    ///     query,
    ///     Some(5),
    ///     None,
    ///     None,
    ///     None
    /// ).await?;
    ///
    /// for (i, ids) in results.ids.iter().enumerate() {
    ///     println!("Query {} found {} neighbors", i, ids.len());
    /// }
    /// # Ok(())
    /// # }
    /// ```
    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, ChromaHttpClientError> {
        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),
        )?;
        let request = request.into_payload()?;
        self.send(true, "query", "query", Method::POST, Some(request))
            .await
    }

    /// Performs hybrid search on the collection using the Search API.
    ///
    /// The Search API provides a powerful, flexible interface for vector similarity search
    /// combined with metadata filtering and custom ranking expressions.
    ///
    /// # Arguments
    ///
    /// * `searches` - One or more search payloads to execute in a single request
    ///
    /// # Returns
    ///
    /// A `SearchResponse` containing results for each search payload
    ///
    /// # Examples
    ///
    /// ## Basic similarity search
    ///
    /// ```
    /// use chroma_types::plan::SearchPayload;
    /// use chroma_types::operator::{RankExpr, QueryVector, Key};
    ///
    /// # async fn example(collection: &chroma::ChromaCollection) -> Result<(), Box<dyn std::error::Error>> {
    /// // Search with a query vector
    /// let search = SearchPayload::default()
    ///     .rank(RankExpr::Knn {
    ///         query: QueryVector::Dense(vec![0.1, 0.2, 0.3]),
    ///         key: Key::Embedding,
    ///         limit: 100,
    ///         default: None,
    ///         return_rank: false,
    ///     })
    ///     .limit(Some(10), 0)
    ///     .select([Key::Document, Key::Score]);
    ///
    /// let results = collection.search(vec![search]).await?;
    /// # Ok(())
    /// # }
    /// ```
    ///
    /// ## Filtered search with metadata
    ///
    /// ```
    /// use chroma_types::plan::SearchPayload;
    /// use chroma_types::operator::{RankExpr, QueryVector, Key};
    ///
    /// # async fn example(collection: &chroma::ChromaCollection) -> Result<(), Box<dyn std::error::Error>> {
    /// // Filter by category and year, then rank by similarity
    /// let search = SearchPayload::default()
    ///     .r#where(
    ///         Key::field("category").eq("science")
    ///             & Key::field("year").gte(2020)
    ///     )
    ///     .rank(RankExpr::Knn {
    ///         query: QueryVector::Dense(vec![0.1, 0.2, 0.3]),
    ///         key: Key::Embedding,
    ///         limit: 200,
    ///         default: None,
    ///         return_rank: false,
    ///     })
    ///     .limit(Some(5), 0)
    ///     .select([Key::Document, Key::Score, Key::field("title")]);
    ///
    /// let results = collection.search(vec![search]).await?;
    /// # Ok(())
    /// # }
    /// ```
    ///
    /// ## Hybrid search with custom ranking
    ///
    /// ```
    /// use chroma_types::plan::SearchPayload;
    /// use chroma_types::operator::{RankExpr, QueryVector, Key};
    ///
    /// # async fn example(collection: &chroma::ChromaCollection) -> Result<(), Box<dyn std::error::Error>> {
    /// // Combine two KNN searches with weights
    /// let dense_knn = RankExpr::Knn {
    ///     query: QueryVector::Dense(vec![0.1, 0.2, 0.3]),
    ///     key: Key::Embedding,
    ///     limit: 200,
    ///     default: None,
    ///     return_rank: false,
    /// };
    ///
    /// let sparse_knn = RankExpr::Knn {
    ///     query: QueryVector::Dense(vec![0.1, 0.2, 0.3]), // Use sparse vector in practice
    ///     key: Key::field("sparse_embedding"),
    ///     limit: 200,
    ///     default: None,
    ///     return_rank: false,
    /// };
    ///
    /// // Weighted combination: 70% dense + 30% sparse
    /// let hybrid_rank = dense_knn * 0.7 + sparse_knn * 0.3;
    ///
    /// let search = SearchPayload::default()
    ///     .rank(hybrid_rank)
    ///     .limit(Some(10), 0)
    ///     .select([Key::Document, Key::Score]);
    ///
    /// let results = collection.search(vec![search]).await?;
    /// # Ok(())
    /// # }
    /// ```
    ///
    /// ## Batch operations
    ///
    /// ```
    /// use chroma_types::plan::SearchPayload;
    /// use chroma_types::operator::{RankExpr, QueryVector, Key};
    ///
    /// # async fn example(collection: &chroma::ChromaCollection) -> Result<(), Box<dyn std::error::Error>> {
    /// // Run multiple searches in one request
    /// let searches = vec![
    ///     SearchPayload::default()
    ///         .r#where(Key::field("category").eq("tech"))
    ///         .rank(RankExpr::Knn {
    ///             query: QueryVector::Dense(vec![0.1, 0.2, 0.3]),
    ///             key: Key::Embedding,
    ///             limit: 100,
    ///             default: None,
    ///             return_rank: false,
    ///         })
    ///         .limit(Some(5), 0),
    ///     SearchPayload::default()
    ///         .r#where(Key::field("category").eq("science"))
    ///         .rank(RankExpr::Knn {
    ///             query: QueryVector::Dense(vec![0.1, 0.2, 0.3]),
    ///             key: Key::Embedding,
    ///             limit: 100,
    ///             default: None,
    ///             return_rank: false,
    ///         })
    ///         .limit(Some(5), 0),
    /// ];
    ///
    /// let results = collection.search(searches).await?;
    /// // results.results[0] contains first search results
    /// // results.results[1] contains second search results
    /// # Ok(())
    /// # }
    /// ```
    pub async fn search(
        &self,
        searches: Vec<SearchPayload>,
    ) -> Result<SearchResponse, ChromaHttpClientError> {
        self.search_with_options(searches, ReadLevel::IndexAndWal)
            .await
    }

    /// Search with custom read level controlling whether to read from the write-ahead log.
    ///
    /// By default, searches read from both the compacted index and the write-ahead log (WAL),
    /// ensuring all committed writes are visible. For higher throughput at the cost of
    /// potentially missing recent uncommitted writes, use `ReadLevel::IndexOnly` to skip
    /// the WAL and read only from the compacted index.
    ///
    /// # Arguments
    ///
    /// * `searches` - Vector of search payloads to execute
    /// * `read_level` - Controls data sources for the query:
    ///   - [`ReadLevel::IndexAndWal`] - Read from both the compacted index and WAL (default).
    ///     All committed writes will be visible.
    ///   - [`ReadLevel::IndexOnly`] - Read only from the compacted index, skipping the WAL.
    ///     Faster, but recent writes that haven't been compacted may not be visible.
    ///   - [`ReadLevel::IndexAndBoundedWal`] - Read from the index and up to a server-configured
    ///     number of WAL entries for bounded query latency.
    ///
    /// # Example
    ///
    /// ```
    /// use chroma::types::{SearchPayload, ReadLevel};
    ///
    /// # async fn example(collection: &chroma::ChromaCollection) -> Result<(), Box<dyn std::error::Error>> {
    /// let search = SearchPayload::default().limit(Some(10), 0);
    ///
    /// // Skip WAL for faster queries (may miss recent writes)
    /// let results = collection
    ///     .search_with_options(vec![search], ReadLevel::IndexOnly)
    ///     .await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn search_with_options(
        &self,
        searches: Vec<SearchPayload>,
        read_level: ReadLevel,
    ) -> Result<SearchResponse, ChromaHttpClientError> {
        let request = SearchRequest::try_new(
            self.collection.tenant.clone(),
            self.collection.database.clone(),
            self.collection.collection_id,
            searches,
            read_level,
        )?;
        let request = request.into_payload();
        self.send(true, "search", "search", Method::POST, Some(request))
            .await
    }

    /// Inserts new records into the collection.
    ///
    /// All provided vectors must have lengths equal: `ids`, optional `embeddings`, and optionally
    /// `documents`, `uris`, and `metadatas`. Records with duplicate IDs will cause an error.
    /// If `embeddings` is `None`, the collection's embedding function computes embeddings from
    /// `documents`.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - Embedding dimensions don't match the collection's schema
    /// - Network communication fails
    /// - Request validation fails
    ///
    /// # Examples
    ///
    /// ```
    /// # use chroma::ChromaCollection;
    /// # async fn example(collection: ChromaCollection) -> Result<(), Box<dyn std::error::Error>> {
    /// let response = collection.add(
    ///     vec!["doc1".to_string(), "doc2".to_string()],
    ///     Some(vec![vec![0.1, 0.2], vec![0.3, 0.4]]),
    ///     Some(vec![Some("First document".to_string()), Some("Second document".to_string())]),
    ///     None,
    ///     None
    /// ).await?;
    /// println!("Added records successfully");
    /// # Ok(())
    /// # }
    /// ```
    pub async fn add(
        &self,
        ids: Vec<String>,
        embeddings: impl IntoOptionalEmbeddings,
        documents: Option<Vec<Option<String>>>,
        uris: Option<Vec<Option<String>>>,
        metadatas: Option<Vec<Option<Metadata>>>,
    ) -> Result<AddCollectionRecordsResponse, ChromaHttpClientError> {
        let embeddings = self
            .resolve_embeddings(embeddings.into_optional_embeddings(), &documents)
            .await?;
        let request = AddCollectionRecordsRequest::try_new(
            self.collection.tenant.clone(),
            self.collection.database.clone(),
            self.collection.collection_id,
            ids,
            embeddings,
            documents,
            uris,
            metadatas,
        )?;
        let request = request.into_payload();
        self.send(false, "add", "add", Method::POST, Some(request))
            .await
    }

    /// Modifies existing records in the collection.
    ///
    /// Updates only the specified fields for records matching the provided IDs. Fields set to
    /// `None` or `Some(None)` remain unchanged. All non-`None` vectors must match the length of
    /// `ids`. If `embeddings` is `None`, `documents` contains values, and an embedding function is
    /// configured, the collection's embedding function computes updated embeddings for those
    /// documents.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - Embedding dimensions don't match the collection's schema
    /// - Network communication fails
    /// - Request validation fails
    ///
    /// # Examples
    ///
    /// ```
    /// # use chroma::ChromaCollection;
    /// # async fn example(collection: ChromaCollection) -> Result<(), Box<dyn std::error::Error>> {
    /// let response = collection.update(
    ///     vec!["doc1".to_string()],
    ///     Some(vec![Some(vec![0.5, 0.6])]),
    ///     Some(vec![Some("Updated document text".to_string())]),
    ///     None,
    ///     None
    /// ).await?;
    /// println!("Updated records successfully");
    /// # Ok(())
    /// # }
    /// ```
    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, ChromaHttpClientError> {
        let embeddings = self
            .resolve_update_embeddings(embeddings, &documents)
            .await?;
        let request = UpdateCollectionRecordsRequest::try_new(
            self.collection.tenant.clone(),
            self.collection.database.clone(),
            self.collection.collection_id,
            ids,
            embeddings,
            documents,
            uris,
            metadatas,
        )?;
        let request = request.into_payload();
        self.send(false, "update", "update", Method::POST, Some(request))
            .await
    }

    /// Inserts new records or updates existing ones based on ID.
    ///
    /// For each ID: if the record exists, updates it; otherwise, inserts a new record.
    /// This combines the semantics of [`add`](Self::add) and [`update`](Self::update) in a single operation.
    /// If `embeddings` is `None`, the collection's embedding function computes embeddings from
    /// `documents`.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - Embedding dimensions don't match the collection's schema
    /// - Network communication fails
    /// - Request validation fails
    ///
    /// # Examples
    ///
    /// ```
    /// # use chroma::ChromaCollection;
    /// # async fn example(collection: ChromaCollection) -> Result<(), Box<dyn std::error::Error>> {
    /// let response = collection.upsert(
    ///     vec!["doc1".to_string(), "doc2".to_string()],
    ///     Some(vec![vec![0.1, 0.2], vec![0.3, 0.4]]),
    ///     Some(vec![Some("Document 1".to_string()), Some("Document 2".to_string())]),
    ///     None,
    ///     None
    /// ).await?;
    /// println!("Upserted records successfully");
    /// # Ok(())
    /// # }
    /// ```
    pub async fn upsert(
        &self,
        ids: Vec<String>,
        embeddings: impl IntoOptionalEmbeddings,
        documents: Option<Vec<Option<String>>>,
        uris: Option<Vec<Option<String>>>,
        metadatas: Option<Vec<Option<UpdateMetadata>>>,
    ) -> Result<UpsertCollectionRecordsResponse, ChromaHttpClientError> {
        let embeddings = self
            .resolve_embeddings(embeddings.into_optional_embeddings(), &documents)
            .await?;
        let request = UpsertCollectionRecordsRequest::try_new(
            self.collection.tenant.clone(),
            self.collection.database.clone(),
            self.collection.collection_id,
            ids,
            embeddings,
            documents,
            uris,
            metadatas,
        )?;
        let request = request.into_payload();
        self.send(false, "upsert", "upsert", Method::POST, Some(request))
            .await
    }

    /// Removes records from the collection by ID or metadata filter.
    ///
    /// At least one of `ids` or `where` must be specified to prevent accidental deletion
    /// of all records.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - Network communication fails
    /// - Request validation fails
    ///
    /// # Examples
    ///
    /// ```
    /// # use chroma::ChromaCollection;
    /// # async fn example(collection: ChromaCollection) -> Result<(), Box<dyn std::error::Error>> {
    /// let response = collection.delete(
    ///     Some(vec!["doc1".to_string(), "doc2".to_string()]),
    ///     None,
    ///     None,
    /// ).await?;
    /// println!("Deleted {} records", response.deleted);
    /// # Ok(())
    /// # }
    /// ```
    pub async fn delete(
        &self,
        ids: Option<Vec<String>>,
        r#where: Option<Where>,
        limit: Option<u32>,
    ) -> Result<DeleteCollectionRecordsResponse, ChromaHttpClientError> {
        let request = DeleteCollectionRecordsRequest::try_new(
            self.collection.tenant.clone(),
            self.collection.database.clone(),
            self.collection.collection_id,
            ids,
            r#where,
            limit,
        )?;
        let request = request.into_payload()?;
        self.send(false, "delete", "delete", Method::POST, Some(request))
            .await
    }

    /// Creates a shallow copy of this collection under a new name.
    ///
    /// The forked collection resides in the same database and tenant but receives a distinct
    /// collection ID and name. The fork shares the same configuration and initial records
    /// as the source collection at the time of forking, but subsequent modifications to either
    /// collection do not affect the other.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - A collection with the target name already exists
    /// - Network communication fails
    /// - The authenticated user lacks sufficient permissions
    ///
    /// # Examples
    ///
    /// ```
    /// # use chroma::ChromaCollection;
    /// # async fn example(collection: ChromaCollection) -> Result<(), Box<dyn std::error::Error>> {
    /// let forked = collection.fork("experimental_version").await?;
    /// assert_eq!(forked.name(), "experimental_version");
    /// # Ok(())
    /// # }
    /// ```
    pub async fn fork(
        &self,
        new_name: impl Into<String>,
    ) -> Result<ChromaCollection, ChromaHttpClientError> {
        let request = ForkCollectionPayload {
            new_name: new_name.into(),
        };
        let collection: Collection = self
            .send(false, "fork", "fork", Method::POST, Some(request))
            .await?;
        Ok(ChromaCollection {
            client: self.client.clone(),
            collection: Arc::new(collection),
            embedding_function: self.embedding_function.clone(),
        })
    }

    /// Returns the number of forks that exist for this collection.
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - The collection no longer exists on the server
    /// - Network communication fails
    /// - The authenticated user lacks sufficient permissions
    ///
    /// # Examples
    ///
    /// ```
    /// # use chroma::ChromaCollection;
    /// # async fn example(collection: ChromaCollection) -> Result<(), Box<dyn std::error::Error>> {
    /// let count = collection.fork_count().await?;
    /// println!("Collection has {} forks", count);
    /// # Ok(())
    /// # }
    /// ```
    pub async fn fork_count(&self) -> Result<usize, ChromaHttpClientError> {
        let response: ForkCountResponse = self
            .send::<(), _>(true, "fork_count", "fork_count", Method::GET, None)
            .await?;
        Ok(response.count)
    }

    /// Attaches a function to this collection.
    ///
    /// Functions execute automatically when records are written to this collection,
    /// producing results in a separate output collection.
    ///
    /// # Arguments
    ///
    /// * `function_id` - The function identifier (e.g., `"statistics"`, `"record_counter"`)
    /// * `name` - A unique name for this attached function instance
    /// * `output_collection` - Name of the collection where function output will be stored
    /// * `params` - Optional JSON parameters for the function
    ///
    /// # Returns
    ///
    /// A tuple of `(response, created)` where `created` is `true` if newly created,
    /// `false` if an attached function with this name already existed (idempotent).
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # use chroma::ChromaCollection;
    /// # async fn example(collection: ChromaCollection) -> Result<(), Box<dyn std::error::Error>> {
    /// let (response, created) = collection
    ///     .attach_function("statistics", "my_stats", "my_stats_output", None)
    ///     .await?;
    /// println!("Attached function ID: {}", response.attached_function.id);
    /// # Ok(())
    /// # }
    /// ```
    pub async fn attach_function(
        &self,
        function_id: impl AsRef<str>,
        name: impl AsRef<str>,
        output_collection: impl AsRef<str>,
        params: Option<serde_json::Value>,
    ) -> Result<(AttachFunctionResponse, bool), ChromaHttpClientError> {
        let body = serde_json::json!({
            "name": name.as_ref(),
            "function_id": function_id.as_ref(),
            "output_collection": output_collection.as_ref(),
            "params": params.unwrap_or(serde_json::json!({})),
        });
        let response: AttachFunctionResponse = self
            .send(
                false,
                "attach_function",
                "functions/attach",
                Method::POST,
                Some(body),
            )
            .await?;
        let created = response.created;
        Ok((response, created))
    }

    /// Detaches a function from this collection, preventing further executions.
    ///
    /// # Arguments
    ///
    /// * `name` - Name of the attached function instance to detach
    /// * `delete_output_collection` - Whether to also delete the output collection
    ///
    /// # Returns
    ///
    /// `true` if the function was successfully detached.
    ///
    /// # Examples
    ///
    /// ```no_run
    /// # use chroma::ChromaCollection;
    /// # async fn example(collection: ChromaCollection) -> Result<(), Box<dyn std::error::Error>> {
    /// let success = collection.detach_function("my_stats", false).await?;
    /// assert!(success);
    /// # Ok(())
    /// # }
    /// ```
    pub async fn detach_function(
        &self,
        name: impl AsRef<str>,
        delete_output_collection: bool,
    ) -> Result<bool, ChromaHttpClientError> {
        let path = format!("attached_functions/{}/detach", name.as_ref());
        let body = serde_json::json!({
            "delete_output": delete_output_collection,
        });
        let response: DetachFunctionResponse = self
            .send(false, "detach_function", &path, Method::POST, Some(body))
            .await?;
        Ok(response.success)
    }

    async fn resolve_embeddings(
        &self,
        embeddings: Option<Vec<Vec<f32>>>,
        documents: &Option<Vec<Option<String>>>,
    ) -> Result<Vec<Vec<f32>>, ChromaHttpClientError> {
        if let Some(embeddings) = embeddings {
            return Ok(embeddings);
        }

        let documents = documents
            .as_ref()
            .ok_or(ChromaHttpClientError::MissingDocumentsForEmbedding)?;
        let input = documents
            .iter()
            .map(|document| {
                document
                    .as_deref()
                    .ok_or(ChromaHttpClientError::MissingDocumentsForEmbedding)
            })
            .collect::<Result<Vec<_>, _>>()?;

        self.embed_documents(&input).await
    }

    async fn resolve_update_embeddings(
        &self,
        embeddings: Option<Vec<Option<Vec<f32>>>>,
        documents: &Option<Vec<Option<String>>>,
    ) -> Result<Option<Vec<Option<Vec<f32>>>>, ChromaHttpClientError> {
        if embeddings.is_some() || documents.is_none() || self.embedding_function.is_none() {
            return Ok(embeddings);
        }

        let documents = documents.as_ref().expect("checked above");
        let input = documents
            .iter()
            .filter_map(|document| document.as_deref())
            .collect::<Vec<_>>();
        if input.is_empty() {
            return Ok(None);
        }

        let mut embeddings = self.embed_documents(&input).await?.into_iter();
        documents
            .iter()
            .map(|document| {
                document
                    .as_ref()
                    .map(|_| {
                        embeddings.next().ok_or_else(|| {
                            ChromaHttpClientError::EmbeddingFunctionError(
                                "Embedding function returned fewer embeddings than documents"
                                    .to_string(),
                            )
                        })
                    })
                    .transpose()
            })
            .collect::<Result<Vec<_>, _>>()
            .map(Some)
    }

    /// Embeds documents with the collection's schema-derived dense embedding
    /// function, applying the *document-side* instruction.
    ///
    /// This is the same path [`add`](Self::add)/[`upsert`](Self::upsert) take
    /// when embeddings are omitted; it is also exposed publicly as the
    /// write-side counterpart of [`embed_query`](Self::embed_query) for callers
    /// that need to embed documents explicitly.
    ///
    /// Returns [`ChromaHttpClientError::MissingEmbeddingFunction`] when the
    /// collection has no dense embedding function configured.
    pub async fn embed_documents(
        &self,
        input: &[&str],
    ) -> Result<Vec<Vec<f32>>, ChromaHttpClientError> {
        self.embed_with(input, false).await
    }

    /// Embeds query text with the collection's schema-derived dense embedding
    /// function, applying the *query-side* instruction.
    ///
    /// This is the read-side mirror of the write path's document auto-embed
    /// ([`add`](Self::add) with omitted embeddings): both pull the same EF off
    /// the collection schema, but documents use the document instruction while
    /// queries use the query instruction. It is public because — unlike writes,
    /// which the server auto-embeds — Chroma's Search API takes pre-computed
    /// query vectors, so search callers must embed the query themselves and
    /// want the exact EF the collection was built with (no config drift).
    ///
    /// Returns [`ChromaHttpClientError::MissingEmbeddingFunction`] when the
    /// collection has no dense embedding function configured.
    ///
    // TODO: when the Search API can auto-embed query text server-side
    // (mirroring write-side auto-embed), callers could pass raw query text on
    // the `$knn` rank and let the EF run on the FE. This explicit client-side
    // path would stay as an option (e.g. local/precomputed embedding), not
    // necessarily be removed.
    pub async fn embed_query(
        &self,
        input: &[&str],
    ) -> Result<Vec<Vec<f32>>, ChromaHttpClientError> {
        self.embed_with(input, true).await
    }

    /// Shared dense-embedding helper. `query` selects the query-side instruction
    /// (`embed_query_strs`) over the document-side one (`embed_strs`).
    async fn embed_with(
        &self,
        input: &[&str],
        query: bool,
    ) -> Result<Vec<Vec<f32>>, ChromaHttpClientError> {
        let embedding_function = self
            .embedding_function
            .as_ref()
            .ok_or(ChromaHttpClientError::MissingEmbeddingFunction)?;
        let embeddings = if query {
            embedding_function.embed_query_strs(input).await
        } else {
            embedding_function.embed_strs(input).await
        }
        .map_err(|err| ChromaHttpClientError::EmbeddingFunctionError(err.to_string()))?;
        if embeddings.len() != input.len() {
            return Err(ChromaHttpClientError::EmbeddingFunctionError(format!(
                "Embedding function returned {} embeddings for {} inputs",
                embeddings.len(),
                input.len()
            )));
        }
        Ok(embeddings)
    }

    /// Internal transport method that constructs collection-specific API paths and delegates to the client.
    async fn send<Body: Serialize, Response: DeserializeOwned>(
        &self,
        read_only: bool,
        operation: &str,
        path: &str,
        method: Method,
        body: Option<Body>,
    ) -> Result<Response, ChromaHttpClientError> {
        self.send_with_query::<Body, (), Response>(
            read_only, operation, path, method, body, None::<()>,
        )
        .await
    }

    /// Internal transport method with query parameter support.
    async fn send_with_query<
        Body: Serialize,
        QueryParams: Serialize,
        Response: DeserializeOwned,
    >(
        &self,
        read_only: bool,
        operation: &str,
        path: &str,
        method: Method,
        body: Option<Body>,
        query_params: Option<QueryParams>,
    ) -> Result<Response, ChromaHttpClientError> {
        let operation_name = format!("collection_{operation}");
        let path = format!(
            "/api/v2/tenants/{}/databases/{}/collections/{}/{}",
            self.collection.tenant, self.collection.database, self.collection.collection_id, path
        );
        let path = path.trim_end_matches("/");

        if read_only {
            self.client
                .send_read_only(&operation_name, method, path, body, query_params)
                .await
        } else {
            self.client
                .send(&operation_name, method, path, body, query_params)
                .await
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::tests::{unique_collection_name, with_client};
    use crate::{
        client::{ChromaAuthMethod, ChromaHttpClientError, ChromaHttpClientOptions},
        embed::{chroma_cloud::ChromaCloudQwenEmbeddingFunction, EmbeddingFunction},
        ChromaCollection, ChromaHttpClient,
    };
    use chroma_types::operator::{Key, QueryVector, RankExpr};
    use chroma_types::plan::{ReadLevel, SearchPayload};
    use chroma_types::{
        Collection, Include, IncludeList, Metadata, MetadataComparison, MetadataExpression,
        MetadataValue, PrimitiveOperator, Schema, UpdateMetadata, UpdateMetadataValue, Where,
    };
    use std::sync::Arc;

    #[derive(Debug, thiserror::Error)]
    #[error("test embedding failed")]
    struct TestEmbeddingError;

    struct TestEmbeddingFunction;

    #[async_trait::async_trait]
    impl EmbeddingFunction for TestEmbeddingFunction {
        type Embedding = Vec<f32>;
        type Error = TestEmbeddingError;

        async fn embed_strs(&self, batches: &[&str]) -> Result<Vec<Vec<f32>>, Self::Error> {
            Ok(batches
                .iter()
                .map(|document| vec![document.len() as f32])
                .collect())
        }
    }

    fn test_collection() -> ChromaCollection {
        let collection = Collection {
            tenant: "tenant".to_string(),
            database: "database".to_string(),
            ..Default::default()
        };
        ChromaCollection {
            client: ChromaHttpClient::default(),
            collection: Arc::new(collection),
            embedding_function: None,
        }
    }

    #[test]
    fn new_auto_attaches_chroma_cloud_qwen_embedding_function() {
        let client = ChromaHttpClient::new(ChromaHttpClientOptions {
            auth_method: ChromaAuthMethod::cloud_api_key("test-api-key").unwrap(),
            ..Default::default()
        });
        let collection = Collection {
            tenant: "tenant".to_string(),
            database: "database".to_string(),
            schema: Some(Schema::default_with_embedding_function(
                ChromaCloudQwenEmbeddingFunction::configuration().build(),
            )),
            ..Default::default()
        };

        let collection = ChromaCollection::new(client, collection);

        assert!(collection.embedding_function.is_some());
    }

    #[tokio::test]
    async fn resolve_embeddings_uses_embedding_function() {
        let mut collection = test_collection();
        collection.set_embedding_function(Some(TestEmbeddingFunction));

        let embeddings = collection
            .resolve_embeddings(
                None,
                &Some(vec![Some("alpha".to_string()), Some("beta".to_string())]),
            )
            .await
            .unwrap();

        assert_eq!(embeddings, vec![vec![5.0], vec![4.0]]);
    }

    #[tokio::test]
    async fn resolve_embeddings_requires_documents() {
        let mut collection = test_collection();
        collection.set_embedding_function(Some(TestEmbeddingFunction));

        let err = collection
            .resolve_embeddings(None, &None)
            .await
            .unwrap_err();

        assert!(matches!(
            err,
            ChromaHttpClientError::MissingDocumentsForEmbedding
        ));
    }

    #[tokio::test]
    async fn resolve_update_embeddings_embeds_only_present_documents() {
        let mut collection = test_collection();
        collection.set_embedding_function(Some(TestEmbeddingFunction));

        let embeddings = collection
            .resolve_update_embeddings(
                None,
                &Some(vec![Some("alpha".to_string()), None, Some("z".to_string())]),
            )
            .await
            .unwrap();

        assert_eq!(
            embeddings,
            Some(vec![Some(vec![5.0]), None, Some(vec![1.0])])
        );
    }

    #[tokio::test]
    async fn resolve_update_embeddings_allows_documents_without_embedding_function() {
        let collection = test_collection();

        let embeddings = collection
            .resolve_update_embeddings(None, &Some(vec![Some("updated document".to_string())]))
            .await
            .unwrap();

        assert_eq!(embeddings, None);
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_accessor_methods() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_accessors").await;
            assert!(!collection.database().is_empty());
            assert_eq!(collection.metadata(), &None);
            assert!(collection.schema().is_some());
            assert!(!collection.tenant().is_empty());
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_count_empty_collection() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_count_empty").await;

            let count = collection.count().await.unwrap();
            println!("Empty collection count: {}", count);
            assert_eq!(count, 0);
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_get_indexing_status() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_indexing_status").await;

            // Test on empty collection
            let status = collection.get_indexing_status().await.unwrap();
            println!("Indexing status: {:?}", status);
            assert_eq!(status.total_ops, 0);
            assert_eq!(status.num_indexed_ops, 0);
            assert_eq!(status.num_unindexed_ops, 0);

            // Add some records
            collection
                .add(
                    vec!["id1".to_string(), "id2".to_string()],
                    vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]],
                    None,
                    None,
                    None,
                )
                .await
                .unwrap();

            // Check status after adding records
            let status = collection.get_indexing_status().await.unwrap();
            println!("Indexing status after add: {:?}", status);
            assert_eq!(status.total_ops, 2);
            assert!(status.op_indexing_progress >= 0.0 && status.op_indexing_progress <= 1.0);
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_add_single_record() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_add_single").await;

            collection
                .add(
                    vec!["id1".to_string()],
                    vec![vec![1.0, 2.0, 3.0]],
                    Some(vec![Some("document1".to_string())]),
                    None,
                    None,
                )
                .await
                .unwrap();

            let count = collection.count().await.unwrap();
            println!("Collection count after add: {}", count);
            assert_eq!(count, 1);
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_add_multiple_records() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_add_multiple").await;

            collection
                .add(
                    vec!["id1".to_string(), "id2".to_string(), "id3".to_string()],
                    vec![
                        vec![1.0, 2.0, 3.0],
                        vec![4.0, 5.0, 6.0],
                        vec![7.0, 8.0, 9.0],
                    ],
                    Some(vec![
                        Some("first document".to_string()),
                        Some("second document".to_string()),
                        Some("third document".to_string()),
                    ]),
                    None,
                    None,
                )
                .await
                .unwrap();

            let count = collection.count().await.unwrap();
            println!("Collection count after adding multiple: {}", count);
            assert_eq!(count, 3);
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_add_with_metadata() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_add_metadata").await;

            let mut metadata = Metadata::new();
            metadata.insert("category".to_string(), "test".into());
            metadata.insert("version".to_string(), 1.into());

            collection
                .add(
                    vec!["id1".to_string()],
                    vec![vec![1.0, 2.0, 3.0]],
                    Some(vec![Some("document with metadata".to_string())]),
                    None,
                    Some(vec![Some(metadata)]),
                )
                .await
                .unwrap();

            let count = collection.count().await.unwrap();
            assert_eq!(count, 1);
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_add_with_uris() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_add_uris").await;

            collection
                .add(
                    vec!["id1".to_string()],
                    vec![vec![1.0, 2.0, 3.0]],
                    Some(vec![Some("document with uri".to_string())]),
                    Some(vec![Some("https://example.com/doc1".to_string())]),
                    None,
                )
                .await
                .unwrap();

            let count = collection.count().await.unwrap();
            assert_eq!(count, 1);
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_get_all_records() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_get_all").await;

            collection
                .add(
                    vec!["id1".to_string(), "id2".to_string()],
                    vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]],
                    Some(vec![Some("first".to_string()), Some("second".to_string())]),
                    None,
                    None,
                )
                .await
                .unwrap();

            let response = collection.get(None, None, None, None, None).await.unwrap();
            println!("Get all response: {:?}", response);
            assert_eq!(response.ids.len(), 2);
            assert!(response.ids.contains(&"id1".to_string()));
            assert!(response.ids.contains(&"id2".to_string()));
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_get_by_ids() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_get_by_ids").await;

            collection
                .add(
                    vec!["id1".to_string(), "id2".to_string(), "id3".to_string()],
                    vec![
                        vec![1.0, 2.0, 3.0],
                        vec![4.0, 5.0, 6.0],
                        vec![7.0, 8.0, 9.0],
                    ],
                    None,
                    None,
                    None,
                )
                .await
                .unwrap();

            let response = collection
                .get(
                    Some(vec!["id1".to_string(), "id3".to_string()]),
                    None,
                    None,
                    None,
                    None,
                )
                .await
                .unwrap();
            println!("Get by ids response: {:?}", response);
            assert_eq!(response.ids.len(), 2);
            assert!(response.ids.contains(&"id1".to_string()));
            assert!(response.ids.contains(&"id3".to_string()));
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_get_with_limit_and_offset() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_get_limit_offset").await;

            collection
                .add(
                    vec![
                        "id1".to_string(),
                        "id2".to_string(),
                        "id3".to_string(),
                        "id4".to_string(),
                    ],
                    vec![
                        vec![1.0, 2.0, 3.0],
                        vec![4.0, 5.0, 6.0],
                        vec![7.0, 8.0, 9.0],
                        vec![10.0, 11.0, 12.0],
                    ],
                    None,
                    None,
                    None,
                )
                .await
                .unwrap();

            let response = collection
                .get(None, None, Some(2), Some(1), None)
                .await
                .unwrap();
            println!("Get with limit and offset response: {:?}", response);
            assert_eq!(response.ids.len(), 2);
            assert!(!response.ids.is_empty());
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_get_with_where_clause() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_get_where").await;

            let mut metadata1 = Metadata::new();
            metadata1.insert("category".to_string(), "a".into());

            let mut metadata2 = Metadata::new();
            metadata2.insert("category".to_string(), "b".into());

            collection
                .add(
                    vec!["id1".to_string(), "id2".to_string()],
                    vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]],
                    None,
                    None,
                    Some(vec![Some(metadata1), Some(metadata2)]),
                )
                .await
                .unwrap();

            let where_clause = Where::Metadata(MetadataExpression {
                key: "category".to_string(),
                comparison: MetadataComparison::Primitive(
                    PrimitiveOperator::Equal,
                    MetadataValue::Str("a".to_string()),
                ),
            });
            let response = collection
                .get(None, Some(where_clause), None, None, None)
                .await
                .unwrap();
            println!("Get with where clause response: {:?}", response);
            assert_eq!(response.ids.len(), 1);
            assert_eq!(response.ids[0], "id1");
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_get_with_include_list() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_get_include").await;

            collection
                .add(
                    vec!["id1".to_string()],
                    vec![vec![1.0, 2.0, 3.0]],
                    Some(vec![Some("test document".to_string())]),
                    None,
                    None,
                )
                .await
                .unwrap();

            let include = IncludeList(vec![
                Include::Document,
                Include::Embedding,
                Include::Metadata,
            ]);
            let response = collection
                .get(None, None, None, None, Some(include))
                .await
                .unwrap();
            println!("Get with include list response: {:?}", response);
            assert_eq!(response.ids.len(), 1);
            assert_eq!(response.ids[0], "id1");
            assert!(response.documents.is_some());
            assert_eq!(
                response.documents.as_ref().unwrap()[0],
                Some("test document".to_string())
            );
            assert!(response.embeddings.is_some());
            assert_eq!(
                response.embeddings.as_ref().unwrap()[0],
                vec![1.0, 2.0, 3.0]
            );
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_query_basic() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_query_basic").await;

            collection
                .add(
                    vec!["id1".to_string(), "id2".to_string(), "id3".to_string()],
                    vec![
                        vec![1.0, 2.0, 3.0],
                        vec![1.1, 2.1, 3.1],
                        vec![10.0, 20.0, 30.0],
                    ],
                    Some(vec![
                        Some("first".to_string()),
                        Some("second".to_string()),
                        Some("third".to_string()),
                    ]),
                    None,
                    None,
                )
                .await
                .unwrap();

            let response = collection
                .query(vec![vec![1.0, 2.0, 3.0]], None, None, None, None)
                .await
                .unwrap();
            println!("Query basic response: {:?}", response);
            assert_eq!(response.ids.len(), 1);
            assert!(!response.ids[0].is_empty());
            assert!(response.ids[0].contains(&"id1".to_string()));
            assert!(response.distances.is_some());
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_query_with_n_results() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_query_n_results").await;

            collection
                .add(
                    vec![
                        "id1".to_string(),
                        "id2".to_string(),
                        "id3".to_string(),
                        "id4".to_string(),
                        "id5".to_string(),
                    ],
                    vec![
                        vec![1.0, 2.0, 3.0],
                        vec![1.1, 2.1, 3.1],
                        vec![1.2, 2.2, 3.2],
                        vec![1.3, 2.3, 3.3],
                        vec![1.4, 2.4, 3.4],
                    ],
                    None,
                    None,
                    None,
                )
                .await
                .unwrap();

            let response = collection
                .query(vec![vec![1.0, 2.0, 3.0]], Some(3), None, None, None)
                .await
                .unwrap();
            println!("Query with n_results response: {:?}", response);
            assert_eq!(response.ids.len(), 1);
            assert_eq!(response.ids[0].len(), 3);
            assert!(response.distances.is_some());
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_query_with_where_clause() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_query_where").await;

            let mut metadata1 = Metadata::new();
            metadata1.insert("category".to_string(), "a".into());

            let mut metadata2 = Metadata::new();
            metadata2.insert("category".to_string(), "b".into());

            collection
                .add(
                    vec!["id1".to_string(), "id2".to_string()],
                    vec![vec![1.0, 2.0, 3.0], vec![1.1, 2.1, 3.1]],
                    None,
                    None,
                    Some(vec![Some(metadata1), Some(metadata2)]),
                )
                .await
                .unwrap();

            let where_clause = Where::Metadata(MetadataExpression {
                key: "category".to_string(),
                comparison: MetadataComparison::Primitive(
                    PrimitiveOperator::Equal,
                    MetadataValue::Str("a".to_string()),
                ),
            });
            let response = collection
                .query(
                    vec![vec![1.0, 2.0, 3.0]],
                    None,
                    Some(where_clause),
                    None,
                    None,
                )
                .await
                .unwrap();
            println!("Query with where clause response: {:?}", response);
            assert_eq!(response.ids.len(), 1);
            assert_eq!(response.ids[0].len(), 1);
            assert_eq!(response.ids[0][0], "id1");
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_query_multiple_embeddings() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_query_multiple").await;

            collection
                .add(
                    vec!["id1".to_string(), "id2".to_string(), "id3".to_string()],
                    vec![
                        vec![1.0, 2.0, 3.0],
                        vec![4.0, 5.0, 6.0],
                        vec![7.0, 8.0, 9.0],
                    ],
                    None,
                    None,
                    None,
                )
                .await
                .unwrap();

            let response = collection
                .query(
                    vec![vec![1.0, 2.0, 3.0], vec![7.0, 8.0, 9.0]],
                    Some(1),
                    None,
                    None,
                    None,
                )
                .await
                .unwrap();
            println!("Query multiple embeddings response: {:?}", response);
            assert_eq!(response.ids.len(), 2);
            assert_eq!(response.ids[0].len(), 1);
            assert_eq!(response.ids[1].len(), 1);
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_search_with_read_levels() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_search_read_level").await;

            collection
                .add(
                    vec!["id1".to_string(), "id2".to_string(), "id3".to_string()],
                    vec![
                        vec![1.0, 2.0, 3.0],
                        vec![1.1, 2.1, 3.1],
                        vec![0.9, 1.9, 2.9],
                    ],
                    Some(vec![
                        Some("first".to_string()),
                        Some("second".to_string()),
                        Some("third".to_string()),
                    ]),
                    None,
                    None,
                )
                .await
                .unwrap();

            let search = SearchPayload::default()
                .rank(RankExpr::Knn {
                    query: QueryVector::Dense(vec![1.0, 2.0, 3.0]),
                    key: Key::Embedding,
                    limit: 10,
                    default: None,
                    return_rank: false,
                })
                .limit(Some(5), 0)
                .select([Key::Document, Key::Score]);

            let index_and_wal = collection
                .search_with_options(vec![search.clone()], ReadLevel::IndexAndWal)
                .await
                .unwrap();
            assert_eq!(index_and_wal.ids.len(), 1);
            assert!(!index_and_wal.ids[0].is_empty());
            assert_eq!(index_and_wal.ids[0].len(), 3);
            assert!(index_and_wal.documents[0].is_some());
            assert!(index_and_wal.scores[0].is_some());

            // IndexOnly may omit recent WAL writes; just ensure the call succeeds
            // and the response structure matches the requested payload.
            let index_only = collection
                .search_with_options(vec![search.clone()], ReadLevel::IndexOnly)
                .await
                .unwrap();
            assert_eq!(index_only.ids.len(), 1);
            assert!(index_only.documents[0].is_some());
            assert!(index_only.scores[0].is_some());

            // IndexAndBoundedWal reads up to a server-configured number of WAL entries;
            // just ensure the call succeeds and the response structure is valid.
            let bounded = collection
                .search_with_options(vec![search], ReadLevel::IndexAndBoundedWal)
                .await
                .unwrap();
            assert_eq!(bounded.ids.len(), 1);
            assert!(bounded.documents[0].is_some());
            assert!(bounded.scores[0].is_some());
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_count_with_read_levels() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_count_read_level").await;

            collection
                .add(
                    vec!["id1".to_string(), "id2".to_string(), "id3".to_string()],
                    vec![
                        vec![1.0, 2.0, 3.0],
                        vec![1.1, 2.1, 3.1],
                        vec![0.9, 1.9, 2.9],
                    ],
                    None,
                    None,
                    None,
                )
                .await
                .unwrap();

            // INDEX_AND_WAL should see all committed writes
            let count = collection
                .count_with_options(ReadLevel::IndexAndWal)
                .await
                .unwrap();
            assert_eq!(count, 3);

            // INDEX_ONLY may omit recent WAL writes; just ensure the call succeeds
            let count = collection
                .count_with_options(ReadLevel::IndexOnly)
                .await
                .unwrap();
            assert!(count <= 3);

            // INDEX_AND_BOUNDED_WAL reads up to a configured limit; just ensure the call succeeds
            let count = collection
                .count_with_options(ReadLevel::IndexAndBoundedWal)
                .await
                .unwrap();
            assert!(count <= 3);
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_update_embeddings() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_update_embeddings").await;

            collection
                .add(
                    vec!["id1".to_string()],
                    vec![vec![1.0, 2.0, 3.0]],
                    Some(vec![Some("original".to_string())]),
                    None,
                    None,
                )
                .await
                .unwrap();

            collection
                .update(
                    vec!["id1".to_string()],
                    Some(vec![Some(vec![4.0, 5.0, 6.0])]),
                    None,
                    None,
                    None,
                )
                .await
                .unwrap();

            let get_response = collection
                .get(
                    Some(vec!["id1".to_string()]),
                    None,
                    None,
                    None,
                    Some(IncludeList(vec![Include::Embedding])),
                )
                .await
                .unwrap();
            println!("Get after update response: {:?}", get_response);
            assert!(get_response.embeddings.is_some());
            assert_eq!(
                get_response.embeddings.as_ref().unwrap()[0],
                vec![4.0, 5.0, 6.0]
            );
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_update_documents() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_update_documents").await;

            collection
                .add(
                    vec!["id1".to_string()],
                    vec![vec![1.0, 2.0, 3.0]],
                    Some(vec![Some("original".to_string())]),
                    None,
                    None,
                )
                .await
                .unwrap();

            collection
                .update(
                    vec!["id1".to_string()],
                    None,
                    Some(vec![Some("updated document".to_string())]),
                    None,
                    None,
                )
                .await
                .unwrap();

            let get_response = collection
                .get(
                    Some(vec!["id1".to_string()]),
                    None,
                    None,
                    None,
                    Some(IncludeList(vec![Include::Document])),
                )
                .await
                .unwrap();
            println!("Get after update response: {:?}", get_response);
            assert!(get_response.documents.is_some());
            assert_eq!(
                get_response.documents.as_ref().unwrap()[0],
                Some("updated document".to_string())
            );
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_update_metadata() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_update_metadata").await;

            let mut original_metadata = Metadata::new();
            original_metadata.insert("version".to_string(), 1.into());

            collection
                .add(
                    vec!["id1".to_string()],
                    vec![vec![1.0, 2.0, 3.0]],
                    None,
                    None,
                    Some(vec![Some(original_metadata)]),
                )
                .await
                .unwrap();

            let mut updated_metadata = UpdateMetadata::new();
            updated_metadata.insert("version".to_string(), UpdateMetadataValue::Int(2));
            updated_metadata.insert(
                "new_field".to_string(),
                UpdateMetadataValue::Str("test".to_string()),
            );

            collection
                .update(
                    vec!["id1".to_string()],
                    None,
                    None,
                    None,
                    Some(vec![Some(updated_metadata)]),
                )
                .await
                .unwrap();

            let get_response = collection
                .get(
                    Some(vec!["id1".to_string()]),
                    None,
                    None,
                    None,
                    Some(IncludeList(vec![Include::Metadata])),
                )
                .await
                .unwrap();
            println!("Get after update response: {:?}", get_response);
            assert!(get_response.metadatas.is_some());
            let metadata = get_response.metadatas.as_ref().unwrap()[0]
                .as_ref()
                .unwrap();
            assert_eq!(metadata.get("version"), Some(&MetadataValue::Int(2)));
            assert_eq!(
                metadata.get("new_field"),
                Some(&MetadataValue::Str("test".to_string()))
            );
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_upsert_insert_new() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_upsert_insert").await;

            collection
                .upsert(
                    vec!["id1".to_string()],
                    vec![vec![1.0, 2.0, 3.0]],
                    Some(vec![Some("new document".to_string())]),
                    None,
                    None,
                )
                .await
                .unwrap();

            let count = collection.count().await.unwrap();
            println!("Count after upsert insert: {}", count);
            assert_eq!(count, 1);
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_upsert_update_existing() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_upsert_update").await;

            collection
                .add(
                    vec!["id1".to_string()],
                    vec![vec![1.0, 2.0, 3.0]],
                    Some(vec![Some("original".to_string())]),
                    None,
                    None,
                )
                .await
                .unwrap();

            collection
                .upsert(
                    vec!["id1".to_string()],
                    vec![vec![4.0, 5.0, 6.0]],
                    Some(vec![Some("updated via upsert".to_string())]),
                    None,
                    None,
                )
                .await
                .unwrap();

            let count = collection.count().await.unwrap();
            println!("Count after upsert update: {}", count);
            assert_eq!(count, 1);
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_upsert_mixed() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_upsert_mixed").await;

            collection
                .add(
                    vec!["id1".to_string()],
                    vec![vec![1.0, 2.0, 3.0]],
                    Some(vec![Some("existing".to_string())]),
                    None,
                    None,
                )
                .await
                .unwrap();

            collection
                .upsert(
                    vec!["id1".to_string(), "id2".to_string()],
                    vec![vec![4.0, 5.0, 6.0], vec![7.0, 8.0, 9.0]],
                    Some(vec![Some("updated".to_string()), Some("new".to_string())]),
                    None,
                    None,
                )
                .await
                .unwrap();

            let count = collection.count().await.unwrap();
            println!("Count after upsert mixed: {}", count);
            assert_eq!(count, 2);
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_delete_by_ids() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_delete_by_ids").await;

            collection
                .add(
                    vec!["id1".to_string(), "id2".to_string(), "id3".to_string()],
                    vec![
                        vec![1.0, 2.0, 3.0],
                        vec![4.0, 5.0, 6.0],
                        vec![7.0, 8.0, 9.0],
                    ],
                    None,
                    None,
                    None,
                )
                .await
                .unwrap();

            collection
                .delete(Some(vec!["id1".to_string(), "id3".to_string()]), None, None)
                .await
                .unwrap();

            let count = collection.count().await.unwrap();
            println!("Count after delete: {}", count);
            assert_eq!(count, 1);
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_delete_by_where() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_delete_by_where").await;

            let mut metadata1 = Metadata::new();
            metadata1.insert("category".to_string(), "a".into());

            let mut metadata2 = Metadata::new();
            metadata2.insert("category".to_string(), "b".into());

            collection
                .add(
                    vec!["id1".to_string(), "id2".to_string()],
                    vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]],
                    None,
                    None,
                    Some(vec![Some(metadata1), Some(metadata2)]),
                )
                .await
                .unwrap();

            let where_clause = Where::Metadata(MetadataExpression {
                key: "category".to_string(),
                comparison: MetadataComparison::Primitive(
                    PrimitiveOperator::Equal,
                    MetadataValue::Str("a".to_string()),
                ),
            });
            collection
                .delete(None, Some(where_clause), None)
                .await
                .unwrap();

            let count = collection.count().await.unwrap();
            println!("Count after delete: {}", count);
            assert_eq!(count, 1);
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_delete_with_limit() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_delete_with_limit").await;

            let mut metadata_a = Metadata::new();
            metadata_a.insert("category".to_string(), "a".into());

            let mut metadata_b = Metadata::new();
            metadata_b.insert("category".to_string(), "b".into());

            collection
                .add(
                    vec![
                        "id1".to_string(),
                        "id2".to_string(),
                        "id3".to_string(),
                        "id4".to_string(),
                        "id5".to_string(),
                    ],
                    vec![
                        vec![1.0, 2.0, 3.0],
                        vec![4.0, 5.0, 6.0],
                        vec![7.0, 8.0, 9.0],
                        vec![10.0, 11.0, 12.0],
                        vec![13.0, 14.0, 15.0],
                    ],
                    None,
                    None,
                    Some(vec![
                        Some(metadata_a.clone()),
                        Some(metadata_a.clone()),
                        Some(metadata_a),
                        Some(metadata_b.clone()),
                        Some(metadata_b),
                    ]),
                )
                .await
                .unwrap();

            // Where matches 3 records (category == "a"), but limit is 2.
            let where_clause = Where::Metadata(MetadataExpression {
                key: "category".to_string(),
                comparison: MetadataComparison::Primitive(
                    PrimitiveOperator::Equal,
                    MetadataValue::Str("a".to_string()),
                ),
            });

            let response = collection
                .delete(None, Some(where_clause), Some(2))
                .await
                .unwrap();

            assert_eq!(response.deleted, 2);

            let count = collection.count().await.unwrap();
            assert_eq!(count, 3);
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_fork_basic() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_fork_source").await;

            collection
                .add(
                    vec!["id1".to_string(), "id2".to_string()],
                    vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]],
                    Some(vec![Some("first".to_string()), Some("second".to_string())]),
                    None,
                    None,
                )
                .await
                .unwrap();

            let target_name = unique_collection_name("test_fork_target");
            let forked = collection.fork(target_name.clone()).await.unwrap();
            client.track(&forked);
            println!("Forked collection: {:?}", forked);

            assert_eq!(forked.collection.name, target_name);
            assert_ne!(
                forked.collection.collection_id,
                collection.collection.collection_id
            );

            let forked_count = forked.count().await.unwrap();
            println!("Forked collection count: {}", forked_count);
            assert_eq!(forked_count, 2);
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_fork_preserves_data() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_fork_preserves_source").await;

            collection
                .add(
                    vec!["id1".to_string()],
                    vec![vec![1.0, 2.0, 3.0]],
                    Some(vec![Some("test document".to_string())]),
                    None,
                    None,
                )
                .await
                .unwrap();

            let target_name = unique_collection_name("test_fork_preserves_target");
            let forked = collection.fork(target_name).await.unwrap();
            client.track(&forked);

            let forked_get_response = forked
                .get(
                    None,
                    None,
                    None,
                    None,
                    Some(IncludeList(vec![Include::Document])),
                )
                .await
                .unwrap();
            println!("Forked collection get response: {:?}", forked_get_response);
            assert_eq!(forked_get_response.ids.len(), 1);
            assert_eq!(forked_get_response.ids[0], "id1");
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_fork_independence() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_fork_independence_source").await;

            collection
                .add(
                    vec!["id1".to_string()],
                    vec![vec![1.0, 2.0, 3.0]],
                    None,
                    None,
                    None,
                )
                .await
                .unwrap();

            let target_name = unique_collection_name("test_fork_independence_target");
            let forked = collection.fork(target_name).await.unwrap();
            client.track(&forked);

            forked
                .add(
                    vec!["id2".to_string()],
                    vec![vec![4.0, 5.0, 6.0]],
                    None,
                    None,
                    None,
                )
                .await
                .unwrap();

            let original_count = collection.count().await.unwrap();
            let forked_count = forked.count().await.unwrap();
            println!(
                "Original count: {}, Forked count: {}",
                original_count, forked_count
            );
            assert_eq!(original_count, 1);
            assert_eq!(forked_count, 2);
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_modify() {
        with_client(|mut client| async move {
            let mut collection = client.new_collection("test_modify").await;

            let mut new_metadata = Metadata::new();
            new_metadata.insert("foo".into(), "bar".into());

            collection
                .modify(None::<String>, Some(new_metadata))
                .await
                .unwrap();
            assert_eq!(
                collection.metadata().as_ref().unwrap().get("foo"),
                Some(&"bar".into())
            );

            let collection = client.get_collection(collection.name()).await.unwrap();
            assert_eq!(
                collection.metadata().as_ref().unwrap().get("foo"),
                Some(&"bar".into())
            );
        })
        .await;
    }

    #[tokio::test]
    #[test_log::test]
    async fn test_k8s_integration_attach_get_detach_function() {
        with_client(|mut client| async move {
            let collection = client.new_collection("test_attach_function").await;

            let attach_name = "my_counter";
            let output_collection_name = unique_collection_name("test_attach_function_output");

            let (response, created) = collection
                .attach_function(
                    chroma_types::FUNCTION_RECORD_COUNTER_NAME,
                    attach_name,
                    output_collection_name.as_str(),
                    None,
                )
                .await
                .unwrap();
            assert!(created);
            assert_eq!(response.attached_function.name, attach_name);
            assert_eq!(
                response.attached_function.function_name,
                chroma_types::FUNCTION_RECORD_COUNTER_NAME
            );

            let retrieved = collection.get_attached_function(attach_name).await.unwrap();
            assert_eq!(retrieved.attached_function.name, attach_name);
            assert_eq!(
                retrieved.attached_function.function_name,
                chroma_types::FUNCTION_RECORD_COUNTER_NAME
            );
            assert_eq!(
                retrieved.attached_function.id.to_string(),
                response.attached_function.id
            );
            assert_eq!(
                retrieved.attached_function.output_collection_name,
                output_collection_name
            );

            let success = collection.detach_function(attach_name, true).await.unwrap();
            assert!(success);
        })
        .await;
    }
}