[ENH] Embed query strings in search api

[jairad26]

Description of changes

Summarize the changes made by this PR.

• Improvements & Bug fixes
  • This PR adds support to embed string queries within Knn if provided & if embedding function can be found for the given key
• 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?

[jairad26]

Warning

This pull request is not mergeable via GitHub because a downstack PR is open. Once all requirements are satisfied, merge this PR as a stack on Graphite.
Learn more

• [ENH] Embed query strings in search api #5599 👈 (View in Graphite)
• [ENH] Add Schema to js client #5621
• [TST] Add schema e2e tests #5616
• [ENH] Support using schema embedding function for dense vectors, allow document source key for sparse #5617
• [ENH] Add hosted splade embedding function to python and js #5610
• main

This stack of pull requests is managed by Graphite. Learn more about stacking.

[github-actions]

Reviewer Checklist

Please leverage this checklist to ensure your code review is thorough before approving

Testing, Bugs, Errors, Logs, Documentation

• Can you think of any use case in which the code does not behave as intended? Have they been tested?
• Can you think of any inputs or external events that could break the code? Is user input validated and safe? Have they been tested?
• If appropriate, are there adequate property based tests?
• If appropriate, are there adequate unit tests?
• Should any logging, debugging, tracing information be added or removed?
• Are error messages user-friendly?
• Have all documentation changes needed been made?
• Have all non-obvious changes been commented?

System Compatibility

• Are there any potential impacts on other parts of the system or backward compatibility?
• Does this change intersect with any items on our roadmap, and if so, is there a plan for fitting them together?

Quality

• Is this code of a unexpectedly high quality (Readability, Modularity, Intuitiveness)

[propel-code-bot]

Embed Query Strings in Search API (Automatic Query Embedding for Hybrid Search)

This PR enhances the search API to allow users to pass string queries directly to Knn expressions (both dense and sparse) without pre-embedding them. Whenever a Knn object is encountered in a search or hybrid ranking expression and the query field is a string, the system will automatically run that query string through the correct embedding function (based on the embedding key) before executing the search. The functionality is deeply recursive: it traverses not only simple searches but also complex, nested hybrid ranking expressions. Extensive test coverage is added to ensure string query embedding works in flat and nested rank expressions across all supported entry points. Documentation and type hints for these hybrid search APIs have been updated to reflect that string queries are now supported and automatically embedded.

Key Changes

• Added _embed_knn_string_queries, _embed_rank_string_queries, and _embed_search_string_queries helper methods to CollectionCommon, enabling recursive embedding of string queries in Knn and hybrid rank expressions.
• Patched Collection and AsyncCollection so that all hybrid searches now pre-process their Search payloads to embed string queries using appropriate embedding functions before backend execution.
• Improved hybrid search API docs and Knn docstrings to clearly state support for direct string input and automatic embedding.
• Extended test suite (chromadb/test/api/test_schema_e2e.py) with new tests for single and nested string query embedding, including instrumentation to verify that only embed_query on embedding functions is used, and all string queries are embedded before reaching search execution.
• Updated type annotations and argument handling in hybrid search code paths to document and enforce acceptance of both string and pre-embedded vector queries.
• Minor bug fix in test suite for a boolean assertion in RRF test cases.

Affected Areas

• chromadb/api/models/CollectionCommon.py (core embedding logic for search queries)
• chromadb/api/models/Collection.py and chromadb/api/models/AsyncCollection.py (search plumbing)
• chromadb/execution/expression/operator.py (Knn and docs)
• chromadb/test/api/test_schema_e2e.py (test coverage)
• chromadb/test/test_api.py (minor assertion fix)

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

[propel-code-bot]

[BestPractice]

Error handling gap: Multiple embedding function calls (self._embed, self._sparse_embed, embedding_func.embed_query) can raise exceptions but aren't wrapped in try-catch blocks. If any embedding function fails (network issues, model errors, invalid input), the error will propagate uncaught and could crash the application.

Consider wrapping embedding calls:

