Merge pull request #219 from marqo-ai/reduce-k-value

reduce k value to result count
marqo-ai · Dec 12, 2022 · 538710a · 538710a
2 parents bb54a78 + feef174
commit 538710a
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 7 deletions.
diff --git a/src/marqo/tensor_search/configs.py b/src/marqo/tensor_search/configs.py
@@ -32,6 +32,6 @@ def default_env_vars() -> dict:
     return {
         EnvVars.MARQO_MAX_INDEX_FIELDS: None,
         EnvVars.MARQO_MAX_DOC_BYTES: 100000,
-        EnvVars.MARQO_MAX_RETRIEVABLE_DOCS: None,
+        EnvVars.MARQO_MAX_RETRIEVABLE_DOCS: 10000,
         EnvVars.MARQO_MODELS_TO_PRELOAD: ['hf/all_datasets_v4_MiniLM-L6', "ViT-L/14"]
     }
diff --git a/src/marqo/tensor_search/tensor_search.py b/src/marqo/tensor_search/tensor_search.py
@@ -732,8 +732,8 @@ def search(config: Config, index_name: str, text: str, result_count: int = 3, hi
     """
     max_docs_limit = utils.read_env_vars_and_defaults(EnvVars.MARQO_MAX_RETRIEVABLE_DOCS)
     check_upper = True if max_docs_limit is None else result_count <= int(max_docs_limit)
-    if not(check_upper and result_count >= 0):
-        upper_bound_explanation = ("The search result limit must be between 0 and the "
+    if not(check_upper and result_count > 0):
+        upper_bound_explanation = ("The search result limit must be greater than 0 and less than or equal to the"
                                   f"MARQO_MAX_RETRIEVABLE_DOCS limit of [{max_docs_limit}]. ")
         above_zero_explanation = "The search result limit must be greater than or equal to 0."
         explanation = upper_bound_explanation if max_docs_limit is not None else above_zero_explanation
@@ -963,7 +963,7 @@ def _vector_text_search(
                         "knn": {
                             f"{TensorField.chunks}.{vector_field}": {
                                 "vector": vectorised_text,
-                                "k": k
+                                "k": result_count
                             }
                         }
                     },

diff --git a/tests/tensor_search/test_search.py b/tests/tensor_search/test_search.py
@@ -227,12 +227,21 @@ def test_result_count_validation(self):
             raise AssertionError
         except IllegalRequestedDocCount as e:
             pass
-        # should work with 0
+        try:
+            # should not work with 0
+            search_res = tensor_search.search(
+                config=self.config, index_name=self.index_name_1, text="Exact match hehehe",
+                searchable_attributes=["other field", "Cool Field 1"], return_doc_ids=True, result_count=0
+            )
+            raise AssertionError
+        except IllegalRequestedDocCount as e:
+            pass
+        # should work with 1:
         search_res = tensor_search.search(
             config=self.config, index_name=self.index_name_1, text="Exact match hehehe",
-            searchable_attributes=["other field", "Cool Field 1"], return_doc_ids=True, result_count=0
+            searchable_attributes=["other field", "Cool Field 1"], return_doc_ids=True, result_count=1
         )
-        assert len(search_res["hits"]) == 0
+        assert len(search_res['hits']) >= 1
 
     def test_highlights_tensor(self):
         tensor_search.add_documents(