Skip to content

Commit

Permalink
[ENH] Add collection_id to quota_enforcer calls (#3537)
Browse files Browse the repository at this point in the history
## Description of changes

*Summarize the changes made by this PR.*
 - Improvements & Bug fixes
   - None
 - New functionality
- Update calls to `quota_enforcer.enforce` to include `collection_id`
for `ADD` and `UPSERT`

## Test plan
*How are these changes tested?*

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

## Documentation Changes
*Are all docstrings for user-facing APIs updated if required? Do we need
to make documentation changes in the [docs
repository](https://github.com/chroma-core/docs)?*
  • Loading branch information
drewkim authored Jan 22, 2025
1 parent 6a67d88 commit b3f2870
Showing 1 changed file with 3 additions and 1 deletion.
4 changes: 3 additions & 1 deletion chromadb/api/segment.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
from chromadb.config import DEFAULT_DATABASE, DEFAULT_TENANT, Settings, System
from chromadb.db.system import SysDB
from chromadb.quota import QuotaEnforcer, Action
from chromadb.rate_limit import RateLimitEnforcer, AsyncRateLimitEnforcer
from chromadb.rate_limit import RateLimitEnforcer
from chromadb.segment import SegmentManager
from chromadb.execution.executor.abstract import Executor
from chromadb.execution.expression.operator import Scan, Filter, Limit, KNN, Projection
Expand Down Expand Up @@ -447,6 +447,7 @@ def _add(
metadatas=metadatas,
documents=documents,
uris=uris,
collection_id=collection_id,
)

self._producer.submit_embeddings(collection_id, records_to_submit)
Expand Down Expand Up @@ -559,6 +560,7 @@ def _upsert(
metadatas=metadatas,
documents=documents,
uris=uris,
collection_id=collection_id,
)

self._producer.submit_embeddings(collection_id, records_to_submit)
Expand Down

0 comments on commit b3f2870

Please sign in to comment.