try:
    embedding = self._embed(input=[query_text], is_query=True)
except Exception as e:
    raise ValueError(f"Failed to embed query '{query_text}': {e}") from e

Context for Agents

[**BestPractice**]

Error handling gap: Multiple embedding function calls (`self._embed`, `self._sparse_embed`, `embedding_func.embed_query`) can raise exceptions but aren't wrapped in try-catch blocks. If any embedding function fails (network issues, model errors, invalid input), the error will propagate uncaught and could crash the application.

Consider wrapping embedding calls:
```python
try:
    embedding = self._embed(input=[query_text], is_query=True)
except Exception as e:
    raise ValueError(f"Failed to embed query '{query_text}': {e}") from e
```

File: chromadb/api/models/CollectionCommon.py
Line: 773

[propel-code-bot]

[BestPractice]

Potential infinite recursion: The method _embed_rank_string_queries recursively processes rank expressions but doesn't have a maximum depth check. For deeply nested or circular rank structures, this could cause a stack overflow.

Consider adding a depth limit:

def _embed_rank_string_queries(self, rank: Any, depth: int = 0, max_depth: int = 100) -> Any:
    if depth > max_depth:
        raise ValueError(f"Maximum recursion depth ({max_depth}) exceeded in rank expression")
    # ... existing logic with depth + 1 passed to recursive calls

Context for Agents

[**BestPractice**]

Potential infinite recursion: The method `_embed_rank_string_queries` recursively processes rank expressions but doesn't have a maximum depth check. For deeply nested or circular rank structures, this could cause a stack overflow.

Consider adding a depth limit:
```python
def _embed_rank_string_queries(self, rank: Any, depth: int = 0, max_depth: int = 100) -> Any:
    if depth > max_depth:
        raise ValueError(f"Maximum recursion depth ({max_depth}) exceeded in rank expression")
    # ... existing logic with depth + 1 passed to recursive calls
```

File: chromadb/api/models/CollectionCommon.py
Line: 868

[propel-code-bot]

[BestPractice]

Potential error handling gap: If _embed_search_string_queries() raises an exception for any search in the list comprehension, the entire operation will fail and no searches will be processed. Consider adding error handling to gracefully handle individual search embedding failures:

embedded_searches = []
for search in searches_list:
    try:
        embedded_searches.append(self._embed_search_string_queries(search))
    except Exception as e:
        logger.warning(f"Failed to embed search: {e}")
        embedded_searches.append(search)  # Use original search as fallback

Context for Agents

[**BestPractice**]

Potential error handling gap: If `_embed_search_string_queries()` raises an exception for any search in the list comprehension, the entire operation will fail and no searches will be processed. Consider adding error handling to gracefully handle individual search embedding failures:

```python
embedded_searches = []
for search in searches_list:
    try:
        embedded_searches.append(self._embed_search_string_queries(search))
    except Exception as e:
        logger.warning(f"Failed to embed search: {e}")
        embedded_searches.append(search)  # Use original search as fallback
```

File: chromadb/api/models/AsyncCollection.py
Line: 362

[propel-code-bot]

[BestPractice]

Same error handling concern applies here. If _embed_search_string_queries() fails for any search in the list, all searches will fail to process. Consider adding individual error handling as suggested for the AsyncCollection version.

Context for Agents

[**BestPractice**]

Same error handling concern applies here. If `_embed_search_string_queries()` fails for any search in the list, all searches will fail to process. Consider adding individual error handling as suggested for the AsyncCollection version.

File: chromadb/api/models/Collection.py
Line: 366

[propel-code-bot]

[BestPractice]

Error handling issue: The try-except AttributeError block only catches AttributeError but the embedding function could fail with other exceptions (network errors, validation errors, etc.). Consider catching broader exceptions:

try:
    embeddings = embedding_func.embed_query(input=[query_text])
except AttributeError:
    # Fallback if embed_query doesn't exist
    embeddings = embedding_func([query_text])
except Exception as e:
    raise ValueError(
        f"Failed to embed string query '{query_text}' using embedding function: {e}"
    ) from e

