feat: add exposure

zhukaihan · zhukaihan · commit cad50a196ab1 · 2025-11-13T05:21:23.000-08:00
diff --git a/src/amplitude_experiment/assignment/assignment.py b/src/amplitude_experiment/assignment/assignment.py
@@ -8,6 +8,9 @@
 
 
 class Assignment:
+    """
+    @deprecated Assignment tracking is deprecated. Use Exposure with ExposureService instead.
+    """
 
     def __init__(self, user: User, results: Dict[str, Variant]):
         self.user = user
diff --git a/src/amplitude_experiment/assignment/assignment_config.py b/src/amplitude_experiment/assignment/assignment_config.py
@@ -2,6 +2,9 @@
 
 
 class AssignmentConfig(amplitude.Config):
+    """
+    @deprecated Assignment tracking is deprecated. Use ExposureConfig with ExposureService instead.
+    """
     def __init__(self, cache_capacity: int = 65536, send_evaluated_props: bool = False, **kw):
         super(AssignmentConfig, self).__init__(**kw)
         self.cache_capacity = cache_capacity
diff --git a/src/amplitude_experiment/assignment/assignment_filter.py b/src/amplitude_experiment/assignment/assignment_filter.py
@@ -4,6 +4,9 @@
 
 
 class AssignmentFilter:
+    """
+    @deprecated Assignment tracking is deprecated. Use ExposureFilter with ExposureService instead.
+    """
     def __init__(self, size: int, ttl_millis: int = DAY_MILLIS):
         self.cache = Cache(size, ttl_millis)
 
diff --git a/src/amplitude_experiment/assignment/assignment_service.py b/src/amplitude_experiment/assignment/assignment_service.py
@@ -9,6 +9,9 @@
 
 
 def to_event(assignment: Assignment, send_evaluated_props: bool) -> BaseEvent:
+    """
+    @deprecated Assignment tracking is deprecated. Use Exposure tracking.
+    """
     event = BaseEvent(event_type='[Experiment] Assignment', user_id=assignment.user.user_id,
                       device_id=assignment.user.device_id, event_properties={}, user_properties={})
 
@@ -58,6 +61,9 @@ def to_event(assignment: Assignment, send_evaluated_props: bool) -> BaseEvent:
 
 
 class AssignmentService:
+    """
+    @deprecated Assignment tracking is deprecated. Use ExposureService with Exposure tracking instead.
+    """
     def __init__(self, amplitude: Amplitude, assignment_filter: AssignmentFilter, send_evaluated_props: bool):
         self.amplitude = amplitude
         self.assignmentFilter = assignment_filter
