update: implement CMAB support in bucketer and decision service, reve…

…rt OptimizelyFactory
optimizely · FarhanAnjum-opti · Jun 4, 2025 · Jun 4, 2025 · Jun 13, 2025 · Jun 13, 2025
commit f74bc8ca805cf964d63c340795a2e20f3c0a16b7
diff --git a/optimizely/bucketer.py b/optimizely/bucketer.py
@@ -164,3 +164,55 @@ def bucket(
             decide_reasons.append(message)
 
         return None, decide_reasons
+
+    def bucket_to_entity_id(
+        self,
+        bucketing_id: str,
+        experiment: Experiment,
+        traffic_allocations: list,
+        parent_id: Optional[str] = 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).
+            parent_id: (optional) Used for mutex group support; if not supplied, experiment.id is used.
+
+        Returns:
+            Tuple of (entity_id or None, list of decide reasons).
+        """
+        decide_reasons = []
+
+        # If experiment is in a mutually exclusive group with random policy, check group bucketing first
+        group_id = getattr(experiment, 'groupId', None)
+        group_policy = getattr(experiment, 'groupPolicy', None)
+        if group_id and group_policy == 'random':
+            bucketing_key = f"{bucketing_id}{group_id}"
+            bucket_number = self._generate_bucket_value(bucketing_key)
+            # Group traffic allocation would need to be passed in or found here
+            # For now, skipping group-level allocation (you can extend this for mutex groups)
+            decide_reasons.append(f'Checked mutex group allocation for group "{group_id}".')
+
+        # Main bucketing for experiment or CMAB dummy entity
+        parent_id = parent_id or experiment.id
+        bucketing_key = f"{bucketing_id}{parent_id}"
+        bucket_number = self._generate_bucket_value(bucketing_key)
+        decide_reasons.append(
+            f'Assigned bucket {bucket_number} to bucketing ID "{bucketing_id}" for parent "{parent_id}".'
+        )
+
+        for allocation in traffic_allocations:
+            end_of_range = allocation.get("end_of_range") or allocation.get("endOfRange")
+            entity_id = allocation.get("entity_id") or allocation.get("entityId")
+            if end_of_range is not None and bucket_number < end_of_range:
+                decide_reasons.append(
+                    f'User with bucketing ID "{bucketing_id}" bucketed into entity "{entity_id}".'
+                )
+                return entity_id, decide_reasons
+
+        decide_reasons.append(
+            f'User with bucketing ID "{bucketing_id}" not bucketed into any entity.'
+        )
+        return None, decide_reasons
diff --git a/optimizely/decision_service.py b/optimizely/decision_service.py
@@ -12,7 +12,7 @@
 # limitations under the License.
 
 from __future__ import annotations
-from typing import TYPE_CHECKING, NamedTuple, Optional, Sequence
+from typing import TYPE_CHECKING, NamedTuple, Optional, Sequence, List, TypedDict
 
 from . import bucketer
 from . import entities
@@ -23,28 +23,41 @@
 from .helpers import validator
 from .optimizely_user_context import OptimizelyUserContext, UserAttributes
 from .user_profile import UserProfile, UserProfileService, UserProfileTracker
+from .cmab.cmab_service import DefaultCmabService, CmabDecision
+from optimizely.helpers.enums import Errors
 
 if TYPE_CHECKING:
     # prevent circular dependenacy by skipping import at runtime
     from .project_config import ProjectConfig
     from .logger import Logger
 
 
+class CmabDecisionResult(TypedDict):
+    error: bool
+    result: Optional[CmabDecision]
+    reasons: List[str]
+
+
 class Decision(NamedTuple):
     """Named tuple containing selected experiment, variation and source.
     None if no experiment/variation was selected."""
     experiment: Optional[entities.Experiment]
     variation: Optional[entities.Variation]
     source: Optional[str]
+    # cmab_uuid: Optional[str]
 
 
 class DecisionService:
     """ Class encapsulating all decision related capabilities. """
 
-    def __init__(self, logger: Logger, user_profile_service: Optional[UserProfileService]):
+    def __init__(self,
+                 logger: Logger,
+                 user_profile_service: Optional[UserProfileService],
+                 cmab_service: DefaultCmabService):
         self.bucketer = bucketer.Bucketer()
         self.logger = logger
         self.user_profile_service = user_profile_service
+        self.cmab_service = cmab_service
 
         # Map of user IDs to another map of experiments to variations.
         # This contains all the forced variations set by the user
@@ -76,6 +89,48 @@ def _get_bucketing_id(self, user_id: str, attributes: Optional[UserAttributes])
 
         return user_id, decide_reasons
 
+    def _get_decision_for_cmab_experiment(
+        self,
+        project_config: ProjectConfig,
+        experiment: entities.Experiment,
+        user_context: OptimizelyUserContext,
+        options: Optional[Sequence[str]] = None
+    ) -> CmabDecisionResult:
+        """
+        Retrieves a decision for a contextual multi-armed bandit (CMAB) experiment.
+
+        Args:
+            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.
+            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.
+            - "reasons": List of strings with reasons or error messages.
+        """
+        try:
+            options_list = list(options) if options is not None else []
+            cmab_decision = self.cmab_service.get_decision(
+                project_config, user_context, experiment.id, options_list
+            )
+            return {
+                "error": False,
+                "result": cmab_decision,
+                "reasons": [],
+            }
+        except Exception as e:
+            error_message = Errors.CMAB_FETCH_FAILED.format(str(e))
+            if self.logger:
+                self.logger.error(error_message)
+            return {
+                "error": True,
+                "result": None,
+                "reasons": [error_message],
+            }
+
     def set_forced_variation(
         self, project_config: ProjectConfig, experiment_key: str,
         user_id: str, variation_key: Optional[str]
@@ -313,7 +368,7 @@ def get_variation(
             else:
                 self.logger.warning('User profile has invalid format.')
 
-        # Bucket user and store the new decision
+        # Check audience conditions
         audience_conditions = experiment.get_audience_conditions_or_ids()
         user_meets_audience_conditions, reasons_received = audience_helper.does_user_meet_audience_conditions(
             project_config, audience_conditions,
@@ -330,8 +385,42 @@ def get_variation(
         # Determine bucketing ID to be used
         bucketing_id, bucketing_id_reasons = self._get_bucketing_id(user_id, user_context.get_user_attributes())
         decide_reasons += bucketing_id_reasons
-        variation, bucket_reasons = self.bucketer.bucket(project_config, experiment, user_id, bucketing_id)
-        decide_reasons += bucket_reasons
+
+        if experiment.cmab:
+            CMAB_DUMMY_ENTITY_ID = "$"
+            # Build the CMAB-specific traffic allocation
+            cmab_traffic_allocation = [{
+                "entity_id": CMAB_DUMMY_ENTITY_ID,
+                "end_of_range": experiment.cmab['trafficAllocation']
+            }]
+
+            # Check if user is in CMAB traffic allocation
+            bucketed_entity_id, bucket_reasons = self.bucketer.bucket_to_entity_id(
+                bucketing_id, experiment, cmab_traffic_allocation
+            )
+            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 None, decide_reasons
+
+            # User is in CMAB allocation, proceed to CMAB decision
+            decision_variation_value = self._get_decision_for_cmab_experiment(project_config,
+                                                                              experiment,
+                                                                              user_context,
+                                                                              options)
+            decide_reasons += decision_variation_value.get('reasons', [])
+            cmab_decision = decision_variation_value.get('result')
+            if not cmab_decision:
+                return None, decide_reasons
+            variation_id = cmab_decision['variation_id']
+            variation = project_config.get_variation_from_id(experiment_key=experiment.key, variation_id=variation_id)
+        else:
+            # Bucket the user
+            variation, bucket_reasons = self.bucketer.bucket(project_config, experiment, user_id, bucketing_id)
+            decide_reasons += bucket_reasons
+
         if isinstance(variation, entities.Variation):
             message = f'User "{user_id}" is in variation "{variation.key}" of experiment {experiment.key}.'
             self.logger.info(message)

diff --git a/optimizely/optimizely.py b/optimizely/optimizely.py
@@ -44,13 +44,18 @@
 from .optimizely_config import OptimizelyConfig, OptimizelyConfigService
 from .optimizely_user_context import OptimizelyUserContext, UserAttributes
 from .project_config import ProjectConfig
-from .cmab.cmab_service import DefaultCmabService
+from .cmab.cmab_client import DefaultCmabClient, CmabRetryConfig
+from .cmab.cmab_service import DefaultCmabService, CmabCacheValue
+from .odp.lru_cache import LRUCache
 
 if TYPE_CHECKING:
     # prevent circular dependency by skipping import at runtime
     from .user_profile import UserProfileService
     from .helpers.event_tag_utils import EventTags
 
+# Default constants for CMAB cache
+DEFAULT_CMAB_CACHE_TIMEOUT = 30 * 60 * 1000  # 30 minutes in milliseconds
+DEFAULT_CMAB_CACHE_SIZE = 1000
 
 class Optimizely:
     """ Class encapsulating all SDK functionality. """
@@ -71,7 +76,6 @@ def __init__(
             default_decide_options: Optional[list[str]] = None,
             event_processor_options: Optional[dict[str, Any]] = None,
             settings: Optional[OptimizelySdkSettings] = None,
-            cmab_service: Optional[DefaultCmabService] = None
     ) -> None:
         """ Optimizely init method for managing Custom projects.
 
@@ -100,7 +104,6 @@ def __init__(
           default_decide_options: Optional list of decide options used with the decide APIs.
           event_processor_options: Optional dict of options to be passed to the default batch event processor.
           settings: Optional instance of OptimizelySdkSettings for sdk configuration.
-          cmab_service: Optional instance of DefaultCmabService for Contextual Multi-Armed Bandit (CMAB) support.
         """
         self.logger_name = '.'.join([__name__, self.__class__.__name__])
         self.is_valid = True
@@ -172,10 +175,21 @@ def __init__(
         self._setup_odp(self.config_manager.get_sdk_key())
 
         self.event_builder = event_builder.EventBuilder()
-        if cmab_service:
-            cmab_service.logger = self.logger
-        self.cmab_service = cmab_service
-        self.decision_service = decision_service.DecisionService(self.logger, user_profile_service, cmab_service)
+
+        # Initialize CMAB components
+
+        self.cmab_client = DefaultCmabClient(
+            retry_config=CmabRetryConfig(),
+            logger=logger
+        )
+        self.cmab_cache: LRUCache[str, CmabCacheValue] = LRUCache(DEFAULT_CMAB_CACHE_SIZE,
+                                                             DEFAULT_CMAB_CACHE_TIMEOUT)
+        self.cmab_service = DefaultCmabService(
+            cmab_cache=self.cmab_cache,
+            cmab_client=self.cmab_client,
+            logger=self.logger
+        )
+        self.decision_service = decision_service.DecisionService(self.logger, user_profile_service, self.cmab_service)
         self.user_profile_service = user_profile_service
 
     def _validate_instantiation_options(self) -> None:

diff --git a/optimizely/optimizely_factory.py b/optimizely/optimizely_factory.py
@@ -22,18 +22,11 @@
 from .event_dispatcher import EventDispatcher, CustomEventDispatcher
 from .notification_center import NotificationCenter
 from .optimizely import Optimizely
-from .cmab.cmab_client import DefaultCmabClient, CmabRetryConfig
-from .cmab.cmab_service import DefaultCmabService, CmabCacheValue
-from .odp.lru_cache import LRUCache
 
 if TYPE_CHECKING:
     # prevent circular dependenacy by skipping import at runtime
     from .user_profile import UserProfileService
 
-# Default constants for CMAB cache
-DEFAULT_CMAB_CACHE_TIMEOUT = 30 * 60 * 1000  # 30 minutes in milliseconds
-DEFAULT_CMAB_CACHE_SIZE = 1000
-
 
 class OptimizelyFactory:
     """ Optimizely factory to provides basic utility to instantiate the Optimizely
@@ -43,8 +36,6 @@ class OptimizelyFactory:
     max_event_flush_interval: Optional[int] = None
     polling_interval: Optional[float] = None
     blocking_timeout: Optional[int] = None
-    cmab_cache_size: int = DEFAULT_CMAB_CACHE_SIZE
-    cmab_cache_timeout: int = DEFAULT_CMAB_CACHE_TIMEOUT
 
     @staticmethod
     def set_batch_size(batch_size: int) -> int:
@@ -113,36 +104,16 @@ def default_instance(sdk_key: str, datafile: Optional[str] = None) -> Optimizely
             notification_center=notification_center,
         )
 
-        # Initialize CMAB components
-        cmab_client = DefaultCmabClient(
-            retry_config=CmabRetryConfig(),
-            logger=logger
-        )
-        cmab_cache: LRUCache[str, CmabCacheValue] = LRUCache(OptimizelyFactory.cmab_cache_size,
-                                                             OptimizelyFactory.cmab_cache_timeout)
-        cmab_service = DefaultCmabService(
-            cmab_cache=cmab_cache,
-            cmab_client=cmab_client,
-            logger=logger
-        )
-
         optimizely = Optimizely(
             datafile, None, logger, error_handler, None, None, sdk_key, config_manager, notification_center,
-            event_processor, cmab_service=cmab_service
+            event_processor
         )
         return optimizely
 
     @staticmethod
     def default_instance_with_config_manager(config_manager: BaseConfigManager) -> Optimizely:
-        # Initialize CMAB components
-        cmab_client = DefaultCmabClient(retry_config=CmabRetryConfig())
-        cmab_cache: LRUCache[str, CmabCacheValue] = LRUCache(OptimizelyFactory.cmab_cache_size,
-                                                             OptimizelyFactory.cmab_cache_timeout)
-        cmab_service = DefaultCmabService(cmab_cache=cmab_cache, cmab_client=cmab_client)
-
         return Optimizely(
-            config_manager=config_manager,
-            cmab_service=cmab_service
+            config_manager=config_manager
         )
 
     @staticmethod
@@ -203,21 +174,7 @@ def custom_instance(
             notification_center=notification_center,
         )
 
-        # Initialize CMAB components
-        cmab_client = DefaultCmabClient(
-            retry_config=CmabRetryConfig(),
-            logger=logger
-        )
-        cmab_cache: LRUCache[str, CmabCacheValue] = LRUCache(OptimizelyFactory.cmab_cache_size,
-                                                             OptimizelyFactory.cmab_cache_timeout)
-        cmab_service = DefaultCmabService(
-            cmab_cache=cmab_cache,
-            cmab_client=cmab_client,
-            logger=logger
-        )
-
         return Optimizely(
             datafile, event_dispatcher, logger, error_handler, skip_json_validation, user_profile_service,
-            sdk_key, config_manager, notification_center, event_processor, settings=settings,
-            cmab_service=cmab_service
-        )
+            sdk_key, config_manager, notification_center, event_processor, settings=settings
+        )