add_task/remove_task

Author: tanujnay112

chromadb/api/__init__.py

@@ -768,5 +768,55 @@ def _delete(
         where_document: Optional[WhereDocument] = None,
         tenant: str = DEFAULT_TENANT,
         database: str = DEFAULT_DATABASE,
-    ) -> None:
+    ) -> IDs:
+        pass
+
+    @abstractmethod
+    def create_task(
+        self,
+        task_name: str,
+        operator_id: str,
+        input_collection_name: str,
+        output_collection_name: str,
+        params: Optional[str] = None,
+        tenant: str = DEFAULT_TENANT,
+        database: str = DEFAULT_DATABASE,
+    ) -> tuple[bool, str]:
+        """Create a recurring task on a collection.
+
+        Args:
+            task_name: Unique name for this task instance
+            operator_id: Built-in operator identifier
+            input_collection_name: Source collection that triggers the task
+            output_collection_name: Target collection where task output is stored
+            params: Optional JSON string with operator-specific parameters
+            tenant: The tenant name
+            database: The database name
+
+        Returns:
+            tuple: (success: bool, task_id: str)
+        """
+        pass
+
+    @abstractmethod
+    def remove_task(
+        self,
+        task_name: str,
+        input_collection_name: str,
+        delete_output: bool = False,
+        tenant: str = DEFAULT_TENANT,
+        database: str = DEFAULT_DATABASE,
+    ) -> bool:
+        """Delete a task and prevent any further runs.
+
+        Args:
+            task_name: Name of the task to remove
+            input_collection_name: Name of the input collection the task is registered on
+            delete_output: Whether to also delete the output collection
+            tenant: The tenant name
+            database: The database name
+
+        Returns:
+            bool: True if successful
+        """
         pass

chromadb/api/fastapi.py

@@ -695,3 +695,64 @@ def get_max_batch_size(self) -> int:
         pre_flight_checks = self.get_pre_flight_checks()
         max_batch_size = cast(int, pre_flight_checks.get("max_batch_size", -1))
         return max_batch_size
+
+    @trace_method("FastAPI.create_task", OpenTelemetryGranularity.ALL)
+    @override
+    def create_task(
+        self,
+        task_name: str,
+        operator_id: str,
+        input_collection_name: str,
+        output_collection_name: str,
+        params: Optional[str] = None,
+        tenant: str = DEFAULT_TENANT,
+        database: str = DEFAULT_DATABASE,
+    ) -> tuple[bool, str]:
+        """Register a recurring task on a collection."""
+        # Get collection ID from name
+        collection = self.get_collection(
+            input_collection_name, tenant=tenant, database=database
+        )
+
+        resp_json = self._make_request(
+            "post",
+            f"/tenants/{tenant}/databases/{database}/collections/{collection.id}/tasks/create",
+            json={
+                "tenant_id": tenant,
+                "database_name": database,
+                "task_name": task_name,
+                "operator_id": operator_id,
+                "input_collection_name": input_collection_name,
+                "output_collection_name": output_collection_name,
+                "params": params,
+            },
+        )
+        return cast(bool, resp_json["success"]), cast(str, resp_json["task_id"])
+
+    @trace_method("FastAPI.remove_task", OpenTelemetryGranularity.ALL)
+    @override
+    def remove_task(
+        self,
+        task_name: str,
+        input_collection_name: str,
+        delete_output: bool = False,
+        tenant: str = DEFAULT_TENANT,
+        database: str = DEFAULT_DATABASE,
+    ) -> bool:
+        """Delete a task and prevent any further runs."""
+        # Get collection ID from name
+        collection = self.get_collection(
+            input_collection_name, tenant=tenant, database=database
+        )
+
+        resp_json = self._make_request(
+            "post",
+            f"/tenants/{tenant}/databases/{database}/collections/{collection.id}/tasks/delete",
+            json={
+                "tenant_id": tenant,
+                "database_name": database,
+                "task_name": task_name,
+                "delete_output": delete_output,
+            },
+        )
+        return cast(bool, resp_json["success"])

chromadb/api/models/Collection.py