diff --git a/src/amplitude_experiment/exposure/__init__.py b/src/amplitude_experiment/exposure/__init__.py
@@ -0,0 +1,5 @@
+from .exposure import Exposure, DAY_MILLIS
+from .exposure_filter import ExposureFilter
+from .exposure_service import ExposureService, to_exposure_events
+from .exposure_config import ExposureConfig
+
diff --git a/src/amplitude_experiment/exposure/exposure.py b/src/amplitude_experiment/exposure/exposure.py
@@ -0,0 +1,31 @@
+import time
+from typing import Dict
+
+from .. import Variant
+from ..user import User
+
+DAY_MILLIS = 24 * 60 * 60 * 1000
+
+
+class Exposure:
+    """
+    Exposure is a class that represents a user's exposure to a set of flags.
+    """
+
+    def __init__(self, user: User, results: Dict[str, Variant]):
+        self.user = user
+        self.results = results
+        self.timestamp = time.time() * 1000
+
+    def canonicalize(self) -> str:
+        user = self.user.user_id.strip() if self.user.user_id else 'None'
+        device = self.user.device_id.strip() if self.user.device_id else 'None'
+        canonical = user + ' ' + device + ' '
+        for flag_key in sorted(self.results):
+            variant = self.results[flag_key]
+            if variant.key is None:
+                continue
+            value = self.results[flag_key].key.strip()
+            canonical += flag_key.strip() + ' ' + value + ' '
+        return canonical
+
diff --git a/src/amplitude_experiment/exposure/exposure_config.py b/src/amplitude_experiment/exposure/exposure_config.py
@@ -0,0 +1,8 @@
+import amplitude
+
+
+class ExposureConfig(amplitude.Config):
+    def __init__(self, cache_capacity: int = 65536, **kw):
+        super(ExposureConfig, self).__init__(**kw)
+        self.cache_capacity = cache_capacity
+
diff --git a/src/amplitude_experiment/exposure/exposure_filter.py b/src/amplitude_experiment/exposure/exposure_filter.py
@@ -0,0 +1,23 @@
+from .exposure import Exposure
+from .exposure import DAY_MILLIS
+from ..util.cache import Cache
+
+
+class ExposureFilter:
+    def __init__(self, size: int, ttl_millis: int = DAY_MILLIS):
+        self.cache = Cache(size, ttl_millis)
+        self.ttl_millis = ttl_millis
+
+    def should_track(self, exposure: Exposure) -> bool:
+        """
+        Determines if an exposure should be tracked based on the dedupe filter.
+        """
+        if not exposure.results:
+            # Don't track empty exposures.
+            return False
+        canonical_exposure = exposure.canonicalize()
+        track = self.cache.get(canonical_exposure) is None
+        if track:
+            self.cache.put(canonical_exposure, object())
+        return track
+
diff --git a/src/amplitude_experiment/exposure/exposure_service.py b/src/amplitude_experiment/exposure/exposure_service.py
@@ -0,0 +1,82 @@
+from amplitude import Amplitude, BaseEvent
+from .exposure import Exposure
+from .exposure import DAY_MILLIS
+from .exposure_filter import ExposureFilter
+from ..util import hash_code
+
+FLAG_TYPE_MUTUAL_EXCLUSION_GROUP = "mutual-exclusion-group"
+
+
+def to_exposure_events(exposure: Exposure, ttl_millis: int) -> list[BaseEvent]:
+    """
+    Convert an Exposure to a list of Amplitude events (one per flag).
+    """
+    events = []
+    canonicalized = exposure.canonicalize()
+    for flag_key in exposure.results:
+        variant = exposure.results[flag_key]
+
+        # TODO: We don't seem to use trackExposure metadata.
+        track_exposure = variant.metadata.get('trackExposure') if variant.metadata is not None else True
+        if track_exposure is False:
+            continue
+
+        is_default = variant.metadata.get('default') if variant.metadata is not None else False
+        if is_default:
+            continue
+
+        # Determine user properties to set and unset.
+        set_props = {}
+        unset_props = {}
+        flag_type = variant.metadata.get('flagType') if variant.metadata is not None else None
+        if flag_type != FLAG_TYPE_MUTUAL_EXCLUSION_GROUP:
+            if is_default:
+                unset_props[f'[Experiment] {flag_key}'] = '-'
+            else:
+                if variant.key:
+                    set_props[f'[Experiment] {flag_key}'] = variant.key
+                elif variant.value:
+                    set_props[f'[Experiment] {flag_key}'] = variant.value
+
+        # Build event properties.
+        event_properties = {}
+        if not is_default:
+            event_properties['[Experiment] Flag Key'] = flag_key
+            if variant.key:
+                event_properties['[Experiment] Variant'] = variant.key
+            elif variant.value:
+                event_properties['[Experiment] Variant'] = variant.value
+            if variant.metadata:
+                event_properties['metadata'] = variant.metadata
+
+        # Build event.
+        event = BaseEvent(
+            event_type='[Experiment] Exposure',
+            user_id=exposure.user.user_id,
+            device_id=exposure.user.device_id,
+            event_properties=event_properties,
+            user_properties={
+                '$set': set_props,
+                '$unset': unset_props,
+            },
+            insert_id=f'{exposure.user.user_id} {exposure.user.device_id} {hash_code(flag_key + " " + canonicalized)} {int(exposure.timestamp / ttl_millis)}'
+        )
+        if exposure.user.groups:
+            event.groups = exposure.user.groups
+
+        events.append(event)
+
+    return events
+
+
+class ExposureService:
+    def __init__(self, amplitude: Amplitude, exposure_filter: ExposureFilter):
+        self.amplitude = amplitude
+        self.exposure_filter = exposure_filter
+
+    def track(self, exposure: Exposure):
+        if self.exposure_filter.should_track(exposure):
+            events = to_exposure_events(exposure, self.exposure_filter.ttl_millis)
+            for event in events:
+                self.amplitude.track(event)
+
diff --git a/src/amplitude_experiment/local/client.py b/src/amplitude_experiment/local/client.py
@@ -5,7 +5,9 @@
 from amplitude import Amplitude
 
 from .config import LocalEvaluationConfig
+from .evaluate_options import EvaluateOptions
 from ..assignment import Assignment, AssignmentFilter, AssignmentService
