diff --git a/chromadb/telemetry/product/events.py b/chromadb/telemetry/product/events.py index e5f6bc688c1..2e1617af4b6 100644 --- a/chromadb/telemetry/product/events.py +++ b/chromadb/telemetry/product/events.py @@ -26,7 +26,7 @@ def __init__(self, collection_uuid: str, embedding_function: str): class CollectionAddEvent(ProductTelemetryEvent): - max_batch_size: ClassVar[int] = 100 + max_batch_size: ClassVar[int] = 1000 batch_size: int collection_uuid: str add_amount: int @@ -67,6 +67,8 @@ def batch(self, other: "ProductTelemetryEvent") -> "CollectionAddEvent": class CollectionUpdateEvent(ProductTelemetryEvent): + max_batch_size: ClassVar[int] = 100 + batch_size: int collection_uuid: str update_amount: int with_embeddings: int @@ -80,6 +82,7 @@ def __init__( with_embeddings: int, with_metadata: int, with_documents: int, + batch_size: int = 1, ): super().__init__() self.collection_uuid = collection_uuid @@ -87,10 +90,28 @@ def __init__( self.with_embeddings = with_embeddings self.with_metadata = with_metadata self.with_documents = with_documents + self.batch_size = batch_size + @property + def batch_key(self) -> str: + return self.collection_uuid + self.name + + def batch(self, other: "ProductTelemetryEvent") -> "CollectionUpdateEvent": + if not self.batch_key == other.batch_key: + raise ValueError("Cannot batch events") + other = cast(CollectionUpdateEvent, other) + total_amount = self.update_amount + other.update_amount + return CollectionUpdateEvent( + collection_uuid=self.collection_uuid, + update_amount=total_amount, + with_documents=self.with_documents + other.with_documents, + with_metadata=self.with_metadata + other.with_metadata, + with_embeddings=self.with_embeddings + other.with_embeddings, + batch_size=self.batch_size + other.batch_size, + ) class CollectionQueryEvent(ProductTelemetryEvent): - max_batch_size: ClassVar[int] = 20 + max_batch_size: ClassVar[int] = 1000 batch_size: int collection_uuid: str query_amount: int @@ -147,6 +168,8 @@ def batch(self, other: "ProductTelemetryEvent") -> "CollectionQueryEvent": class CollectionGetEvent(ProductTelemetryEvent): + max_batch_size: ClassVar[int] = 100 + batch_size: int collection_uuid: str ids_count: int limit: int @@ -160,6 +183,7 @@ def __init__( limit: int, include_metadata: int, include_documents: int, + batch_size: int = 1, ): super().__init__() self.collection_uuid = collection_uuid @@ -167,6 +191,25 @@ def __init__( self.limit = limit self.include_metadata = include_metadata self.include_documents = include_documents + self.batch_size = batch_size + + @property + def batch_key(self) -> str: + return self.collection_uuid + self.name + str(self.limit) + + def batch(self, other: "ProductTelemetryEvent") -> "CollectionGetEvent": + if not self.batch_key == other.batch_key: + raise ValueError("Cannot batch events") + other = cast(CollectionGetEvent, other) + total_amount = self.ids_count + other.ids_count + return CollectionGetEvent( + collection_uuid=self.collection_uuid, + ids_count=total_amount, + limit=self.limit, + include_metadata=self.include_metadata + other.include_metadata, + include_documents=self.include_documents + other.include_documents, + batch_size=self.batch_size + other.batch_size, + ) class CollectionDeleteEvent(ProductTelemetryEvent):