Support collapse_fields param in both create_index and search endpoints (#296)

* Support collapse fields params in both create_index and search endpoints.

* downgrade minimum version

Author: papa99do

src/marqo/index.py

@@ -101,6 +101,7 @@ def create(config: Config,
                wait_for_readiness: bool = True,
                text_chunk_prefix: Optional[str] = None,
                text_query_prefix: Optional[str] = None,
+               collapse_fields: Optional[List[marqo_index.CollapseField]] = None,
                ) -> Dict[str, Any]:
         """Create the index. Please refer to the marqo cloud to see options for inference and storage node types.
         Creates CreateIndexSettings object and then uses it to create the index.
@@ -137,6 +138,7 @@ def create(config: Config,
             number_of_inferences: number of inferences for the index
             number_of_shards: number of shards for the index
             number_of_replicas: number of replicas for the index
+            collapse_fields: list of fields that can be collapsed on at query time
         Note:
             wait_for_readiness, inference_type, storage_class, number_of_inferences,
             number_of_shards, number_of_replicas are Marqo Cloud specific parameters,
@@ -166,6 +168,7 @@ def create(config: Config,
                 annParameters=ann_parameters,
                 textChunkPrefix=text_chunk_prefix,
                 textQueryPrefix=text_query_prefix,
+                collapseFields=collapse_fields,
             )
             return req.post(f"indexes/{index_name}", body=local_create_index_settings.generate_request_body())
 
@@ -195,6 +198,7 @@ def create(config: Config,
                 storageClass=storage_class,
                 textChunkPrefix=text_chunk_prefix,
                 textQueryPrefix=text_query_prefix,
+                collapseFields=collapse_fields,
             )
 
             response = req.post(f"indexes/{index_name}", body=cloud_index_settings.generate_request_body())
@@ -228,7 +232,8 @@ def search(self, q: Optional[Union[str, dict]] = None, searchable_attributes: Op
                language: Optional[str] = None,
                sort_by: Optional[dict] = None,
                relevance_cutoff: Optional[dict] = None,
-               interpolation_method: Optional[str] = None
+               interpolation_method: Optional[str] = None,
+               collapse_fields: Optional[List[dict]] = None,
                ) -> Dict[str, Any]:
         """Search the index.
 
@@ -275,6 +280,7 @@ def search(self, q: Optional[Union[str, dict]] = None, searchable_attributes: Op
             sort_by: a dictionary of the sort_by parameters to be used for sorting the results
             relevance_cutoff: a dictionary of the relevance cutoff parameters
             interpolation_method: the interpolation method to use for combining query & context embeddings.
+            collapse_fields: a list of fields to collapse/group on
 
         Returns:
             Dictionary with hits and other metadata
@@ -317,7 +323,8 @@ def search(self, q: Optional[Union[str, dict]] = None, searchable_attributes: Op
             "language": language,
             "sortBy": sort_by,
             "relevanceCutoff": relevance_cutoff,
-            "interpolationMethod": interpolation_method
+            "interpolationMethod": interpolation_method,
+            "collapseFields": collapse_fields,
         }
 
         body = {k: v for k, v in body.items() if v is not None}

src/marqo/models/create_index_settings.py

@@ -26,6 +26,7 @@ class IndexSettings(MarqoBaseModel):
         imagePreprocessing: The image preprocessing method to use.
         vectorNumericType: The numeric type of the vector.
         annParameters: The ANN parameters to use.
+        collapseFields: list of fields that can be collapsed on at query time
 
     Please note, we don't note set default values in the py-marqo side. All the
     values are set to be None and will not be sent to Marqo in the HttpRequest.
@@ -50,6 +51,7 @@ class IndexSettings(MarqoBaseModel):
     annParameters: Optional[marqo_index.AnnParameters] = None
     textQueryPrefix: Optional[str] = None
     textChunkPrefix: Optional[str] = None
+    collapseFields: Optional[List[marqo_index.CollapseField]] = None
 
     def generate_request_body(self) -> dict:
         """A json encoded string of the request body"""

src/marqo/models/marqo_index.py

@@ -104,3 +104,8 @@ class FieldRequest(StrictBaseModel):
 class AnnParameters(StrictBaseModel):
     spaceType: Optional[DistanceMetric] = Field(None, alias="space_type")
     parameters: Optional[HnswConfig] = None
+
+
+class CollapseField(StrictBaseModel):
+    name: str
+    minGroups: int = 500

unit_tests/test_index_management_endpoints.py

@@ -0,0 +1,34 @@
+from unittest.mock import patch, MagicMock
+
+from marqo.config import Config
+from marqo.default_instance_mappings import DefaultInstanceMappings
+from marqo.index import Index
+from marqo.models.marqo_index import CollapseField
+from unit_tests.marqo_unit_tests import MarqoUnitTests
+
+
+class TestIndexCreationEndpoint(MarqoUnitTests):
+    @classmethod
+    def setUpClass(cls):
+        """
+        Set up the class by creating a test index.
+        """
+        cls.config = Config(
+            instance_mappings=DefaultInstanceMappings(url="http://unit-tests-url:8882"),
+        )
+
+    @patch('marqo._httprequests.HttpRequests.send_request')
+    def test_collapse_field_parameter(self, mock_send_request):
+        mock_response = MagicMock()
+        mock_send_request.return_value = mock_response
+
+        Index.create(self.config, index_name='test_index',
+                     collapse_fields=[
+                         CollapseField(name='parent_id'),  # test default minGroups value is 500
+                         CollapseField(name='color', minGroups=10),
+                     ])
+
+        self.assertEqual(1, mock_send_request.call_count)
+        body = mock_send_request.call_args[0][2]
+        self.assertListEqual([{'name': 'parent_id', 'minGroups': 500}, {'name': 'color', 'minGroups': 10}],
+                             body["collapseFields"])

unit_tests/test_index_search_endpoint.py

@@ -179,3 +179,30 @@ def test_interpolation_method_parameter_not_exists(self, mock_send_request):
         self.assertEqual(1, mock_send_request.call_count)
         body = mock_send_request.call_args[0][2]
         self.assertNotIn("interpolationMethod", body)
+
+
+class TestIndexSearchEndpointCollapseFields(MarqoUnitTests):
+    """
+    Test class for the Index search endpoint collapse_fields parameter.
+    """
+    @classmethod
+    def setUpClass(cls):
+        """
+        Set up the class by creating a test index.
+        """
+        config = Config(
+            instance_mappings=DefaultInstanceMappings(url="http://unit-tests-url:8882"),
+        )
+        cls.index = Index(config=config, index_name="test_index")
+
+    @patch('marqo._httprequests.HttpRequests.send_request')
+    def test_collapse_field_parameter(self, mock_send_request):
+        mock_response = MagicMock()
+        mock_send_request.return_value = mock_response
+
+        collapse_fields = [{"name": "parent_id"}]
+        self.index.search(q="test", collapse_fields=collapse_fields)
+
+        self.assertEqual(1, mock_send_request.call_count)
+        body = mock_send_request.call_args[0][2]
+        self.assertEqual(collapse_fields, body["collapseFields"])