+from ..exposure import Exposure, ExposureFilter, ExposureService
 from ..cohort.cohort import USER_GROUP_TYPE
 from ..cohort.cohort_download_api import DirectCohortDownloadApi
 from ..cohort.cohort_loader import CohortLoader
@@ -47,6 +49,13 @@ def __init__(self, api_key: str, config: LocalEvaluationConfig = None):
             instance = Amplitude(config.assignment_config.api_key, config.assignment_config)
             self.assignment_service = AssignmentService(instance, AssignmentFilter(
                 config.assignment_config.cache_capacity), config.assignment_config.send_evaluated_props)
+        
+        # Exposure service is always instantiated, using deployment key if no api key provided
+        exposure_config = self.config.exposure_config
+        if not exposure_config.api_key:
+            exposure_config.api_key = api_key
+        exposure_instance = Amplitude(exposure_config.api_key, exposure_config)
+        self.exposure_service = ExposureService(exposure_instance, ExposureFilter(exposure_config.cache_capacity))
         self.logger = logging.getLogger("Amplitude")
         self.logger.addHandler(logging.StreamHandler())
         if self.config.debug:
@@ -81,7 +90,8 @@ def start(self):
         """
         self.deployment_runner.start()
 
-    def evaluate_v2(self, user: User, flag_keys: Set[str] = None) -> Dict[str, Variant]:
+    # TODO: python backwards compatibility for evaluate_v2 to be looked at again
+    def evaluate_v2(self, user: User, flag_keys: Set[str] = None, options: EvaluateOptions = None) -> Dict[str, Variant]:
         """
         Locally evaluates flag variants for a user.
 