@@ -327,29 +327,29 @@ def search(
             from chromadb.execution.expression import (
                 Search, Key, K, Knn, Val
             )
-            
+
             # Note: K is an alias for Key, so K.DOCUMENT == Key.DOCUMENT
             search = (Search()
                 .where((K("category") == "science") & (K("score") > 0.5))
                 .rank(Knn(query=[0.1, 0.2, 0.3]) * 0.8 + Val(0.5) * 0.2)
                 .limit(10, offset=0)
                 .select(K.DOCUMENT, K.SCORE, "title"))
-            
+
             # Direct construction
             from chromadb.execution.expression import (
                 Search, Eq, And, Gt, Knn, Limit, Select, Key
             )
-            
+
             search = Search(
                 where=And([Eq("category", "science"), Gt("score", 0.5)]),
                 rank=Knn(query=[0.1, 0.2, 0.3]),
                 limit=Limit(offset=0, limit=10),
                 select=Select(keys={Key.DOCUMENT, Key.SCORE, "title"})
             )
-            
+
             # Single search
             result = collection.search(search)
-            
+
             # Multiple searches at once
             searches = [
                 Search().where(K("type") == "article").rank(Knn(query=[0.1, 0.2])),
@@ -490,3 +490,64 @@ def delete(
             tenant=self.tenant,
             database=self.database,
         )
+
+    def create_task(
+        self,
+        name: str,
+        operator_id: str,
+        output_collection: str,
+        params: Optional[str] = None,
+    ) -> tuple[bool, str]:
+        """Create a recurring task that processes this collection.
+
+        Args:
+            name: Unique name for this task instance
+            operator_id: Built-in operator identifier (e.g., "record_counter")
+            output_collection: Name of the collection where task output will be stored
+            params: Optional JSON string with operator-specific parameters
+
+        Returns:
+            tuple: (success: bool, task_id: str)
+
+        Example:
+            >>> success, task_id = collection.create_task(
+            ...     name="count_docs",
+            ...     operator_id="record_counter",
+            ...     output_collection="doc_counts",
+            ...     params=None
+            ... )
+        """
+        return self._client.create_task(
+            task_name=name,
+            operator_id=operator_id,
+            input_collection_name=self.name,
+            output_collection_name=output_collection,
+            params=params,
+            tenant=self.tenant,
+            database=self.database,
+        )
+
+    def remove_task(
+        self,
+        name: str,
+        delete_output: bool = False,
+    ) -> bool:
+        """Delete a task and prevent any further runs.
+
+        Args:
+            name: Name of the task to remove
+            delete_output: Whether to also delete the output collection. Defaults to False.
+
+        Returns:
+            bool: True if successful
+
+        Example:
+            >>> success = collection.remove_task("count_docs", delete_output=True)
+        """
+        return self._client.remove_task(
+            task_name=name,
+            input_collection_name=self.name,
+            delete_output=delete_output,
+            tenant=self.tenant,
+            database=self.database,
+        )

chromadb/api/rust.py

@@ -320,9 +320,7 @@ def _search(
         tenant: str = DEFAULT_TENANT,
         database: str = DEFAULT_DATABASE,
     ) -> SearchResult:
-        raise NotImplementedError(
-            "Search is not implemented for Local Chroma"
-        )
+        raise NotImplementedError("Search is not implemented for Local Chroma")
 
     @override
     def _count(
@@ -583,6 +581,38 @@ def get_settings(self) -> Settings:
     def get_max_batch_size(self) -> int:
         return self.bindings.get_max_batch_size()
 
+    @override
+    def create_task(
+        self,
+        task_name: str,
+        operator_id: str,
+        input_collection_name: str,
+        output_collection_name: str,
+        params: Optional[str] = None,
+        tenant: str = DEFAULT_TENANT,
+        database: str = DEFAULT_DATABASE,
+    ) -> tuple[bool, str]:
+        """Tasks are not supported in the Rust bindings (local embedded mode)."""
+        raise NotImplementedError(
+            "Tasks are only supported when connecting to a Chroma server via HttpClient. "
+            "The Rust bindings (embedded mode) do not support task operations."
+        )
+
+    @override
+    def remove_task(
+        self,
+        task_name: str,
+        input_collection_name: str,
+        delete_output: bool = False,
+        tenant: str = DEFAULT_TENANT,
+        database: str = DEFAULT_DATABASE,
+    ) -> bool:
+        """Tasks are not supported in the Rust bindings (local embedded mode)."""
+        raise NotImplementedError(
+            "Tasks are only supported when connecting to a Chroma server via HttpClient. "
+            "The Rust bindings (embedded mode) do not support task operations."
+        )
+
     # TODO: Remove this if it's not planned to be used
     @override
     def get_user_identity(self) -> UserIdentity:

chromadb/api/segment.py

@@ -902,6 +902,38 @@ def get_settings(self) -> Settings:
     def get_max_batch_size(self) -> int:
         return self._producer.max_batch_size
 
+    @override
+    def create_task(
+        self,
+        task_name: str,
+        operator_id: str,
+        input_collection_name: str,
+        output_collection_name: str,
+        params: Optional[str] = None,
+        tenant: str = DEFAULT_TENANT,
+        database: str = DEFAULT_DATABASE,
+    ) -> tuple[bool, str]:
+        """Tasks are not supported in the Segment API (local embedded mode)."""
+        raise NotImplementedError(
+            "Tasks are only supported when connecting to a Chroma server via HttpClient. "
+            "The Segment API (embedded mode) does not support task operations."
+        )
+
+    @override
+    def remove_task(
+        self,
+        task_name: str,
+        input_collection_name: str,
+        delete_output: bool = False,
+        tenant: str = DEFAULT_TENANT,
+        database: str = DEFAULT_DATABASE,
+    ) -> bool:
+        """Tasks are not supported in the Segment API (local embedded mode)."""
+        raise NotImplementedError(
+            "Tasks are only supported when connecting to a Chroma server via HttpClient. "
+            "The Segment API (embedded mode) does not support task operations."
+        )
+
     # TODO: This could potentially cause race conditions in a distributed version of the
     # system, since the cache is only local.
     # TODO: promote collection -> topic to a base class method so that it can be