[ENH]: Support Key object in schema

[sanketkedia]

Description of changes

Summarize the changes made by this PR.

• Improvements & Bug fixes
  • Support Key type (introduced by search api) in schema. Users can specify this type as schema keys as well as source_key
  • Adds validation for allowed values for key and source_key in schema on both client and server side
  • Adds a small helper to pretty print schema json
• New functionality
  • ...

Test plan

Added plenty of unit tests and end to end tests

• Tests pass locally with pytest for python, yarn test for js, cargo test for rust

Migration plan

None

Observability plan

Staging

Documentation Changes

None

[sanketkedia]

• [ENH]: modify collection updates schema #5730
• [ENH]: Support Key object in schema #5721 👈 (View in Graphite)
• 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]

Schema Enhancements: Support for Key Object, Robust Validation, and Pretty Printing

This PR introduces full support for using the Key object (as introduced in the search API) within the schema system, enabling users to specify Key instances instead of raw strings both as schema keys and as source_key for index configurations. It includes thorough client-side and server-side validation to enforce allowed patterns for keys and source_key across Python and Rust, preventing misuse-such as use of keys beginning with # that aren't reserved for system use (except for core system keys like #embedding and #document). New unit and end-to-end tests are provided to validate this behavior, including handling of both valid and invalid uses. A new utility for pretty-printing JSON schema representations has also been added.

Key Changes

• Extended Schema.create_index and Schema.delete_index to accept both str and Key types as the key argument, converting Key objects as necessary.
• Updated VectorIndexConfig and SparseVectorIndexConfig to allow both str and Key types for the source_key parameter via Pydantic validators, with validation ensuring only acceptable special keys (i.e., #document) are supported.
• Added robust validation logic in Python (chromadb/api/types.py) and Rust (rust/types/src/validators.rs) to disallow keys (or source_key) starting with #, unless they are core system keys. Rejection errors are raised at both the client and server level.
• Expanded test suites (chromadb/test/api/test_schema.py and chromadb/test/api/test_schema_e2e.py) to ensure all combinations of key usage are correctly supported or rejected, including equivalence between using str and Key, and validation of type and value.
• Added a helper method Collection.pretty_schema for pretty-printing the serialized JSON schema.
• Synchronized schema behavior and validation with server-side Rust logic to prevent acceptance of unsupported key patterns regardless of entry point.

Affected Areas

• chromadb/api/types.py
• chromadb/types.py
• chromadb/test/api/test_schema.py
• chromadb/test/api/test_schema_e2e.py
• rust/types/src/validators.rs

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