Context for Agents

[**BestPractice**]

Error handling issue: The `try-except AttributeError` block only catches `AttributeError` but the embedding function could fail with other exceptions (network errors, validation errors, etc.). Consider catching broader exceptions:

```python
try:
    embeddings = embedding_func.embed_query(input=[query_text])
except AttributeError:
    # Fallback if embed_query doesn't exist
    embeddings = embedding_func([query_text])
except Exception as e:
    raise ValueError(
        f"Failed to embed string query '{query_text}' using embedding function: {e}"
    ) from e
```

File: chromadb/api/models/CollectionCommon.py
Line: 843

[propel-code-bot]

[BestPractice]

Error handling gap: Similar to the main embedding case, if self._sparse_embed() fails, the error lacks context about which query and key failed. With multiple searches containing different sparse embedding keys, this makes debugging difficult.

Consider adding context:

Suggested Change

Suggested change

	# Embed the query
	sparse_embedding = self._sparse_embed(
	# Embed the query
	try:
	sparse_embedding = self._sparse_embed(
	input=[query_text],
	sparse_embedding_function=embedding_func,
	is_query=True,
	)
	except Exception as e:
	raise ValueError(
	f"Failed to embed string query '{query_text}' for sparse key '{key}': {e}"
	) from e

⚡ Committable suggestion

Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.

Context for Agents

[**BestPractice**]

Error handling gap: Similar to the main embedding case, if `self._sparse_embed()` fails, the error lacks context about which query and key failed. With multiple searches containing different sparse embedding keys, this makes debugging difficult.

Consider adding context:

<details>
<summary>Suggested Change</summary>

```suggestion
                    # Embed the query
                    try:
                        sparse_embedding = self._sparse_embed(
                            input=[query_text],
                            sparse_embedding_function=embedding_func,
                            is_query=True,
                        )
                    except Exception as e:
                        raise ValueError(
                            f"Failed to embed string query '{query_text}' for sparse key '{key}': {e}"
                        ) from e
```

⚡ **Committable suggestion**

Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation.

</details>

File: chromadb/api/models/CollectionCommon.py
Line: 812

[jairad26]

3 Jobs Failed:

PR checks / all-required-pr-checks-passed failed on "Decide whether the needed jobs succeeded or failed"

[...]
}
EOM
)"
shell: /usr/bin/bash --noprofile --norc -e -o pipefail {0}
env:
  GITHUB_REPO_NAME: chroma-core/chroma
  PYTHONPATH: /home/runner/_work/_actions/re-actors/alls-green/release/v1/src
# ❌ Some of the required to succeed jobs failed 😢😢😢

📝 Job statuses:
📝 python-tests → ❌ failure [required to succeed or be skipped]
📝 python-vulnerability-scan → ✓ success [required to succeed or be skipped]
📝 javascript-client-tests → ✓ success [required to succeed or be skipped]
📝 rust-tests → ✓ success [required to succeed or be skipped]
📝 go-tests → ✓ success [required to succeed or be skipped]
📝 lint → ✓ success [required to succeed]
📝 check-helm-version-bump → ⬜ skipped [required to succeed or be skipped]
📝 delete-helm-comment → ✓ success [required to succeed or be skipped]
Error: Process completed with exit code 1.