@@ -91,16 +101,23 @@ def evaluate_v2(self, user: User, flag_keys: Set[str] = None) -> Dict[str, Varia
 
             Parameters:
                 user (User): The user to evaluate
-                flag_keys (List[str]): The flags to evaluate with the user. If empty, all flags are evaluated.
+                flag_keys (Set[str]): The flags to evaluate with the user. If empty, all flags are evaluated.
+                options (EvaluateOptions): Optional evaluation options.
 
             Returns:
                 The evaluated variants.
         """
+        # Handle backwards compatibility: if options is None, create default
+        if options is None:
+            options = EvaluateOptions(flag_keys=flag_keys)
+        # Use flag_keys from options if provided, otherwise fall back to parameter
+        flag_keys_to_use = options.flag_keys if options.flag_keys is not None else flag_keys
+        
         flag_configs = self.flag_config_storage.get_flag_configs()
         if flag_configs is None or len(flag_configs) == 0:
             return {}
         self.logger.debug(f"[Experiment] Evaluate: user={user} - Flags: {flag_configs}")
-        sorted_flags = topological_sort(flag_configs, flag_keys and list(flag_keys))
+        sorted_flags = topological_sort(flag_configs, flag_keys_to_use and list(flag_keys_to_use))
         if not sorted_flags:
             return {}
 
@@ -120,7 +137,10 @@ def evaluate_v2(self, user: User, flag_keys: Set[str] = None) -> Dict[str, Varia
             ) for k, v in result.items()
         }
         self.logger.debug(f"[Experiment] Evaluate Result: {variants}")
+        if options.tracks_exposure is True:
+            self.exposure_service.track(Exposure(user, variants))
         if self.assignment_service is not None:
+            # @deprecated Assignment tracking is deprecated. Use ExposureService with Exposure tracking instead.
             self.assignment_service.track(Assignment(user, variants))
         return variants
 
diff --git a/src/amplitude_experiment/local/config.py b/src/amplitude_experiment/local/config.py
@@ -1,4 +1,5 @@
 from ..assignment import AssignmentConfig
+from ..exposure import ExposureConfig
 from ..cohort.cohort_sync_config import CohortSyncConfig, DEFAULT_COHORT_SYNC_URL, EU_COHORT_SYNC_URL
 from ..server_zone import ServerZone
 
@@ -21,6 +22,7 @@ def __init__(self, debug: bool = False,
                  stream_server_url: str = DEFAULT_STREAM_URL,
                  stream_flag_conn_timeout: int = 1500,
                  assignment_config: AssignmentConfig = None,
+                 exposure_config: ExposureConfig = ExposureConfig(),
                  cohort_sync_config: CohortSyncConfig = None):
         """
         Initialize a config
@@ -32,7 +34,8 @@ def __init__(self, debug: bool = False,
                   configurations.
                 flag_config_poller_request_timeout_millis (int): The request timeout, in milliseconds, used when
                   fetching flag configurations.
-                assignment_config (AssignmentConfig): The assignment configuration.
+                assignment_config (AssignmentConfig): The assignment configuration. @deprecated use exposure_config instead.
+                exposure_config (ExposureConfig): The exposure configuration.
                 cohort_sync_config (CohortSyncConfig): The cohort sync configuration.
 
            Returns:
@@ -57,3 +60,4 @@ def __init__(self, debug: bool = False,
         self.stream_updates = stream_updates
         self.stream_flag_conn_timeout = stream_flag_conn_timeout
         self.assignment_config = assignment_config
+        self.exposure_config = exposure_config
diff --git a/src/amplitude_experiment/local/evaluate_options.py b/src/amplitude_experiment/local/evaluate_options.py
@@ -0,0 +1,18 @@
+from typing import Optional, Set
+
+
+class EvaluateOptions:
+    """
+    Options for evaluating variants for a user.
+    """
+    def __init__(self, flag_keys: Optional[Set[str]] = None, tracks_exposure: Optional[bool] = None):
+        """
+        Initialize EvaluateOptions.
+        
+        Args:
+            flag_keys: The flags to evaluate with the user. If None, all flags are evaluated.
+            tracks_exposure: Whether to track exposure event for the evaluation.
+        """
+        self.flag_keys = flag_keys
+        self.tracks_exposure = tracks_exposure
+
diff --git a/tests/local/client_test.py b/tests/local/client_test.py
@@ -4,6 +4,7 @@
 
 from src.amplitude_experiment import LocalEvaluationClient, LocalEvaluationConfig, User, Variant
 from src.amplitude_experiment.cohort.cohort_sync_config import CohortSyncConfig
+from src.amplitude_experiment.local.evaluate_options import EvaluateOptions
 from dotenv import load_dotenv
 import os
 
@@ -111,6 +112,47 @@ def test_evaluation_cohorts_not_in_storage_with_sync_config(self):
                 )
                 self.assertTrue(any(re.match(log_message, message) for message in log.output))
 
+    def test_evaluate_v2_with_tracks_exposure_tracks_non_default_variants(self):
+        # Mock the amplitude client's track method
+        with mock.patch.object(self._local_evaluation_client.exposure_service.amplitude, 'track') as mock_track:
+            # Perform evaluation with tracks_exposure=True
+            options = EvaluateOptions(tracks_exposure=True)
+            variants = self._local_evaluation_client.evaluate_v2(test_user, None, options)
+            
+            # Verify that track was called
+            self.assertTrue(mock_track.called, 'Amplitude track should be called when tracks_exposure is True')
+            
+            # Count non-default variants
+            non_default_variants = {
+                flag_key: variant for flag_key, variant in variants.items()
+                if not (variant.metadata and variant.metadata.get('default'))
+            }
+            
+            # Get all tracked events
+            tracked_events = [call[0][0] for call in mock_track.call_args_list]
+            
+            # Verify that we have one event per non-default variant
+            self.assertEqual(len(tracked_events), len(non_default_variants),
+                           f'Expected {len(non_default_variants)} exposure events, got {len(tracked_events)}')
+            
+            # Verify each event has the correct structure
+            tracked_flag_keys = set()
+            for event in tracked_events:
+                self.assertEqual(event.event_type, '[Experiment] Exposure')
+                self.assertEqual(event.user_id, test_user.user_id)
+                flag_key = event.event_properties.get('[Experiment] Flag Key')
+                self.assertIsNotNone(flag_key, 'Event should have flag key')
+                tracked_flag_keys.add(flag_key)
+                # Verify the variant is not default
+                variant = variants.get(flag_key)
+                self.assertIsNotNone(variant, f'Variant for {flag_key} should exist')
+                self.assertFalse(variant.metadata and variant.metadata.get('default'),
+                               f'Variant for {flag_key} should not be default')
+            
+            # Verify all non-default variants were tracked
+            self.assertEqual(tracked_flag_keys, set(non_default_variants.keys()),
+                           'All non-default variants should be tracked')
+
 
 class LocalEvaluationClientStreamingTestCase(LocalEvaluationClientTestCase):
     _stream_update: bool = True
diff --git a/tests/local/exposure/__init__.py b/tests/local/exposure/__init__.py
diff --git a/tests/local/exposure/exposure_filter_test.py b/tests/local/exposure/exposure_filter_test.py
diff --git a/tests/local/exposure/exposure_service_test.py b/tests/local/exposure/exposure_service_test.py
