Support target country codes

examples/free_text_order.py

@@ -19,7 +19,7 @@ def new_free_text_order(rapi: RapidataClient):
         .feature_flags(
             FeatureFlags().free_text_minimum_characters(15).alert_on_fast_response(5)
         )
-        .target_country_codes(CountryCodes.ENGLISH_SPEAKING)
+        .country_filter(CountryCodes.ENGLISH_SPEAKING)
     ).create()
 
 

rapidata/rapidata_client/order/rapidata_order_builder.py

@@ -3,9 +3,13 @@
 from rapidata.api_client.models.create_order_model_referee import (
     CreateOrderModelReferee,
 )
+from rapidata.api_client.models.create_order_model_user_filters_inner import (
+    CreateOrderModelUserFiltersInner,
+)
 from rapidata.api_client.models.create_order_model_workflow import (
     CreateOrderModelWorkflow,
 )
+from rapidata.api_client.models.country_user_filter_model import CountryUserFilterModel
 from rapidata.rapidata_client.feature_flags import FeatureFlags
 from rapidata.rapidata_client.metadata.base_metadata import Metadata
 from rapidata.rapidata_client.dataset.rapidata_dataset import RapidataDataset
@@ -16,6 +20,7 @@
 from rapidata.rapidata_client.referee import Referee
 from rapidata.service.openapi_service import OpenAPIService
 
+
 class RapidataOrderBuilder:
     """
     Builder object for creating Rapidata orders.
@@ -42,26 +47,30 @@ def __init__(
         self._aggregator: AggregatorType | None = None
         self._validation_set_id: str | None = None
         self._feature_flags: FeatureFlags | None = None
+        self._country_codes: list[str] | None = None
 
-    def create(self, submit=True) -> RapidataOrder:
-        """
-        Actually makes the API calls to create the order based on how the order builder was configures. Returns a RapidataOrder instance based on the created order with order_id and dataset_id.
-
-        :return: The created RapidataOrder instance.
-        :rtype: RapidataOrder
-        :raises ValueError: If no workflow is provided.
-        """
+    def _to_model(self) -> CreateOrderModel:
         if self._workflow is None:
             raise ValueError("You must provide a blueprint to create an order.")
 
         if self._referee is None:
             print("No referee provided, using default NaiveReferee.")
             self._referee = NaiveReferee()
+        if self._country_codes is None:
+            country_filter = None
+        else:
+            country_filter = CountryUserFilterModel(
+                _t="CountryFilter", countries=self._country_codes
+            )
 
-        order_model = CreateOrderModel(
+        return CreateOrderModel(
             orderName=self._name,
             workflow=CreateOrderModelWorkflow(self._workflow.to_model()),
-            userFilters=[],
+            userFilters=(
+                [CreateOrderModelUserFiltersInner(country_filter)]
+                if country_filter
+                else []
+            ),
             referee=CreateOrderModelReferee(self._referee.to_model()),
             validationSetId=self._validation_set_id,
             featureFlags=(
@@ -71,6 +80,16 @@ def create(self, submit=True) -> RapidataOrder:
             ),
         )
 
+    def create(self, submit=True) -> RapidataOrder:
+        """
+        Actually makes the API calls to create the order based on how the order builder was configures. Returns a RapidataOrder instance based on the created order with order_id and dataset_id.
+
+        :return: The created RapidataOrder instance.
+        :rtype: RapidataOrder
+        :raises ValueError: If no workflow is provided.
+        """
+        order_model = self._to_model()
+
         result = self._openapi_service.order_api.order_create_post(
             create_order_model=order_model
         )
@@ -143,7 +162,7 @@ def feature_flags(self, feature_flags: FeatureFlags):
         self._feature_flags = feature_flags
         return self
 
-    def target_country_codes(self, country_codes: list[str]):
+    def country_filter(self, country_codes: list[str]):
         """
         Set the target country codes for the order.
 
@@ -152,12 +171,7 @@ def target_country_codes(self, country_codes: list[str]):
         :return: The updated RapidataOrderBuilder instance.
         :rtype: RapidataOrderBuilder
         """
-        if self._workflow is None:
-            raise ValueError(
-                "You must set the workflow before setting the target country codes."
-            )
-
-        self._workflow.target_country_codes(country_codes)
+        self._country_codes = country_codes
         return self
 
     def aggregator(self, aggregator: AggregatorType):