[ENH] Add python client support to query on subset of IDs

Author: jairad26

chromadb/api/__init__.py

@@ -270,6 +270,7 @@ def _query(
         self,
         collection_id: UUID,
         query_embeddings: Embeddings,
+        ids: Optional[IDs] = None,
         n_results: int = 10,
         where: Optional[Where] = None,
         where_document: Optional[WhereDocument] = None,
@@ -724,6 +725,7 @@ def _query(
         self,
         collection_id: UUID,
         query_embeddings: Embeddings,
+        ids: Optional[IDs] = None,
         n_results: int = 10,
         where: Optional[Where] = None,
         where_document: Optional[WhereDocument] = None,

chromadb/api/async_api.py

@@ -264,6 +264,7 @@ async def _query(
         self,
         collection_id: UUID,
         query_embeddings: Embeddings,
+        ids: Optional[IDs] = None,
         n_results: int = 10,
         where: Optional[Where] = None,
         where_document: Optional[WhereDocument] = None,
@@ -718,6 +719,7 @@ async def _query(
         self,
         collection_id: UUID,
         query_embeddings: Embeddings,
+        ids: Optional[IDs] = None,
         n_results: int = 10,
         where: Optional[Where] = None,
         where_document: Optional[WhereDocument] = None,

chromadb/api/async_client.py

@@ -403,6 +403,7 @@ async def _query(
         self,
         collection_id: UUID,
         query_embeddings: Embeddings,
+        ids: Optional[IDs] = None,
         n_results: int = 10,
         where: Optional[Where] = None,
         where_document: Optional[WhereDocument] = None,
@@ -411,6 +412,7 @@ async def _query(
         return await self._server._query(
             collection_id=collection_id,
             query_embeddings=query_embeddings,
+            ids=ids,
             n_results=n_results,
             where=where,
             where_document=where_document,

chromadb/api/async_fastapi.py

@@ -623,6 +623,7 @@ async def _query(
         self,
         collection_id: UUID,
         query_embeddings: Embeddings,
+        ids: Optional[IDs] = None,
         n_results: int = 10,
         where: Optional[Where] = None,
         where_document: Optional[WhereDocument] = None,

chromadb/api/client.py

@@ -374,13 +374,15 @@ def _query(
         self,
         collection_id: UUID,
         query_embeddings: Embeddings,
+        ids: Optional[IDs] = None,
         n_results: int = 10,
         where: Optional[Where] = None,
         where_document: Optional[WhereDocument] = None,
         include: Include = IncludeMetadataDocumentsDistances,
     ) -> QueryResult:
         return self._server._query(
             collection_id=collection_id,
+            ids=ids,
             tenant=self.tenant,
             database=self.database,
             query_embeddings=query_embeddings,

chromadb/api/fastapi.py

@@ -591,6 +591,7 @@ def _query(
         self,
         collection_id: UUID,
         query_embeddings: Embeddings,
+        ids: Optional[IDs] = None,
         n_results: int = 10,
         where: Optional[Where] = None,
         where_document: Optional[WhereDocument] = None,
@@ -606,6 +607,7 @@ def _query(
             "post",
             f"/tenants/{tenant}/databases/{database}/collections/{collection_id}/query",
             json={
+                "ids": ids,
                 "query_embeddings": convert_np_embeddings_to_list(query_embeddings)
                 if query_embeddings is not None
                 else None,

chromadb/api/models/AsyncCollection.py

@@ -169,6 +169,7 @@ async def query(
         query_texts: Optional[OneOrMany[Document]] = None,
         query_images: Optional[OneOrMany[Image]] = None,
         query_uris: Optional[OneOrMany[URI]] = None,
+        ids: Optional[IDs] = None,
         n_results: int = 10,
         where: Optional[Where] = None,
         where_document: Optional[WhereDocument] = None,
@@ -184,6 +185,7 @@ async def query(
             query_embeddings: The embeddings to get the closes neighbors of. Optional.
             query_texts: The document texts to get the closes neighbors of. Optional.
             query_images: The images to get the closes neighbors of. Optional.
+            ids: A subset of ids to search within. Optional.
             n_results: The number of neighbors to return for each query_embedding or query_texts. Optional.
             where: A Where type dict used to filter results by. E.g. `{"$and": ["color" : "red", "price": {"$gte": 4.20}]}`. Optional.
             where_document: A WhereDocument type dict used to filter by the documents. E.g. `{$contains: {"text": "hello"}}`. Optional.
@@ -213,6 +215,7 @@ async def query(
 
         query_results = await self._client._query(
             collection_id=self.id,
+            ids=ids,
             query_embeddings=query_request["embeddings"],
             n_results=query_request["n_results"],
             where=query_request["where"],

chromadb/api/models/Collection.py

@@ -172,6 +172,7 @@ def query(
         query_texts: Optional[OneOrMany[Document]] = None,
         query_images: Optional[OneOrMany[Image]] = None,
         query_uris: Optional[OneOrMany[URI]] = None,
+        ids: Optional[IDs] = None,
         n_results: int = 10,
         where: Optional[Where] = None,
         where_document: Optional[WhereDocument] = None,
@@ -188,6 +189,7 @@ def query(
             query_texts: The document texts to get the closes neighbors of. Optional.
             query_images: The images to get the closes neighbors of. Optional.
             query_uris: The URIs to be used with data loader. Optional.
+            ids: A subset of ids to search within. Optional.
             n_results: The number of neighbors to return for each query_embedding or query_texts. Optional.
             where: A Where type dict used to filter results by. E.g. `{"$and": [{"color" : "red"}, {"price": {"$gte": 4.20}}]}`. Optional.
             where_document: A WhereDocument type dict used to filter by the documents. E.g. `{$contains: {"text": "hello"}}`. Optional.
@@ -217,6 +219,7 @@ def query(
 
         query_results = self._client._query(
             collection_id=self.id,
+            ids=ids,
             query_embeddings=query_request["embeddings"],
             n_results=query_request["n_results"],
             where=query_request["where"],

chromadb/api/rust.py

@@ -472,6 +472,7 @@ def _query(
         self,
         collection_id: UUID,
         query_embeddings: Embeddings,
+        ids: Optional[IDs] = None,
         n_results: int = 10,
         where: Optional[Where] = None,
         where_document: Optional[WhereDocument] = None,
@@ -496,6 +497,7 @@ def _query(
 
         rust_response = self.bindings.query(
             str(collection_id),
+            ids,
             query_embeddings,
             n_results,
             json.dumps(where) if where else None,

chromadb/api/segment.py

@@ -775,6 +775,7 @@ def _query(
         self,
         collection_id: UUID,
         query_embeddings: Embeddings,
+        ids: Optional[IDs] = None,
         n_results: int = 10,
         where: Optional[Where] = None,
         where_document: Optional[WhereDocument] = None,