[ENH]: (Rust client): add EF config #5809

codetheweb · 2025-11-04T22:33:35Z

Description of changes

Main changes:

There are now separate sparse and dense EF traits.
EF traits have a config GAT which must implement TryInto<EmbeddingFunctionConfiguration>.

This does not introduce any machinery to automatically persist/hydrate EFs--currently, users must call get_config() / build_from_config() themselves (or .try_into() on the config). We can add that as a follow-up later but it's pretty non-trivial to build:

must implement a registry system to map EF names to implementations, allowing third-party crates to register custom EFs
need to remove GATs or type erase EFs so we can store refs to generic EFs during hydration
signatures of methods like query() will have to change which is a significant breaking change

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?

github-actions · 2025-11-04T22:33:52Z

codetheweb · 2025-11-05T01:48:25Z

rust/chroma/src/collection.rs

 }

 impl ChromaCollection {
+    pub(crate) fn new(client: ChromaHttpClient, collection: Collection) -> Self {


this change ended up not being strictly necessary for this PR but makes things slightly cleaner so I kept it

propel-code-bot · 2025-11-05T02:04:39Z

Embed-Function Config GATs, Dense/Sparse Traits & Collection Helpers

Introduces explicit DenseEmbeddingFunction and SparseEmbeddingFunction traits with a config GAT, enabling round-trip serialise / deserialise of embedding-function parameters. Built-in BM25 and Ollama implementations are updated to comply, including config structs, TryFrom ⇄ TryInto<EmbeddingFunctionConfiguration> bridges, and helper builders. Supporting refactors touch collection/client helpers, schema validation and several utility impls (Key: AsRef<str>).

Key Changes

• Split generic EmbeddingFunction into DenseEmbeddingFunction and SparseEmbeddingFunction in rust/chroma/src/embed/mod.rs with new required fns build_from_config, get_config, get_name.
• Added config structs BM25Config and OllamaEmbeddingFunctionConfig + TryFrom → EmbeddingFunctionConfiguration implementations for persistence.
• Refactored BM25SparseEmbeddingFunction and OllamaEmbeddingFunction to GAT-based config, added error variants for (de)serialisation, simplified encode path (BM25 now infallible).
• Client/collection QoL: centralised ChromaCollection::new, replaced ad-hoc struct literal construction; ChromaHttpClient list/create now use that helper.
• Execution layer: added duplicate impl AsRef<str> for Key (and used in frontend key validation).
• Frontend fix: schema validation now calls key.as_ref() instead of allocating to_string().
• Misc clean-ups: hard-coded timeouts noted, TODOs added, minor doc / test updates.

Affected Areas

• rust/chroma/src/embed/* (traits + implementations)
• rust/chroma/src/client/chroma_http_client.rs
• rust/chroma/src/collection.rs
• rust/types/src/execution/operator.rs
• rust/frontend/src/impls/service_based_frontend.rs

This summary was automatically generated by @propel-code-bot

Sicheng-Pan · 2025-11-05T02:12:19Z

rust/chroma/src/embed/bm25.rs

+    token_max_length: usize,
+}
+
+impl TryInto<EmbeddingFunctionConfiguration> for BM25Config {


I thought TryFrom is always preferred over TryInto. Also why is it possible to return error?

will replace with TryFrom

error is possible because of serialization

propel-code-bot · 2025-11-05T02:12:33Z

rust/chroma/src/embed/ollama.rs

+    fn get_config(&self) -> Result<Self::Config, Self::Error> {
+        Ok(OllamaEmbeddingFunctionConfig {
+            url: self.host.clone(),
+            model_name: self.model.clone(),
+            timeout: 60,
+        })
+    }


[CriticalError]

The timeout is hardcoded here, which breaks the configuration round-trip property. If an OllamaEmbeddingFunction is created with a non-default timeout via build_from_config, this function will return an incorrect timeout value.

To fix this, the timeout should be stored on the OllamaEmbeddingFunction struct. This requires updating the struct definition, new(), and build_from_config() to handle the timeout field consistently.

For example, this function should be updated to use the stored value:

fn get_config(&self) -> Result<Self::Config, Self::Error> { Ok(OllamaEmbeddingFunctionConfig { url: self.host.clone(), model_name: self.model.clone(), timeout: self.timeout, // Use stored value }) }

Context for Agents

[**CriticalError**] The timeout is hardcoded here, which breaks the configuration round-trip property. If an `OllamaEmbeddingFunction` is created with a non-default timeout via `build_from_config`, this function will return an incorrect timeout value. To fix this, the `timeout` should be stored on the `OllamaEmbeddingFunction` struct. This requires updating the struct definition, `new()`, and `build_from_config()` to handle the `timeout` field consistently. For example, this function should be updated to use the stored value: ```rust fn get_config(&self) -> Result<Self::Config, Self::Error> { Ok(OllamaEmbeddingFunctionConfig { url: self.host.clone(), model_name: self.model.clone(), timeout: self.timeout, // Use stored value }) } ``` File: rust/chroma/src/embed/ollama.rs Line: 182

Sicheng-Pan · 2025-11-05T02:13:11Z

rust/chroma/src/embed/bm25.rs

-    H: TokenHasher + Send + Sync + 'static,
-{
-    type Embedding = SparseVector;
+impl SparseEmbeddingFunction for BM25SparseEmbeddingFunction<Bm25Tokenizer, Murmur3AbsHasher> {


why is the generics pinned here?

Sicheng-Pan · 2025-11-05T02:14:32Z

rust/chroma/src/embed/bm25.rs

+        Ok(Self {
+            tokenizer,
+            hasher: Murmur3AbsHasher::default(),
+            k: config.k,
+            b: config.b,
+            avg_len: config.avg_doc_length,
+        })


This forces people to use the builtin hasher and tokenizer. Personally I do not like this

if we supported custom hashers and tokenizers how would we build it in the other languages? i think having the default one is fine no? if they want to have a custom one they can write their own ef

For the hasher and tokenizers we provide in Rust, we need to copy the same logic in other language. We could have a default ef, but not in this way because people cannot impl EmbeddingFunction for BM25SparseEmbeddingFunction<?, ?> themselves without wrapping in around.

Sicheng-Pan · 2025-11-05T02:16:21Z

rust/chroma/src/embed/ollama.rs

+    timeout: u64,
+}
+
+impl TryInto<EmbeddingFunctionConfiguration> for OllamaEmbeddingFunctionConfig {


Sicheng-Pan · 2025-11-05T02:18:22Z

rust/types/src/execution/operator.rs

 }

+impl AsRef<str> for Key {
+    fn as_ref(&self) -> &str {


nit: we have multiple identical impl for Key -> &str, could be better if we unify them

jairad26 · 2025-11-05T02:13:35Z

rust/chroma/src/embed/mod.rs

+    where
+        Self: Sized;
+}
+


i know this isn't necessary at this point for what this PR accomplishes, but can we add a todo for the other trait functions (default space, supported spaces)?

jairad26 · 2025-11-05T02:14:54Z

rust/chroma/src/embed/mod.rs

+    /// # Ok(())
+    /// # }
+    /// ```
+    async fn embed_strs(&self, batches: &[&str]) -> Result<Vec<Vec<f32>>, Self::Error>;


another TODO for the embed_query_strs would be helpful for tracking

jairad26 · 2025-11-05T02:15:49Z

rust/chroma/src/embed/mod.rs

    /// # }
    /// ```
-    async fn embed_strs(&self, batches: &[&str]) -> Result<Vec<Self::Embedding>, Self::Error>;
+    async fn embed_strs(&self, batches: &[&str]) -> Result<Vec<SparseVector>, Self::Error>;


todo for embed_query_strs here as well

jairad26 · 2025-11-05T02:19:29Z

rust/types/src/execution/operator.rs

+impl AsRef<str> for Key {
+    fn as_ref(&self) -> &str {
+        match self {
+            Key::Document => "#document",


side note, can we make it so schema supports using Key::Document in source key?

rescrv

I'll defer to Sicheng on sparse vector feel. I think this makes sense.

codetheweb force-pushed the feat-rust-client-ef-config branch 2 times, most recently from 7dd143f to 481b589 Compare November 5, 2025 01:47

codetheweb commented Nov 5, 2025

View reviewed changes

codetheweb force-pushed the feat-rust-client-ef-config branch from 481b589 to 1842640 Compare November 5, 2025 01:51

codetheweb changed the title ~~[ENH]: (Rust client): add EF config & auto-hydration~~ [ENH]: (Rust client): add EF config Nov 5, 2025

codetheweb force-pushed the feat-rust-client-ef-config branch from 1842640 to e705502 Compare November 5, 2025 02:03

codetheweb marked this pull request as ready for review November 5, 2025 02:03

codetheweb requested a review from rescrv as a code owner November 5, 2025 02:03

codetheweb requested review from Sicheng-Pan and jairad26 November 5, 2025 02:03

codetheweb force-pushed the feat-rust-client-ef-config branch from e705502 to 126ac80 Compare November 5, 2025 02:08

Sicheng-Pan reviewed Nov 5, 2025

View reviewed changes

propel-code-bot bot reviewed Nov 5, 2025

View reviewed changes

Sicheng-Pan reviewed Nov 5, 2025

View reviewed changes

jairad26 reviewed Nov 5, 2025

View reviewed changes

codetheweb force-pushed the feat-rust-client-ef-config branch from 126ac80 to e14aaec Compare November 5, 2025 02:20

[ENH]: (Rust client): add EF config & auto-hydration

a3e645f

codetheweb force-pushed the feat-rust-client-ef-config branch from e14aaec to a3e645f Compare November 5, 2025 02:26

rescrv approved these changes Nov 10, 2025

View reviewed changes

[ENH]: (Rust client): add EF config #5809

Are you sure you want to change the base?

[ENH]: (Rust client): add EF config #5809

Uh oh!

Conversation

codetheweb commented Nov 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description of changes

Test plan

Migration plan

Observability plan

Documentation Changes

Uh oh!

github-actions bot commented Nov 4, 2025

Reviewer Checklist

Testing, Bugs, Errors, Logs, Documentation

System Compatibility

Quality

Uh oh!

Choose a reason for hiding this comment

Uh oh!

propel-code-bot bot commented Nov 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

propel-code-bot bot Nov 5, 2025

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Sicheng-Pan Nov 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jairad26 Nov 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rescrv left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

codetheweb commented Nov 4, 2025 •

edited

Loading

propel-code-bot bot commented Nov 5, 2025 •

edited

Loading

Sicheng-Pan Nov 5, 2025 •

edited

Loading

jairad26 Nov 5, 2025 •

edited

Loading