[ENH] Generate IDs when not given in add

chromadb/api/__init__.py

@@ -24,6 +24,7 @@
     URIs,
     Where,
     QueryResult,
+    AddResult,
     GetResult,
     WhereDocument,
 )
@@ -115,13 +116,13 @@ def delete_collection(
     @abstractmethod
     def _add(
         self,
-        ids: IDs,
         collection_id: UUID,
         embeddings: Embeddings,
+        ids: Optional[IDs] = None,
         metadatas: Optional[Metadatas] = None,
         documents: Optional[Documents] = None,
         uris: Optional[URIs] = None,
-    ) -> bool:
+    ) -> AddResult:
         """[Internal] Add embeddings to a collection specified by UUID.
         If (some) ids already exist, only the new embeddings will be added.
 

chromadb/api/async_api.py

@@ -24,6 +24,7 @@
     Where,
     QueryResult,
     GetResult,
+    AddResult,
     WhereDocument,
 )
 from chromadb.config import Component, Settings
@@ -106,13 +107,13 @@ async def delete_collection(
     @abstractmethod
     async def _add(
         self,
-        ids: IDs,
         collection_id: UUID,
         embeddings: Embeddings,
+        ids: Optional[IDs] = None,
         metadatas: Optional[Metadatas] = None,
         documents: Optional[Documents] = None,
         uris: Optional[URIs] = None,
-    ) -> bool:
+    ) -> AddResult:
         """[Internal] Add embeddings to a collection specified by UUID.
         If (some) ids already exist, only the new embeddings will be added.
 

chromadb/api/async_client.py

@@ -14,6 +14,7 @@
     EmbeddingFunction,
     Embeddings,
     GetResult,
+    AddResult,
     IDs,
     Include,
     Loadable,
@@ -260,13 +261,13 @@ async def delete_collection(
     @override
     async def _add(
         self,
-        ids: IDs,
         collection_id: UUID,
         embeddings: Embeddings,
+        ids: Optional[IDs] = None,
         metadatas: Optional[Metadatas] = None,
         documents: Optional[Documents] = None,
         uris: Optional[URIs] = None,
-    ) -> bool:
+    ) -> AddResult:
         return await self._server._add(
             ids=ids,
             collection_id=collection_id,

chromadb/api/async_fastapi.py

@@ -2,7 +2,7 @@
 from uuid import UUID
 import urllib.parse
 import orjson
-from typing import Any, Optional, cast, Tuple, Sequence, Dict
+from typing import Any, Optional, cast, Sequence, Dict
 import logging
 import httpx
 from overrides import override
@@ -30,9 +30,11 @@
     Where,
     WhereDocument,
     GetResult,
+    AddResult,
     QueryResult,
     CollectionMetadata,
-    validate_batch,
+    validate_batch_size,
+    RecordSet,
 )
 
 
