refactor code to extract out unpacking embeddings set from existing v…

…alidation logic
chroma-core · Aug 27, 2024 · 88fde75 · 88fde75
1 parent bc12ba7
commit 88fde75
Show file tree

Hide file tree

Showing 4 changed files with 244 additions and 214 deletions.
diff --git a/chromadb/api/models/AsyncCollection.py b/chromadb/api/models/AsyncCollection.py
@@ -2,13 +2,16 @@
     TYPE_CHECKING,
     Optional,
     Union,
+    cast,
 )
-import numpy as np
+from numpy.typing import NDArray
 
 from chromadb.api.types import (
     URI,
     CollectionMetadata,
     Embedding,
+    Embeddings,
+    EmbeddingDType,
     Include,
     Metadata,
     Document,
@@ -35,7 +38,7 @@ async def add(
         embeddings: Optional[
             Union[
                 OneOrMany[Embedding],
-                OneOrMany[np.ndarray],
+                OneOrMany[NDArray[EmbeddingDType]],
             ]
         ] = None,
         metadatas: Optional[OneOrMany[Metadata]] = None,
@@ -63,17 +66,23 @@ async def add(
             ValueError: If you provide an id that already exists
 
         """
-        (
+        record_set = self._process_add_request(
             ids,
             embeddings,
             metadatas,
             documents,
+            images,
             uris,
-        ) = self._validate_and_prepare_embedding_set(
-            ids, embeddings, metadatas, documents, images, uris
         )
 
-        await self._client._add(ids, self.id, embeddings, metadatas, documents, uris)
+        await self._client._add(
+            record_set["ids"],
+            self.id,
+            cast(Embeddings, record_set["embeddings"]),
+            record_set["metadatas"],
+            record_set["documents"],
+            record_set["uris"],
+        )
 
     async def count(self) -> int:
         """The total number of embeddings added to the database
@@ -91,7 +100,7 @@ async def get(
         limit: Optional[int] = None,
         offset: Optional[int] = None,
         where_document: Optional[WhereDocument] = None,
-        include: Include = ["metadatas", "documents"],
+        include: Include = ["metadatas", "documents"],  # type: ignore[list-item]
     ) -> GetResult:
         """Get embeddings and their associate data from the data store. If no ids or where filter is provided returns
         all embeddings up to limit starting at offset.
@@ -144,7 +153,7 @@ async def query(
         query_embeddings: Optional[
             Union[
                 OneOrMany[Embedding],
-                OneOrMany[np.ndarray],
+                OneOrMany[NDArray[EmbeddingDType]],
             ]
         ] = None,
         query_texts: Optional[OneOrMany[Document]] = None,
@@ -153,7 +162,7 @@ async def query(
         n_results: int = 10,
         where: Optional[Where] = None,
         where_document: Optional[WhereDocument] = None,
-        include: Include = ["metadatas", "documents", "distances"],
+        include: Include = ["metadatas", "documents", "distances"],  # type: ignore[list-item]
     ) -> QueryResult:
         """Get the n_results nearest neighbor embeddings for provided query_embeddings or query_texts.
 
@@ -232,7 +241,7 @@ async def update(
         embeddings: Optional[
             Union[
                 OneOrMany[Embedding],
-                OneOrMany[np.ndarray],
+                OneOrMany[NDArray[EmbeddingDType]],
             ]
         ] = None,
         metadatas: Optional[OneOrMany[Metadata]] = None,
@@ -251,25 +260,32 @@ async def update(
         Returns:
             None
         """
-        (
+        record_set = self._process_upsert_or_update_request(
             ids,
             embeddings,
             metadatas,
             documents,
+            images,
             uris,
-        ) = self._validate_and_prepare_update_request(
-            ids, embeddings, metadatas, documents, images, uris
+            require_embeddings_or_data=False,
         )
 
-        await self._client._update(self.id, ids, embeddings, metadatas, documents, uris)
+        await self._client._update(
+            self.id,
+            record_set["ids"],
+            cast(Embeddings, record_set["embeddings"]),
+            record_set["metadatas"],
+            record_set["documents"],
+            record_set["uris"],
+        )
 
     async def upsert(
         self,
         ids: OneOrMany[ID],
         embeddings: Optional[
             Union[
                 OneOrMany[Embedding],
-                OneOrMany[np.ndarray],
+                OneOrMany[NDArray[EmbeddingDType]],
             ]
         ] = None,
         metadatas: Optional[OneOrMany[Metadata]] = None,
@@ -288,23 +304,23 @@ async def upsert(
         Returns:
             None
         """
-        (
+        record_set = self._process_upsert_or_update_request(
             ids,
             embeddings,
             metadatas,
             documents,
+            images,
             uris,
-        ) = self._validate_and_prepare_upsert_request(
-            ids, embeddings, metadatas, documents, images, uris
+            require_embeddings_or_data=True,
         )
 
         await self._client._upsert(
             collection_id=self.id,
-            ids=ids,
-            embeddings=embeddings,
-            metadatas=metadatas,
-            documents=documents,
-            uris=uris,
+            ids=record_set["ids"],
+            embeddings=cast(Embeddings, record_set["embeddings"]),
+            metadatas=record_set["metadatas"],
+            documents=record_set["documents"],
+            uris=record_set["uris"],
         )
 
     async def delete(

diff --git a/chromadb/api/models/Collection.py b/chromadb/api/models/Collection.py
@@ -1,11 +1,14 @@
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING, Optional, Union, cast
 import numpy as np
+from numpy.typing import NDArray
 
 from chromadb.api.models.CollectionCommon import CollectionCommon
 from chromadb.api.types import (
     URI,
     CollectionMetadata,
     Embedding,
+    Embeddings,
+    EmbeddingDType,
     Include,
     Metadata,
     Document,
@@ -40,10 +43,10 @@ def count(self) -> int:
     def add(
         self,
         ids: OneOrMany[ID],
-        embeddings: Optional[  # type: ignore[type-arg]
+        embeddings: Optional[
             Union[
                 OneOrMany[Embedding],
-                OneOrMany[np.ndarray],
+                OneOrMany[NDArray[EmbeddingDType]],
             ]
         ] = None,
         metadatas: Optional[OneOrMany[Metadata]] = None,
@@ -71,17 +74,23 @@ def add(
             ValueError: If you provide an id that already exists
 
         """
-        (
+        record_set = self._process_add_request(
             ids,
             embeddings,
             metadatas,
             documents,
+            images,
             uris,
-        ) = self._validate_and_prepare_embedding_set(
-            ids, embeddings, metadatas, documents, images, uris
         )
 
-        self._client._add(ids, self.id, embeddings, metadatas, documents, uris)
+        self._client._add(
+            record_set["ids"],
+            self.id,
+            cast(Embeddings, record_set["embeddings"]),
+            record_set["metadatas"],
+            record_set["documents"],
+            record_set["uris"],
+        )
 
     def get(
         self,
@@ -90,7 +99,7 @@ def get(
         limit: Optional[int] = None,
         offset: Optional[int] = None,
         where_document: Optional[WhereDocument] = None,
-        include: Include = ["metadatas", "documents"],
+        include: Include = ["metadatas", "documents"],  # type: ignore[list-item]
     ) -> GetResult:
         """Get embeddings and their associate data from the data store. If no ids or where filter is provided returns
         all embeddings up to limit starting at offset.
@@ -152,7 +161,7 @@ def query(
         n_results: int = 10,
         where: Optional[Where] = None,
         where_document: Optional[WhereDocument] = None,
-        include: Include = ["metadatas", "documents", "distances"],
+        include: Include = ["metadatas", "documents", "distances"],  # type: ignore[list-item]
     ) -> QueryResult:
         """Get the n_results nearest neighbor embeddings for provided query_embeddings or query_texts.
 
@@ -250,17 +259,24 @@ def update(
         Returns:
             None
         """
-        (
+        record_set = self._process_upsert_or_update_request(
             ids,
             embeddings,
             metadatas,
             documents,
+            images,
             uris,
-        ) = self._validate_and_prepare_update_request(
-            ids, embeddings, metadatas, documents, images, uris
+            require_embeddings_or_data=False,
         )
 
-        self._client._update(self.id, ids, embeddings, metadatas, documents, uris)
+        self._client._update(
+            self.id,
+            record_set["ids"],
+            cast(Embeddings, record_set["embeddings"]),
+            record_set["metadatas"],
+            record_set["documents"],
+            record_set["uris"],
+        )
 
     def upsert(
         self,
@@ -287,23 +303,23 @@ def upsert(
         Returns:
             None
         """
-        (
+        record_set = self._process_upsert_or_update_request(
             ids,
             embeddings,
             metadatas,
             documents,
+            images,
             uris,
-        ) = self._validate_and_prepare_upsert_request(
-            ids, embeddings, metadatas, documents, images, uris
+            require_embeddings_or_data=True,
         )
 
         self._client._upsert(
-            collection_id=self.id,
-            ids=ids,
-            embeddings=embeddings,
-            metadatas=metadatas,
-            documents=documents,
-            uris=uris,
+            self.id,
+            record_set["ids"],
+            cast(Embeddings, record_set["embeddings"]),
+            record_set["metadatas"],
+            record_set["documents"],
+            record_set["uris"],
         )
 
     def delete(