PR checks / Python tests / test-rust-bindings (3.9, --ignore-glob 'chromadb/test/property/*' --ignore-glob 'chromadb/test/st... failed on "Test"

[...]
  /opt/hostedtoolcache/Python/3.9.23/x64/lib/python3.9/site-packages/httpx/_config.py:51: DeprecationWarning: `verify=<str>` is deprecated. Use `verify=ssl.create_default_context(cafile=...)` or `verify=ssl.create_default_context(capath=...)` instead.
    warnings.warn(message, DeprecationWarning)

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
============================= slowest 10 durations =============================
16.00s call     chromadb/test/test_multithreaded.py::test_interleaved_add_query[rust_sqlite_ephemeral]
14.11s call     chromadb/test/test_multithreaded.py::test_interleaved_add_query[rust_sqlite_persistent]
11.84s call     chromadb/test/api/test_limit_offset.py::test_get_limit_offset[rust_sqlite_ephemeral]
11.51s call     chromadb/test/ef/test_default_ef.py::test_invalid_sha256
11.39s call     chromadb/test/api/test_limit_offset.py::test_get_limit_offset[rust_sqlite_persistent]
10.98s call     chromadb/test/test_cli.py::test_vacuum_errors_if_locked[rust_sqlite_persistent]
7.70s call     chromadb/test/test_multithreaded.py::test_multithreaded_add[rust_sqlite_ephemeral]
5.90s call     chromadb/test/test_multithreaded.py::test_multithreaded_add[rust_sqlite_persistent]
5.15s setup    chromadb/test/test_api.py::test_ssl_self_signed[fastapi_ssl]
5.08s call     chromadb/test/test_cli.py::test_app
=========================== short test summary info ============================
FAILED chromadb/test/api/test_schema_e2e.py::test_schema_delete_index_and_restore[rust_sqlite_ephemeral] - Failed: DID NOT RAISE <class 'Exception'>
FAILED chromadb/test/api/test_schema_e2e.py::test_schema_delete_index_and_restore[rust_sqlite_persistent] - Failed: DID NOT RAISE <class 'Exception'>
= 2 failed, 490 passed, 60 skipped, 1 xfailed, 246 warnings in 144.69s (0:02:24) =
Error: Process completed with exit code 1.

PR checks / Python tests / test-rust-single-node-integration (3.9, --ignore-glob 'chromadb/test/property/*' --ignore-glob 'c... failed on "Rust Integration Test"

[...]
  /home/runner/_work/chroma/chroma/chromadb/api/async_fastapi.py:306: UserWarning: space foobar is not supported by default. Supported spaces: ['cosine', 'l2', 'ip']
    create_collection_configuration_to_json(configuration, metadata)

chromadb/test/test_api.py::test_ssl_self_signed[fastapi_ssl]
chromadb/test/test_api.py::test_ssl_self_signed[fastapi_ssl]
chromadb/test/test_api.py::test_ssl_self_signed[fastapi_ssl]
chromadb/test/test_api.py::test_ssl_self_signed_without_ssl_verify[fastapi_ssl]
chromadb/test/test_api.py::test_ssl_self_signed_without_ssl_verify[fastapi_ssl]
  /opt/hostedtoolcache/Python/3.9.23/x64/lib/python3.9/site-packages/httpx/_config.py:51: DeprecationWarning: `verify=<str>` is deprecated. Use `verify=ssl.create_default_context(cafile=...)` or `verify=ssl.create_default_context(capath=...)` instead.
    warnings.warn(message, DeprecationWarning)

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
=========================== short test summary info ============================
FAILED chromadb/test/api/test_schema_e2e.py::test_schema_delete_index_and_restore[integration] - Failed: DID NOT RAISE <class 'Exception'>
FAILED chromadb/test/api/test_schema_e2e.py::test_schema_delete_index_and_restore[async_integration] - Failed: DID NOT RAISE <class 'Exception'>
= 2 failed, 485 passed, 63 skipped, 1 xfailed, 353 warnings in 157.44s (0:02:37) =
  2025-10-17T16:25:03.132348Z ERROR  , name: "BatchSpanProcessor.Flush.ExportError", reason: "ExportFailed(Status { code: Unavailable, message: \", detailed error message: dns error: failed to lookup address information: Temporary failure in name resolution\" })", Failed during the export process
    at /home/runner/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/opentelemetry_sdk-0.27.0/src/trace/span_processor.rs:334

Error: Process completed with exit code 1.

---

Summary: 1 successful workflow, 1 failed workflow

• ✅ Check PR Title (1 job succeeded)
• ❌ PR checks (52 jobs succeeded, 3 jobs failed)

Last updated: 2025-10-17 16:48:15 UTC