Wire up APIs

Author: Sicheng Pan

Diff for: chromadb/api/__init__.py

@@ -92,23 +92,6 @@ def _modify(
         """
         pass
 
-    @abstractmethod
-    def _fork(
-        self,
-        collection_id: UUID,
-        new_name: str
-    ) -> CollectionModel:
-        """[Internal] Fork the current collection under a new name. The returning collection should contain identical data to the current collection.
-        This is an experimental API that only works for Hosted Chroma for now.
-
-        Args:
-            new_name: The name of the new collection.
-
-        Returns:
-            Collection: A new collection with the specified name and containing identical data to the current collection.
-        """
-        pass
-
     @abstractmethod
     def delete_collection(
         self,
@@ -478,6 +461,24 @@ def get_or_create_collection(
         """
         pass
 
+    @abstractmethod
+    def _fork(
+        self,
+        collection_id: UUID,
+        new_name: str
+    ) -> Collection:
+        """[Internal] Fork the current collection under a new name. The returning collection should contain identical data to the current collection.
+        This is an experimental API that only works for Hosted Chroma for now.
+
+        Args:
+            new_name: The name of the new collection.
+
+        Returns:
+            Collection: A new collection with the specified name and containing identical data to the current collection.
+        """
+        pass
+
+
     @abstractmethod
     def set_tenant(self, tenant: str, database: str = DEFAULT_DATABASE) -> None:
         """Set the tenant and database for the client. Raises an error if the tenant or
@@ -654,7 +655,6 @@ def _modify(
         pass
 
     @abstractmethod
-    @override
     def _fork(
         self,
         collection_id: UUID,

Diff for: chromadb/api/async_api.py

@@ -86,23 +86,6 @@ async def _modify(
         """
         pass
 
-    @abstractmethod
-    async def _fork(
-        self,
-        collection_id: UUID,
-        new_name: str
-    ) -> CollectionModel:
-        """[Internal] Fork the current collection under a new name. The returning collection should contain identical data to the current collection.
-        This is an experimental API that only works for Hosted Chroma for now.
-
-        Args:
-            new_name: The name of the new collection.
-
-        Returns:
-            Collection: A new collection with the specified name and containing identical data to the current collection.
-        """
-        pass
-
     @abstractmethod
     async def delete_collection(
         self,
@@ -472,6 +455,27 @@ async def get_or_create_collection(
         """
         pass
 
+    @abstractmethod
+    async def _fork(
+        self,
+        collection_id: UUID,
+        new_name: str,
+        embedding_function: Optional[
+            EmbeddingFunction[Embeddable]
+        ] = ef.DefaultEmbeddingFunction(),  # type: ignore
+        data_loader: Optional[DataLoader[Loadable]] = None,
+    ) -> AsyncCollection:
+        """[Internal] Fork the current collection under a new name. The returning collection should contain identical data to the current collection.
+        This is an experimental API that only works for Hosted Chroma for now.
+
+        Args:
+            new_name: The name of the new collection.
+
+        Returns:
+            Collection: A new collection with the specified name and containing identical data to the current collection.
+        """
+        pass
+
     @abstractmethod
     async def set_tenant(self, tenant: str, database: str = DEFAULT_DATABASE) -> None:
         """Set the tenant and database for the client. Raises an error if the tenant or
@@ -648,7 +652,6 @@ async def _modify(
         pass
 
     @abstractmethod
-    @override
     async def _fork(
         self,
         collection_id: UUID,

Diff for: chromadb/api/async_client.py

@@ -268,6 +268,27 @@ async def _modify(
             database=self.database,
         )
 
+    @override
+    async def _fork(
+        self,
+        collection_id: UUID,
+        new_name: str,
+        embedding_function: Optional[
+            EmbeddingFunction[Embeddable]
+        ] = ef.DefaultEmbeddingFunction(),  # type: ignore
+        data_loader: Optional[DataLoader[Loadable]] = None,
+    ) -> AsyncCollection:
+        model = await self._server._fork(
+            collection_id=collection_id,
+            new_name=new_name,
+        )
+        return AsyncCollection(
+            client=self._server,
+            model=model,
+            embedding_function=embedding_function,
+            data_loader=data_loader,
+        )
+
     @override
     async def delete_collection(
         self,

Diff for: chromadb/api/async_fastapi.py

@@ -410,7 +410,7 @@ async def _fork(
     ) -> CollectionModel:
         resp_json = await self._make_request(
             "post",
-            f"/tenants/{tenant}/databases/{database}/collections/{id}/fork",
+            f"/tenants/{tenant}/databases/{database}/collections/{collection_id}/fork",
             json={"new_name": new_name},
         )
         model = CollectionModel.from_json(resp_json)

Diff for: chromadb/api/client.py

@@ -235,6 +235,27 @@ def _modify(
             new_configuration=new_configuration,
         )
 
+    @override
+    def _fork(
+        self,
+        collection_id: UUID,
+        new_name: str,
+        embedding_function: Optional[
+            EmbeddingFunction[Embeddable]
+        ] = ef.DefaultEmbeddingFunction(),  # type: ignore
+        data_loader: Optional[DataLoader[Loadable]] = None,
+    ) -> Collection:
+        model = self._server._fork(
+            collection_id=collection_id,
+            new_name=new_name,
+        )
+        return Collection(
+            client=self._server,
+            model=model,
+            embedding_function=embedding_function,
+            data_loader=data_loader,
+        )
+
     @override
     def delete_collection(
         self,

Diff for: chromadb/api/fastapi.py

@@ -363,7 +363,7 @@ def _fork(
         """Forks a collection"""
         resp_json = self._make_request(
             "post",
-            f"/tenants/{tenant}/databases/{database}/collections/{id}/fork",
+            f"/tenants/{tenant}/databases/{database}/collections/{collection_id}/fork",
             json={"new_name": new_name},
         )
         model = CollectionModel.from_json(resp_json)