Skip to content

Commit a4cdaca

Browse files
authored
working (#283)
1 parent 35fca79 commit a4cdaca

File tree

3 files changed

+68
-7
lines changed

3 files changed

+68
-7
lines changed

nucleus/__init__.py

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -742,14 +742,23 @@ def create_custom_index(
742742
)
743743

744744
@deprecated("Prefer calling Dataset.delete_custom_index instead.")
745-
def delete_custom_index(self, dataset_id: str):
745+
def delete_custom_index(self, dataset_id: str, image: bool):
746746
# TODO: deprecate in favor of Dataset.delete_custom_index invocation
747747
return self.make_request(
748-
{},
748+
{"image": image},
749749
f"indexing/{dataset_id}",
750750
requests_command=requests.delete,
751751
)
752752

753+
@deprecated("Prefer calling Dataset.set_primary_index instead.")
754+
def set_primary_index(self, dataset_id: str, image: bool, custom: bool):
755+
# TODO: deprecate in favor of Dataset.set_primary_index invocation
756+
return self.make_request(
757+
{"image": image, "custom": custom},
758+
f"indexing/{dataset_id}/setPrimary",
759+
requests_command=requests.post,
760+
)
761+
753762
@deprecated("Prefer calling Dataset.set_continuous_indexing instead.")
754763
def set_continuous_indexing(self, dataset_id: str, enable: bool = True):
755764
# TODO: deprecate in favor of Dataset.set_continuous_indexing invocation

nucleus/dataset.py

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -944,7 +944,7 @@ def create_custom_index(
944944
self._client,
945945
)
946946

947-
def delete_custom_index(self):
947+
def delete_custom_index(self, image: bool = True):
948948
"""Deletes the custom index uploaded to the dataset.
949949
950950
Returns:
@@ -956,7 +956,18 @@ def delete_custom_index(self):
956956
"message": str
957957
}
958958
"""
959-
return self._client.delete_custom_index(self.id)
959+
return self._client.delete_custom_index(self.id, image)
960+
961+
def set_primary_index(self, image: bool = True, custom: bool = False):
962+
"""Sets the primary index used for Autotag and Similarity Search on this dataset.
963+
964+
Returns:
965+
966+
{
967+
"success": bool,
968+
}
969+
"""
970+
return self._client.set_primary_index(self.id, image, custom)
960971

961972
def set_continuous_indexing(self, enable: bool = True):
962973
"""Toggle whether embeddings are automatically generated for new data.
@@ -985,9 +996,10 @@ def set_continuous_indexing(self, enable: bool = True):
985996
MESSAGE_KEY: preprocessed_response[MESSAGE_KEY],
986997
}
987998
if enable:
988-
response[BACKFILL_JOB_KEY] = (
989-
AsyncJob.from_json(preprocessed_response, self._client),
999+
response[BACKFILL_JOB_KEY] = AsyncJob.from_json(
1000+
preprocessed_response, self._client
9901001
)
1002+
9911003
return response
9921004

9931005
def create_image_index(self):

tests/test_indexing.py

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
from nucleus import DatasetItem
44
from nucleus.constants import (
5+
BACKFILL_JOB_KEY,
56
ERROR_PAYLOAD,
67
JOB_ID_KEY,
78
MESSAGE_KEY,
@@ -38,7 +39,31 @@ def dataset(CLIENT):
3839

3940

4041
@pytest.mark.integration
41-
def test_index_integration(dataset):
42+
def test_set_continuous_indexing(dataset):
43+
resp = dataset.set_continuous_indexing(True)
44+
job = resp[BACKFILL_JOB_KEY]
45+
print(job)
46+
assert job
47+
assert job.job_id
48+
assert job.job_last_known_status
49+
assert job.job_type
50+
assert job.job_creation_time
51+
52+
job_status_response = job.status()
53+
assert STATUS_KEY in job_status_response
54+
assert JOB_ID_KEY in job_status_response
55+
assert MESSAGE_KEY in job_status_response
56+
57+
58+
@pytest.mark.integration
59+
def test_set_primary_index(dataset):
60+
dataset.set_continuous_indexing()
61+
resp = dataset.set_primary_index(image=True, custom=False)
62+
assert resp["success"]
63+
64+
65+
@pytest.mark.integration
66+
def test_create_custom_index(dataset):
4267
signed_embeddings_url = TEST_INDEX_EMBEDDINGS_FILE
4368
job = dataset.create_custom_index([signed_embeddings_url], embedding_dim=3)
4469
assert job.job_id
@@ -52,6 +77,21 @@ def test_index_integration(dataset):
5277
assert MESSAGE_KEY in job_status_response
5378

5479

80+
@pytest.mark.integration
81+
def test_create_and_delete_custom_index(dataset):
82+
# Creates image index
83+
resp = dataset.set_continuous_indexing(True)
84+
85+
# Starts custom indexing job
86+
signed_embeddings_url = TEST_INDEX_EMBEDDINGS_FILE
87+
dataset.create_custom_index([signed_embeddings_url], embedding_dim=3)
88+
89+
resp = dataset.set_primary_index(image=True, custom=True)
90+
assert resp["success"]
91+
92+
dataset.delete_custom_index(image=True)
93+
94+
5595
@pytest.mark.skip(reason="Times out consistently")
5696
def test_generate_image_index_integration(dataset):
5797
job = dataset.create_image_index()

0 commit comments

Comments
 (0)