[ENH]: Add token storage for chroma cloud splade sparse vectors (#5944)

## Description of changes

Added token/label storage for Splade sparse vectors.

- Improvements & Bug fixes
  - ...
- New functionality
  - ...

## Test plan

_How are these changes tested?_

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

## Migration plan

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

## Observability plan

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

## Documentation Changes

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

Author: tanujnay112

chromadb/utils/embedding_functions/chroma_cloud_splade_embedding_function.py

@@ -3,7 +3,7 @@
     SparseVectors,
     Documents,
 )
-from typing import Dict, Any
+from typing import Dict, Any, List, Optional
 from enum import Enum
 from chromadb.utils.embedding_functions.schemas import validate_config_schema
 from chromadb.utils.sparse_embedding_utils import normalize_sparse_vector
@@ -21,6 +21,7 @@ def __init__(
         self,
         api_key_env_var: str = "CHROMA_API_KEY",
         model: ChromaCloudSpladeEmbeddingModel = ChromaCloudSpladeEmbeddingModel.SPLADE_PP_EN_V1,
+        store_tokens: bool = False,
     ):
         """
         Initialize the ChromaCloudSpladeEmbeddingFunction.
@@ -49,6 +50,7 @@ def __init__(
                 f"or in any existing client instances"
             )
         self.model = model
+        self.store_tokens = bool(store_tokens)
         self._api_url = "https://embed.trychroma.com/embed_sparse"
         self._session = httpx.Client()
         self._session.headers.update(
@@ -87,6 +89,7 @@ def __call__(self, input: Documents) -> SparseVectors:
             "texts": list(input),
             "task": "",
             "target": "",
+            "fetch_tokens": "true" if self.store_tokens is True else "false",
         }
 
         try:
@@ -120,13 +123,17 @@ def _parse_response(self, response: Any) -> SparseVectors:
             if isinstance(emb, dict):
                 indices = emb.get("indices", [])
                 values = emb.get("values", [])
+                raw_labels = emb.get("labels") if self.store_tokens else None
+                labels: Optional[List[str]] = raw_labels if raw_labels else None
             else:
                 # Already a SparseVector, extract its data
+                assert(isinstance(emb, SparseVector))
                 indices = emb.indices
                 values = emb.values
+                labels = emb.labels if self.store_tokens else None
 
             normalized_vectors.append(
-                normalize_sparse_vector(indices=indices, values=values)
+                normalize_sparse_vector(indices=indices, values=values, labels=labels)
             )
 
         return normalized_vectors
@@ -148,18 +155,23 @@ def build_from_config(
         return ChromaCloudSpladeEmbeddingFunction(
             api_key_env_var=api_key_env_var,
             model=ChromaCloudSpladeEmbeddingModel(model),
+            store_tokens=config.get("store_tokens", False),
         )
 
     def get_config(self) -> Dict[str, Any]:
-        return {"api_key_env_var": self.api_key_env_var, "model": self.model.value}
+        return {
+            "api_key_env_var": self.api_key_env_var,
+            "model": self.model.value,
+            "store_tokens": self.store_tokens,
+        }
 
     def validate_config_update(
         self, old_config: Dict[str, Any], new_config: Dict[str, Any]
     ) -> None:
-        if "model" in new_config:
-            raise ValueError(
-                "model cannot be changed after the embedding function has been initialized"
-            )
+        immutable_keys = {"store_tokens", "model"}
+        for key in immutable_keys:
+            if key in new_config and new_config[key] != old_config.get(key):
+                raise ValueError(f"Updating '{key}' is not supported for chroma-cloud-splade")
 
     @staticmethod
     def validate_config(config: Dict[str, Any]) -> None:

schemas/embedding_functions/chroma-cloud-splade.json

@@ -16,6 +16,11 @@
             "type": "string",
             "description": "Environment variable name that contains your API key for the Chroma Embedding API",
             "default": "CHROMA_API_KEY"
+        },
+        "store_tokens": {
+            "type": "boolean",
+            "description": "Whether to store token labels in the sparse vector output",
+            "default": false
         }
     },
     "required": [