[BUG] use keys instead of key_overrides in query embedding strings (#5754)

## Description of changes

_Summarize the changes made by this PR._

- Improvements & Bug fixes
- There was a bug in sparse embedding query texts where `key_overrides`
was used during query embedding instead of `keys`, as per the rename.
This PR fixes that, and adds tests to ensure the happy path of embedding
query strings works as intended
- New functionality
  - ...

## Test plan

_How are these changes tested?_

- [x ] 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: jairad26

Cargo.lock

@@ -17,7 +17,6 @@
 from uuid import UUID
 
 from chromadb.api.types import (
-    EMBEDDING_KEY,
     URI,
     Schema,
     SparseVectorIndexConfig,
@@ -840,22 +839,22 @@ def _embed_knn_string_queries(self, knn: Any) -> Any:
 
         # Handle metadata field with potential sparse embedding
         schema = self.schema
-        if schema is None or key not in schema.key_overrides:
+        if schema is None or key not in schema.keys:
             raise ValueError(
                 f"Cannot embed string query for key '{key}': "
                 f"key not found in schema. Please provide an embedded vector or "
                 f"configure an embedding function for this key in the schema."
             )
 
-        value_type = schema.key_overrides[key]
+        value_type = schema.keys[key]
 
         # Check for sparse vector with embedding function
         if value_type.sparse_vector is not None:
             sparse_index = value_type.sparse_vector.sparse_vector_index
             if sparse_index is not None and sparse_index.enabled:
-                config = sparse_index.config
-                if config.embedding_function is not None:
-                    embedding_func = config.embedding_function
+                sparse_config = sparse_index.config
+                if sparse_config.embedding_function is not None:
+                    embedding_func = sparse_config.embedding_function
                     if not isinstance(embedding_func, SparseEmbeddingFunction):
                         embedding_func = cast(
                             SparseEmbeddingFunction[Any], embedding_func
@@ -887,9 +886,9 @@ def _embed_knn_string_queries(self, knn: Any) -> Any:
         if value_type.float_list is not None:
             vector_index = value_type.float_list.vector_index
             if vector_index is not None and vector_index.enabled:
-                config = vector_index.config
-                if config.embedding_function is not None:
-                    embedding_func = config.embedding_function
+                dense_config = vector_index.config
+                if dense_config.embedding_function is not None:
+                    embedding_func = dense_config.embedding_function
                     validate_embedding_function(embedding_func)
 
                     # Embed the query using the schema's embedding function

chromadb/api/models/CollectionCommon.py

@@ -1,4 +1,4 @@
-from chromadb.api import ClientAPI
+from chromadb.api import ClientAPI, ServerAPI
 from chromadb.api.types import (
     Schema,
     SparseVectorIndexConfig,
@@ -30,8 +30,10 @@
     register_sparse_embedding_function,
 )
 from chromadb.api.models.Collection import Collection
+from chromadb.api.models.CollectionCommon import CollectionCommon
 from chromadb.errors import InvalidArgumentError, InternalError
 from chromadb.execution.expression import Knn, Search
+from chromadb.types import Collection as CollectionModel
 from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple, cast
 from uuid import uuid4
 import numpy as np
@@ -606,6 +608,36 @@ def test_search_embeds_string_knn_queries(
     assert embedded_rank.query == [11.0, 12.5]
 
 
+@pytest.mark.skipif(is_spann_disabled_mode, reason=skip_reason_spann_disabled)
+def test_search_embeds_string_knn_queries_with_sparse_embedding_function(
+    client_factories: "ClientFactories",
+) -> None:
+    """_embed_search_string_queries should embed string KNN queries using collection EF."""
+
+    sparse_ef = DeterministicSparseEmbeddingFunction(label="sparse")
+    schema = Schema().create_index(
+        key="sparse_metadata",
+        config=SparseVectorIndexConfig(
+            source_key="raw_text", embedding_function=sparse_ef
+        ),
+    )
+    collection, _ = _create_isolated_collection(client_factories, schema=schema)
+
+    search = Search().rank(Knn(key="sparse_metadata", query="hello world"))
+
+    embedded_search = collection._embed_search_string_queries(search)
+
+    assert isinstance(search._rank, Knn)
+    assert search._rank.key == "sparse_metadata"
+    assert search._rank.query == "hello world"
+
+    embedded_rank = embedded_search._rank
+    assert isinstance(embedded_rank, Knn)
+    assert embedded_rank.key == "sparse_metadata"
+    print(embedded_rank.query)
+    assert embedded_rank.query == SparseVector(indices=[0], values=[11.0])
+
+
 @pytest.mark.skipif(is_spann_disabled_mode, reason=skip_reason_spann_disabled)
 def test_search_embeds_string_queries_in_nested_ranks(
     client_factories: "ClientFactories",
@@ -2048,7 +2080,9 @@ def test_sparse_vector_index_config_with_key_types(
     )
 
     # Verify sparse embeddings were generated from text_field
-    result2 = collection2.get(ids=["sparse-key-1", "sparse-key-2"], include=["metadatas"])
+    result2 = collection2.get(
+        ids=["sparse-key-1", "sparse-key-2"], include=["metadatas"]
+    )
     assert result2["metadatas"] is not None
     assert "sparse2" in result2["metadatas"][0]
     assert "sparse2" in result2["metadatas"][1]
@@ -2072,7 +2106,9 @@ def test_schema_rejects_special_key_in_create_index() -> None:
         schema.create_index(config=StringInvertedIndexConfig(), key="#custom_field")
 
     # Test with Key object starting with #
-    with pytest.raises(ValueError, match="Cannot create index on special key '#embedding'"):
+    with pytest.raises(
+        ValueError, match="Cannot create index on special key '#embedding'"
+    ):
         schema.create_index(config=StringInvertedIndexConfig(), key=Key.EMBEDDING)
 
 
@@ -2084,7 +2120,9 @@ def test_schema_rejects_special_key_in_delete_index() -> None:
         schema.delete_index(config=StringInvertedIndexConfig(), key="#custom_field")
 
     # Test with Key object starting with #
-    with pytest.raises(ValueError, match="Cannot delete index on special key '#document'"):
+    with pytest.raises(
+        ValueError, match="Cannot delete index on special key '#document'"
+    ):
         schema.delete_index(config=StringInvertedIndexConfig(), key=Key.DOCUMENT)
 
 
@@ -2122,7 +2160,13 @@ def test_server_validates_schema_with_special_keys(
     # This should be caught server-side by validate_schema()
     schema = Schema()
     # Bypass client-side validation by directly manipulating schema.keys
-    from chromadb.api.types import ValueTypes, StringValueType, StringInvertedIndexType, StringInvertedIndexConfig
+    from chromadb.api.types import (
+        ValueTypes,
+        StringValueType,
+        StringInvertedIndexType,
+        StringInvertedIndexConfig,
+    )
+
     schema.keys["#invalid_key"] = ValueTypes(
         string=StringValueType(
             string_inverted_index=StringInvertedIndexType(
@@ -2138,7 +2182,9 @@ def test_server_validates_schema_with_special_keys(
 
     # Verify server caught the invalid key
     error_msg = str(exc_info.value)
-    assert "#" in error_msg or "key" in error_msg.lower() or "invalid" in error_msg.lower()
+    assert (
+        "#" in error_msg or "key" in error_msg.lower() or "invalid" in error_msg.lower()
+    )
 
 
 @pytest.mark.skipif(is_spann_disabled_mode, reason=skip_reason_spann_disabled)
@@ -2153,14 +2199,18 @@ def test_server_validates_invalid_source_key_in_sparse_vector_config(
 
     # Create schema with invalid source_key
     # Bypass client-side validation by directly creating the config
-    from chromadb.api.types import ValueTypes, SparseVectorValueType, SparseVectorIndexType
+    from chromadb.api.types import (
+        ValueTypes,
+        SparseVectorValueType,
+        SparseVectorIndexType,
+    )
 
     schema = Schema()
     # Manually construct config with invalid source_key using model_construct to bypass validation
     invalid_config = SparseVectorIndexConfig.model_construct(
         embedding_function=None,
         source_key="#embedding",  # Invalid - should be rejected
-        bm25=None
+        bm25=None,
     )
 
     schema.keys["test_sparse"] = ValueTypes(
@@ -2178,11 +2228,17 @@ def test_server_validates_invalid_source_key_in_sparse_vector_config(
 
     # Verify server caught the invalid source_key
     error_msg = str(exc_info.value)
-    assert "source_key" in error_msg.lower() or "#" in error_msg or "document" in error_msg.lower()
+    assert (
+        "source_key" in error_msg.lower()
+        or "#" in error_msg
+        or "document" in error_msg.lower()
+    )
 
 
 @pytest.mark.skipif(is_spann_disabled_mode, reason=skip_reason_spann_disabled)
-def test_modify_collection_no_initial_config_creates_default_schema(client: ClientAPI) -> None:
+def test_modify_collection_no_initial_config_creates_default_schema(
+    client: ClientAPI,
+) -> None:
     """Test that modifying a collection without initial config/schema creates and updates default schema."""
     collection_name = f"test_modify_no_init_{uuid4()}"
 
@@ -2259,7 +2315,9 @@ def test_modify_collection_with_initial_spann_schema(client: ClientAPI) -> None:
 
 
 @pytest.mark.skipif(is_spann_disabled_mode, reason=skip_reason_spann_disabled)
-def test_modify_collection_updates_schema_spann_multiple_fields(client: ClientAPI) -> None:
+def test_modify_collection_updates_schema_spann_multiple_fields(
+    client: ClientAPI,
+) -> None:
     """Test that modifying multiple SPANN fields updates schema correctly."""
     collection_name = f"test_modify_schema_multi_{uuid4()}"
 
@@ -2372,3 +2430,37 @@ def test_modify_collection_preserves_other_schema_fields(client: ClientAPI) -> N
     assert refreshed_schema.defaults.boolean is not None
     assert refreshed_schema.defaults.sparse_vector is not None
     assert refreshed_schema.keys["#document"] is not None
+
+
+def test_embeds_using_schema_embedding_function() -> None:
+    """Test that embeddings are using the schema embedding function."""
+    schema = Schema().create_index(
+        config=VectorIndexConfig(embedding_function=SimpleEmbeddingFunction()),
+    )
+
+    collection_model = CollectionModel(
+        id=uuid4(),
+        name="schema_only_collection",
+        configuration_json={},
+        serialized_schema=schema.serialize_to_json(),
+        metadata=None,
+        dimension=4,
+        tenant="tenant",
+        database="database",
+        version=0,
+        log_position=0,
+    )
+
+    collection = CollectionCommon(
+        client=cast(ServerAPI, object()),
+        model=collection_model,
+        embedding_function=None,
+    )
+
+    assert collection._embedding_function is None
+    assert collection.configuration is not None
+    assert collection.configuration.get("embedding_function") is None
+
+    embeddings = collection._embed(["hello world"])
+    assert embeddings is not None
+    assert np.allclose(embeddings[0], [0.0, 1.0, 2.0, 3.0])

chromadb/test/api/test_schema_e2e.py

@@ -4,7 +4,7 @@ import type { SearchResponse, SparseVector } from "../src/api";
 import { CollectionImpl } from "../src/collection";
 import type { CollectionConfiguration } from "../src/collection-configuration";
 import type { ChromaClient } from "../src/chroma-client";
-import type { EmbeddingFunction } from "../src/embedding-function";
+import type { EmbeddingFunction, SparseEmbeddingFunction } from "../src/embedding-function";
 
 class QueryMockEmbedding implements EmbeddingFunction {
   public readonly name = "query_mock";
@@ -402,4 +402,94 @@ describe("search expression DSL", () => {
     expect(knnPayload.key).toBe("#embedding");
     expect(knnPayload.limit).toBe(7);
   });
+
+  test("search auto-embeds string knn queries with sparse embedding function", async () => {
+    const queryText = "hello world";
+
+    class DeterministicSparseEmbedding implements SparseEmbeddingFunction {
+      public readonly name = "deterministic_sparse";
+
+      constructor(private readonly label = "sparse") { }
+
+      async generate(texts: string[]): Promise<SparseVector[]> {
+        return texts.map((text) => {
+          if (text === "hello world") {
+            return { indices: [0], values: [11.0] };
+          }
+          return { indices: [], values: [] };
+        });
+      }
+
+      getConfig(): Record<string, any> {
+        return { label: this.label };
+      }
+
+      static buildFromConfig(config: Record<string, any>): DeterministicSparseEmbedding {
+        return new DeterministicSparseEmbedding(config.label);
+      }
+    }
+
+    const sparseEf = new DeterministicSparseEmbedding("sparse");
+    const generateSpy = jest.spyOn(sparseEf, "generate");
+
+    const { Schema, SparseVectorIndexConfig } = await import("../src/schema");
+    const schema = new Schema().createIndex(
+      new SparseVectorIndexConfig({
+        sourceKey: "raw_text",
+        embeddingFunction: sparseEf,
+      }),
+      "sparse_metadata",
+    );
+
+    let capturedBody: any;
+    const mockChromaClient = {
+      getMaxBatchSize: jest.fn<() => Promise<number>>().mockResolvedValue(1000),
+      supportsBase64Encoding: jest.fn<() => Promise<boolean>>().mockResolvedValue(false),
+      _path: jest.fn<() => Promise<{ path: string; tenant: string; database: string }>>().mockResolvedValue({ path: "/api/v1", tenant: "default_tenant", database: "default_database" }),
+    };
+
+    const mockApiClient = {
+      post: jest.fn().mockImplementation(async (options: any) => {
+        capturedBody = options.body;
+        return {
+          data: {
+            ids: [],
+            documents: [],
+            embeddings: [],
+            metadatas: [],
+            scores: [],
+            select: [],
+          } as SearchResponse,
+        };
+      }),
+    };
+
+    const collection = new CollectionImpl({
+      chromaClient: mockChromaClient as unknown as ChromaClient,
+      apiClient: mockApiClient as any,
+      id: "col-id",
+      name: "test",
+      configuration: {} as CollectionConfiguration,
+      metadata: undefined,
+      embeddingFunction: undefined,
+      schema,
+    });
+
+    await collection.search(
+      new Search().rank(Knn({ key: "sparse_metadata", query: queryText, limit: 10 })),
+    );
+
+    expect(mockApiClient.post).toHaveBeenCalledTimes(1);
+    expect(generateSpy).toHaveBeenCalledTimes(1);
+    expect(generateSpy).toHaveBeenCalledWith([queryText]);
+
+    expect(capturedBody).toBeDefined();
+    expect(Array.isArray(capturedBody.searches)).toBe(true);
+    expect(capturedBody.searches).toHaveLength(1);
+
+    const knnPayload = capturedBody.searches[0].rank.$knn;
+    expect(knnPayload.query).toEqual({ indices: [0], values: [11.0] });
+    expect(knnPayload.key).toBe("sparse_metadata");
+    expect(knnPayload.limit).toBe(10);
+  });
 });