update: enhance decision service to handle error states and improve bucketing logic

FarhanAnjum-opti · FarhanAnjum-opti · commit c637878f763e · 2025-07-03T16:31:59.000+06:00
diff --git a/optimizely/bucketer.py b/optimizely/bucketer.py
@@ -28,7 +28,7 @@
 if TYPE_CHECKING:
     # prevent circular dependenacy by skipping import at runtime
     from .project_config import ProjectConfig
-    from .entities import Experiment, Variation, Group
+    from .entities import Experiment, Variation
     from .helpers.types import TrafficAllocation
 
 
@@ -119,6 +119,34 @@ def bucket(
             and array of log messages representing decision making.
      */.
         """
+        variation_id, decide_reasons = self.bucket_to_entity_id(project_config, experiment, user_id, bucketing_id)
+        if variation_id:
+            variation = project_config.get_variation_from_id_by_experiment_id(experiment.id, variation_id)
+            return variation, decide_reasons
+
+        elif not decide_reasons:
+            message = 'Bucketed into an empty traffic range. Returning nil.'
+            project_config.logger.info(message)
+            decide_reasons.append(message)
+
+        return None, decide_reasons
+
+    def bucket_to_entity_id(
+        self, project_config: ProjectConfig,
+        experiment: Experiment, user_id: str, bucketing_id: str
+    ) -> tuple[Optional[str], list[str]]:
+        """
+        For a given experiment and bucketing ID determines variation ID to be shown to user.
+
+        Args:
+            project_config: Instance of ProjectConfig.
+            experiment: The experiment object (used for group/groupPolicy logic if needed).
+            user_id: The user ID string.
+            bucketing_id: The bucketing ID string for the user.
+
+        Returns:
+            Tuple of (entity_id or None, list of decide reasons).
+        """
         decide_reasons: list[str] = []
         if not experiment:
             return None, decide_reasons
@@ -154,77 +182,5 @@ def bucket(
         # Bucket user if not in white-list and in group (if any)
         variation_id = self.find_bucket(project_config, bucketing_id,
                                         experiment.id, experiment.trafficAllocation)
-        if variation_id:
-            variation = project_config.get_variation_from_id_by_experiment_id(experiment.id, variation_id)
-            return variation, decide_reasons
-
-        else:
-            message = 'Bucketed into an empty traffic range. Returning nil.'
-            project_config.logger.info(message)
-            decide_reasons.append(message)
-
-        return None, decide_reasons
-
-    def bucket_to_entity_id(
-        self,
-        bucketing_id: str,
-        experiment: Experiment,
-        traffic_allocations: list[TrafficAllocation],
-        group: Optional[Group] = None
-    ) -> tuple[Optional[str], list[str]]:
-        """
-        Buckets the user and returns the entity ID (for CMAB experiments).
-        Args:
-            bucketing_id: The bucketing ID string for the user.
-            experiment: The experiment object (for group/groupPolicy logic if needed).
-            traffic_allocations: List of traffic allocation dicts (should have 'entity_id' and 'end_of_range' keys).
-            group: (optional) Group object for mutex group support.
 
-        Returns:
-            Tuple of (entity_id or None, list of decide reasons).
-        """
-        decide_reasons = []
-
-        group_id = getattr(experiment, 'groupId', None)
-        if group_id and group and getattr(group, 'policy', None) == 'random':
-            bucket_key = bucketing_id + group_id
-            bucket_val = self._generate_bucket_value(bucket_key)
-            decide_reasons.append(f'Generated group bucket value {bucket_val} for key "{bucket_key}".')
-
-            matched = False
-            for allocation in group.trafficAllocation:
-                end_of_range = allocation['endOfRange']
-                entity_id = allocation['entityId']
-                if bucket_val < end_of_range:
-                    matched = True
-                    if entity_id != experiment.id:
-                        decide_reasons.append(
-                            f'User not bucketed into experiment "{experiment.id}" (got "{entity_id}").'
-                        )
-                        return None, decide_reasons
-                    decide_reasons.append(
-                        f'User is bucketed into experiment "{experiment.id}" within group "{group_id}".'
-                    )
-                    break
-            if not matched:
-                decide_reasons.append(
-                    f'User not bucketed into any experiment in group "{group_id}".'
-                )
-                return None, decide_reasons
-
-        # Main experiment bucketing
-        bucket_key = bucketing_id + experiment.id
-        bucket_val = self._generate_bucket_value(bucket_key)
-        decide_reasons.append(f'Generated experiment bucket value {bucket_val} for key "{bucket_key}".')
-
-        for allocation in traffic_allocations:
-            end_of_range = allocation['endOfRange']
-            entity_id = allocation['entityId']
-            if bucket_val < end_of_range:
-                decide_reasons.append(
-                    f'User bucketed into entity id "{entity_id}".'
-                )
-                return entity_id, decide_reasons
-
-        decide_reasons.append('User not bucketed into any entity id.')
-        return None, decide_reasons
+        return variation_id, decide_reasons
diff --git a/optimizely/decision/optimizely_decision.py b/optimizely/decision/optimizely_decision.py
@@ -48,3 +48,23 @@ def as_json(self) -> dict[str, Any]:
             'user_context': self.user_context.as_json() if self.user_context else None,
             'reasons': self.reasons
         }
+
+    @classmethod
+    def new_error_decision(cls, key: str, user: OptimizelyUserContext, reasons: list[str]) -> OptimizelyDecision:
+        """Create a new OptimizelyDecision representing an error state.
+        Args:
+            key: The flag key
+            user: The user context
+            reasons: List of reasons explaining the error
+        Returns:
+            OptimizelyDecision with error state values
+        """
+        return cls(
+            variation_key=None,
+            enabled=False,
+            variables={},
+            rule_key=None,
+            flag_key=key,
+            user_context=user,
+            reasons=[reasons[-1]] if reasons else []
+        )
diff --git a/optimizely/decision_service.py b/optimizely/decision_service.py
@@ -30,7 +30,6 @@
     # prevent circular dependenacy by skipping import at runtime
     from .project_config import ProjectConfig
     from .logger import Logger
-    from .helpers.types import TrafficAllocation
 
 
 class CmabDecisionResult(TypedDict):
@@ -134,6 +133,7 @@ def _get_decision_for_cmab_experiment(
         project_config: ProjectConfig,
         experiment: entities.Experiment,
         user_context: OptimizelyUserContext,
+        bucketing_id: str,
         options: Optional[Sequence[str]] = None
     ) -> CmabDecisionResult:
         """
@@ -143,14 +143,36 @@ def _get_decision_for_cmab_experiment(
             project_config: Instance of ProjectConfig.
             experiment: The experiment object for which the decision is to be made.
             user_context: The user context containing user id and attributes.
+            bucketing_id: The bucketing ID to use for traffic allocation.
             options: Optional sequence of decide options.
 
         Returns:
             A dictionary containing:
             - "error": Boolean indicating if there was an error.
-            - "result": The CmabDecision result or empty dict if error.
+            - "result": The CmabDecision result or None if error.
             - "reasons": List of strings with reasons or error messages.
         """
+        decide_reasons: list[str] = []
+        user_id = user_context.user_id
+
+        # Check if user is in CMAB traffic allocation
+        bucketed_entity_id, bucket_reasons = self.bucketer.bucket_to_entity_id(
+            project_config, experiment, user_id, bucketing_id
+        )
+        decide_reasons.extend(bucket_reasons)
+
+        if not bucketed_entity_id:
+            message = f'User "{user_context.user_id}" not in CMAB experiment ' \
+                      f'"{experiment.key}" due to traffic allocation.'
+            self.logger.info(message)
+            decide_reasons.append(message)
+            return {
+                "error": False,
+                "result": None,
+                "reasons": decide_reasons,
+            }
+
+        # User is in CMAB allocation, proceed to CMAB decision
         try:
             options_list = list(options) if options is not None else []
             cmab_decision = self.cmab_service.get_decision(
@@ -159,7 +181,7 @@ def _get_decision_for_cmab_experiment(
             return {
                 "error": False,
                 "result": cmab_decision,
-                "reasons": [],
+                "reasons": decide_reasons,
             }
         except Exception as e:
             error_message = Errors.CMAB_FETCH_FAILED_DETAILED.format(
@@ -468,49 +490,25 @@ def get_variation(
         # If so, handle CMAB-specific traffic allocation and decision logic.
         # Otherwise, proceed with standard bucketing logic for non-CMAB experiments.
         if experiment.cmab:
-            CMAB_DUMMY_ENTITY_ID = "$"
-            # Build the CMAB-specific traffic allocation
-            cmab_traffic_allocation: list[TrafficAllocation] = [{
-                "entityId": CMAB_DUMMY_ENTITY_ID,
-                "endOfRange": experiment.cmab['trafficAllocation']
-            }]
-
-            # Check if user is in CMAB traffic allocation
-            group = None
-            if experiment.groupId:
-                group = project_config.get_group(group_id=experiment.groupId)
-            bucketed_entity_id, bucket_reasons = self.bucketer.bucket_to_entity_id(
-                bucketing_id, experiment, cmab_traffic_allocation, group
-            )
-            decide_reasons += bucket_reasons
-            if bucketed_entity_id != CMAB_DUMMY_ENTITY_ID:
-                message = f'User "{user_id}" not in CMAB experiment "{experiment.key}" due to traffic allocation.'
-                self.logger.info(message)
-                decide_reasons.append(message)
-                return {
-                    'cmab_uuid': None,
-                    'error': False,
-                    'reasons': decide_reasons,
-                    'variation': None
-                }
-
-            # User is in CMAB allocation, proceed to CMAB decision
+            experiment.cmab
             cmab_decision_result = self._get_decision_for_cmab_experiment(project_config,
                                                                           experiment,
                                                                           user_context,
+                                                                          bucketing_id,
                                                                           options)
             decide_reasons += cmab_decision_result.get('reasons', [])
             cmab_decision = cmab_decision_result.get('result')
-            if not cmab_decision or cmab_decision_result['error']:
+            if cmab_decision_result['error']:
                 return {
                     'cmab_uuid': None,
                     'error': True,
                     'reasons': decide_reasons,
                     'variation': None
                 }
-            variation_id = cmab_decision['variation_id']
-            cmab_uuid = cmab_decision['cmab_uuid']
-            variation = project_config.get_variation_from_id(experiment_key=experiment.key, variation_id=variation_id)
+            variation_id = cmab_decision['variation_id'] if cmab_decision else None
+            cmab_uuid = cmab_decision['cmab_uuid'] if cmab_decision else None
+            variation = project_config.get_variation_from_id(experiment_key=experiment.key,
+                                                             variation_id=variation_id) if variation_id else None
         else:
             # Bucket the user
             variation, bucket_reasons = self.bucketer.bucket(project_config, experiment, user_id, bucketing_id)
diff --git a/optimizely/optimizely.py b/optimizely/optimizely.py
@@ -1369,14 +1369,19 @@ def _decide_for_keys(
             user_context,
             merged_decide_options
         )
-
+        print("here")
         for i in range(0, len(flags_without_forced_decision)):
             decision = decision_list[i]['decision']
             reasons = decision_list[i]['reasons']
-            # Can catch errors now. Not used as decision logic implicitly handles error decision.
-            # Will be required for impression events
-            # error = decision_list[i]['error']
+            error = decision_list[i]['error']
             flag_key = flags_without_forced_decision[i].key
+            # store error decision against key and remove key from valid keys
+            if error:
+                optimizely_decision = OptimizelyDecision.new_error_decision(flags_without_forced_decision[i].key,
+                                                                            user_context, reasons)
+                decisions[flag_key] = optimizely_decision
+                if flag_key in valid_keys:
+                    valid_keys.remove(flag_key)
             flag_decisions[flag_key] = decision
             decision_reasons_dict[flag_key] += reasons
 
diff --git a/tests/test_decision_service.py b/tests/test_decision_service.py
diff --git a/tests/test_optimizely.py b/tests/test_optimizely.py
