diff --git a/chromadb/api/fastapi.py b/chromadb/api/fastapi.py index f53dc5bed28..3a67eaa012d 100644 --- a/chromadb/api/fastapi.py +++ b/chromadb/api/fastapi.py @@ -1,4 +1,5 @@ import json +from typing import Optional, cast, Tuple import logging from typing import Optional, cast, Tuple from typing import Sequence @@ -35,6 +36,7 @@ from urllib.parse import urlparse, urlunparse, quote logger = logging.getLogger(__name__) +from chromadb.utils.batch_utils import create_batches class FastAPI(API): diff --git a/chromadb/utils/batch_utils.py b/chromadb/utils/batch_utils.py index e9688ef38e0..04258e4e97f 100644 --- a/chromadb/utils/batch_utils.py +++ b/chromadb/utils/batch_utils.py @@ -1,6 +1,5 @@ from typing import Optional, Tuple, List -from chromadb.api import API from chromadb.api.types import ( Documents, Embeddings, @@ -10,7 +9,7 @@ def create_batches( - api: API, + max_batch_size: int, ids: IDs, embeddings: Optional[Embeddings] = None, metadatas: Optional[Metadatas] = None, @@ -19,15 +18,15 @@ def create_batches( _batches: List[ Tuple[IDs, Embeddings, Optional[Metadatas], Optional[Documents]] ] = [] - if len(ids) > api.max_batch_size: + if len(ids) > max_batch_size: # create split batches - for i in range(0, len(ids), api.max_batch_size): + for i in range(0, len(ids), max_batch_size): _batches.append( ( # type: ignore - ids[i : i + api.max_batch_size], - embeddings[i : i + api.max_batch_size] if embeddings else None, - metadatas[i : i + api.max_batch_size] if metadatas else None, - documents[i : i + api.max_batch_size] if documents else None, + ids[i : i + max_batch_size], + embeddings[i : i + max_batch_size] if embeddings else None, + metadatas[i : i + max_batch_size] if metadatas else None, + documents[i : i + max_batch_size] if documents else None, ) ) else: diff --git a/docs/CIP_5_Large_Batch_Handling_Improvements.md b/docs/CIP_5_Large_Batch_Handling_Improvements.md index 9b03d080f0f..f0be464a74f 100644 --- a/docs/CIP_5_Large_Batch_Handling_Improvements.md +++ b/docs/CIP_5_Large_Batch_Handling_Improvements.md @@ -56,4 +56,6 @@ New tests: ## **Rejected Alternatives** -N/A +Exposing `max_batch_size` and throwing an exception - We decided against this because submitting +batches (especially large ones) comes with monetary or, at the very least, compute cost. Instead, we want the API to +gracefully handle large batches by splitting them up.