[ENH]: Change attach_function API to be clearer

Author: tanujnay112

chromadb/__init__.py

@@ -53,6 +53,13 @@
     # Reciprocal Rank Fusion for combining rankings
     Rrf,
 )
+
+# Import attachable functions
+from chromadb.api.functions import (
+    Function,
+    StatisticsFunction,
+    RecordCounterFunction,
+)
 from pathlib import Path
 import os
 
@@ -97,6 +104,10 @@
     "IntInvertedIndexConfig",
     "FloatInvertedIndexConfig",
     "BoolInvertedIndexConfig",
+    # Attachable Functions
+    "Function",
+    "StatisticsFunction",
+    "RecordCounterFunction",
 ]
 
 from chromadb.types import CloudClientArg

chromadb/api/models/Collection.py

@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Optional, Union, List, cast, Dict, Any
+from typing import TYPE_CHECKING, Optional, Union, List, cast
 
 from chromadb.api.models.CollectionCommon import CollectionCommon
 from chromadb.api.types import (
@@ -27,6 +27,7 @@
 
 if TYPE_CHECKING:
     from chromadb.api.models.AttachedFunction import AttachedFunction
+    from chromadb.api.functions import Function
 
 logger = logging.getLogger(__name__)
 
@@ -500,36 +501,34 @@ def delete(
 
     def attach_function(
         self,
-        function_id: str,
+        function: "Function",
         name: str,
         output_collection: str,
-        params: Optional[Dict[str, Any]] = None,
     ) -> "AttachedFunction":
         """Attach a function to this collection.
 
         Args:
-            function_id: Built-in function identifier (e.g., "record_counter")
+            function: The function to attach (e.g. StatisticsFunction(), RecordCounterFunction())
             name: Unique name for this attached function
             output_collection: Name of the collection where function output will be stored
-            params: Optional dictionary with function-specific parameters
 
         Returns:
             AttachedFunction: Object representing the attached function
 
         Example:
+            >>> from chromadb.api.functions import StatisticsFunction
             >>> attached_fn = collection.attach_function(
-            ...     function_id="record_counter",
+            ...     function=StatisticsFunction(),
             ...     name="mycoll_stats_fn",
-            ...     output_collection="mycoll_stats",
-            ...     params={"threshold": 100}
+            ...     output_collection="mycoll_stats"
             ... )
         """
         return self._client.attach_function(
-            function_id=function_id,
+            function_id=function.name,
             name=name,
             input_collection_id=self.id,
             output_collection=output_collection,
-            params=params,
+            params=function.params,
             tenant=self.tenant,
             database=self.database,
         )

chromadb/test/distributed/test_task_api.py

@@ -7,6 +7,7 @@
 
 import pytest
 from chromadb.api.client import Client as ClientCreator
+from chromadb.api.functions import Function, RecordCounterFunction
 from chromadb.config import System
 from chromadb.errors import ChromaError, NotFoundError
 from chromadb.test.utils.wait_for_version_increase import (
@@ -29,10 +30,9 @@ def test_count_function_attach_and_detach(basic_http_client: System) -> None:
 
     # Create a task that counts records in the collection
     attached_fn = collection.attach_function(
+        function=RecordCounterFunction(),
         name="count_my_docs",
-        function_id="record_counter",  # Built-in operator that counts records
         output_collection="my_documents_counts",
-        params=None,
     )
 
     # Verify task creation succeeded
@@ -66,6 +66,14 @@ def test_count_function_attach_and_detach(basic_http_client: System) -> None:
     assert success is True
 
 
+class InvalidFunction(Function):
+    """A function with an invalid name for testing error handling."""
+
+    @property
+    def name(self) -> str:
+        return "nonexistent_function"
+
+
 def test_task_with_invalid_function(basic_http_client: System) -> None:
     """Test that creating a task with an invalid function raises an error"""
     client = ClientCreator.from_system(basic_http_client)
@@ -77,10 +85,9 @@ def test_task_with_invalid_function(basic_http_client: System) -> None:
     # Attempt to create task with non-existent function should raise ChromaError
     with pytest.raises(ChromaError, match="function not found"):
         collection.attach_function(
+            function=InvalidFunction(),
             name="invalid_task",
-            function_id="nonexistent_function",
             output_collection="output_collection",
-            params=None,
         )
 
 
@@ -94,10 +101,9 @@ def test_attach_function_returns_function_name(basic_http_client: System) -> Non
 
     # Attach a function and verify function_name field in response
     attached_fn = collection.attach_function(
+        function=RecordCounterFunction(),
         name="my_counter",
-        function_id="record_counter",
         output_collection="output_collection",
-        params=None,
     )
 
     # Verify the attached function has function_name (not function_id UUID)
@@ -122,10 +128,9 @@ def test_function_multiple_collections(basic_http_client: System) -> None:
     collection1.add(ids=["id1", "id2"], documents=["doc1", "doc2"])
 
     attached_fn1 = collection1.attach_function(
+        function=RecordCounterFunction(),
         name="task_1",
-        function_id="record_counter",
         output_collection="output_1",
-        params=None,
     )
 
     assert attached_fn1 is not None
@@ -135,10 +140,9 @@ def test_function_multiple_collections(basic_http_client: System) -> None:
     collection2.add(ids=["id3", "id4"], documents=["doc3", "doc4"])
 
     attached_fn2 = collection2.attach_function(
+        function=RecordCounterFunction(),
         name="task_2",
-        function_id="record_counter",
         output_collection="output_2",
-        params=None,
     )
 
     assert attached_fn2 is not None
@@ -170,10 +174,9 @@ def test_functions_one_attached_function_per_collection(
 
     # Create first task on the collection
     attached_fn1 = collection.attach_function(
+        function=RecordCounterFunction(),
         name="task_1",
-        function_id="record_counter",
         output_collection="output_1",
-        params=None,
     )
 
     assert attached_fn1 is not None
@@ -182,19 +185,17 @@ def test_functions_one_attached_function_per_collection(
     # (only one attached function allowed per collection)
     with pytest.raises(ChromaError, match="already has an attached function"):
         collection.attach_function(
+            function=RecordCounterFunction(),
             name="task_2",
-            function_id="record_counter",
             output_collection="output_2",
-            params=None,
         )
 
     # Attempt to create a task with the same name but different params should also fail
     with pytest.raises(ChromaError, match="already exists"):
         collection.attach_function(
+            function=RecordCounterFunction(),
             name="task_1",
-            function_id="record_counter",
             output_collection="output_different",  # Different output collection
-            params=None,
         )
 
     # Detach the first function
@@ -205,10 +206,9 @@ def test_functions_one_attached_function_per_collection(
 
     # Now we should be able to attach a new function
     attached_fn2 = collection.attach_function(
+        function=RecordCounterFunction(),
         name="task_2",
-        function_id="record_counter",
         output_collection="output_2",
-        params=None,
     )
 
     assert attached_fn2 is not None
@@ -229,10 +229,9 @@ def test_function_remove_nonexistent(basic_http_client: System) -> None:
     collection = client.create_collection(name="test_collection")
     collection.add(ids=["id1"], documents=["test"])
     attached_fn = collection.attach_function(
+        function=RecordCounterFunction(),
         name="test_function",
-        function_id="record_counter",
         output_collection="output_collection",
-        params=None,
     )
 
     collection.detach_function(attached_fn.name, delete_output_collection=True)

chromadb/utils/statistics.py

@@ -30,6 +30,7 @@
 from collections import defaultdict
 
 from chromadb.api.types import Where
+from chromadb.api.functions import StatisticsFunction
 
 if TYPE_CHECKING:
     from chromadb.api.models.Collection import Collection
@@ -75,10 +76,9 @@ def attach_statistics_function(
         stats_collection_name = f"{collection.name}_statistics"
 
     return collection.attach_function(
+        function=StatisticsFunction(),
         name=get_statistics_fn_name(collection),
-        function_id="statistics",
         output_collection=stats_collection_name,
-        params=None,
     )
 
 

examples/task_api_example.py

@@ -7,6 +7,7 @@
 """
 
 import chromadb
+from chromadb import RecordCounterFunction
 import time
 
 # Connect to Chroma server
@@ -37,10 +38,9 @@
 # Attach a function that counts records in the collection
 # The 'record_counter' function processes each record and outputs {"count": N}
 attached_fn = collection.attach_function(
-    function_id="record_counter",  # Built-in function that counts records
+    function=RecordCounterFunction(),
     name="count_my_docs",
-    output_collection="my_documents_counts",  # Auto-created
-    params=None,  # No additional parameters needed
+    output_collection="my_documents_counts",
 )
 
 print("✅ Function attached successfully!")