@@ -411,16 +413,10 @@ async def _delete(
 
         return cast(IDs, resp_json)
 
-    @trace_method("AsyncFastAPI._submit_batch", OpenTelemetryGranularity.ALL)
-    async def _submit_batch(
+    @trace_method("AsyncFastAPI._submit_record_set", OpenTelemetryGranularity.ALL)
+    async def _submit_record_set(
         self,
-        batch: Tuple[
-            IDs,
-            Optional[Embeddings],
-            Optional[Metadatas],
-            Optional[Documents],
-            Optional[URIs],
-        ],
+        record_set: RecordSet,
         url: str,
     ) -> Any:
         """
@@ -430,29 +426,53 @@ async def _submit_batch(
             "post",
             url,
             json={
-                "ids": batch[0],
-                "embeddings": batch[1],
-                "metadatas": batch[2],
-                "documents": batch[3],
-                "uris": batch[4],
+                "ids": record_set["ids"],
+                "embeddings": record_set["embeddings"],
+                "metadatas": record_set["metadatas"],
+                "documents": record_set["documents"],
+                "uris": record_set["uris"],
             },
         )
 
     @trace_method("AsyncFastAPI._add", OpenTelemetryGranularity.ALL)
     @override
     async def _add(
         self,
-        ids: IDs,
         collection_id: UUID,
         embeddings: Embeddings,
+        ids: Optional[IDs] = None,
         metadatas: Optional[Metadatas] = None,
         documents: Optional[Documents] = None,
         uris: Optional[URIs] = None,
-    ) -> bool:
-        batch = (ids, embeddings, metadatas, documents, uris)
-        validate_batch(batch, {"max_batch_size": await self.get_max_batch_size()})
-        await self._submit_batch(batch, "/collections/" + str(collection_id) + "/add")
-        return True
+    ) -> AddResult:
+        record_set: RecordSet = {
+            "ids": ids,
+            "embeddings": embeddings,
+            "metadatas": metadatas,
+            "documents": documents,
+            "uris": uris,
+            "images": None,
+        }
+
+        validate_batch_size(
+            record_set, {"max_batch_size": await self.get_max_batch_size()}
+        )
+
+        resp_json = await self._make_request(
+            "post",
+            "/collections/" + str(collection_id) + "/add",
+            json={
+                "ids": record_set["ids"],
+                "embeddings": record_set["embeddings"],
+                "metadatas": record_set["metadatas"],
+                "documents": record_set["documents"],
+                "uris": record_set["uris"],
+            },
+        )
+
+        return AddResult(
+            ids=resp_json["ids"],
+        )
 
     @trace_method("AsyncFastAPI._update", OpenTelemetryGranularity.ALL)
     @override
@@ -465,11 +485,21 @@ async def _update(
         documents: Optional[Documents] = None,
         uris: Optional[URIs] = None,
     ) -> bool:
-        batch = (ids, embeddings, metadatas, documents, uris)
-        validate_batch(batch, {"max_batch_size": await self.get_max_batch_size()})
+        record_set: RecordSet = {
+            "ids": ids,
+            "embeddings": embeddings,
+            "metadatas": metadatas,
+            "documents": documents,
+            "uris": uris,
+            "images": None,
+        }
+
+        validate_batch_size(
+            record_set, {"max_batch_size": await self.get_max_batch_size()}
+        )
 
-        await self._submit_batch(
-            batch, "/collections/" + str(collection_id) + "/update"
+        await self._submit_record_set(
+            record_set, "/collections/" + str(collection_id) + "/update"
         )
 
         return True
@@ -485,11 +515,23 @@ async def _upsert(
         documents: Optional[Documents] = None,
         uris: Optional[URIs] = None,
     ) -> bool:
-        batch = (ids, embeddings, metadatas, documents, uris)
-        validate_batch(batch, {"max_batch_size": await self.get_max_batch_size()})
-        await self._submit_batch(
-            batch, "/collections/" + str(collection_id) + "/upsert"
+        record_set: RecordSet = {
+            "ids": ids,
+            "embeddings": embeddings,
+            "metadatas": metadatas,
+            "documents": documents,
+            "uris": uris,
+            "images": None,
+        }
+
+        validate_batch_size(
+            record_set, {"max_batch_size": await self.get_max_batch_size()}
+        )
+
+        await self._submit_record_set(
+            record_set, "/collections/" + str(collection_id) + "/upsert"
         )
+
         return True
 
     @trace_method("AsyncFastAPI._query", OpenTelemetryGranularity.ALL)

chromadb/api/client.py

@@ -19,6 +19,7 @@
     Loadable,
     Metadatas,
     QueryResult,
+    AddResult,
     URIs,
 )
 from chromadb.config import Settings, System
@@ -208,13 +209,13 @@ def delete_collection(
     @override
     def _add(
         self,
-        ids: IDs,
         collection_id: UUID,
         embeddings: Embeddings,
+        ids: Optional[IDs] = None,
         metadatas: Optional[Metadatas] = None,
         documents: Optional[Documents] = None,
         uris: Optional[URIs] = None,
-    ) -> bool:
+    ) -> AddResult:
         return self._server._add(
             ids=ids,
             collection_id=collection_id,