Updates to flag providers

efahk · efahk · commit 906d90232c1a · 2025-10-20T23:37:48.000Z
diff --git a/mixpanel/flags/local_feature_flags.py b/mixpanel/flags/local_feature_flags.py
@@ -4,7 +4,7 @@
 import time
 import threading
 from datetime import datetime, timedelta
-from typing import Dict, Any, Callable, Optional
+from typing import List, Dict, Any, Callable, Optional
 from .types import (
     ExperimentationFlag,
     ExperimentationFlags,
@@ -17,6 +17,7 @@
     normalized_hash,
     prepare_common_query_params,
     EXPOSURE_EVENT,
+    generate_traceparent,
 )
 
 logger = logging.getLogger(__name__)
@@ -170,15 +171,32 @@ def is_enabled(self, flag_key: str, context: Dict[str, Any]) -> bool:
         variant_value = self.get_variant_value(flag_key, False, context)
         return bool(variant_value)
 
+    def get_all_variants(self, context: Dict[str, Any], reportExposureEvents: bool = False) -> List[SelectedVariant]:
+        """
+        Gets the selected variant for all feature flags that the current user context is in the rollout for.
+        :param Dict[str, Any] context: The user context to evaluate against the feature flags
+        :param bool reportExposureEvents: Whether to immediately report exposure events to your Mixpanel project for each flag evaluated. Defaults to False.
+        """
+        variants: Dict[str, SelectedVariant] = {}
+        fallback = SelectedVariant(variant_key=None, variant_value=None)
+
+        for flag_key in self._flag_definitions.keys():
+            variant = self.get_variant(flag_key, fallback, context, report_exposure=reportExposureEvents)
+            if variant.variant_key is not None:
+                variants[flag_key] = variant
+
+        return variants
+
     def get_variant(
-        self, flag_key: str, fallback_value: SelectedVariant, context: Dict[str, Any]
+        self, flag_key: str, fallback_value: SelectedVariant, context: Dict[str, Any], report_exposure: bool = True
     ) -> SelectedVariant:
         """
         Gets the selected variant for a feature flag
 
         :param str flag_key: The key of the feature flag to evaluate
         :param SelectedVariant fallback_value: The default variant to return if evaluation fails
         :param Dict[str, Any] context: Context dictionary containing user's distinct_id and any other attributes needed for rollout evaluation
+        :param bool report_exposure: Whether to track an exposure event for this flag evaluation. Defaults to True.
         """
         start_time = time.perf_counter()
         flag_definition = self._flag_definitions.get(flag_key)
@@ -205,7 +223,8 @@ def get_variant(
                 flag_definition, context_value, flag_key, rollout
             )
             end_time = time.perf_counter()
-            self._track_exposure(flag_key, variant, end_time - start_time, context)
+            if report_exposure:
+                self._track_exposure(flag_key, variant, end_time - start_time, context)
             return variant
 
         logger.info(
@@ -241,12 +260,17 @@ def _get_assigned_variant(
             ):
                 return variant
 
-        variants = flag_definition.ruleset.variants
 
         hash_input = str(context_value) + flag_name
 
         variant_hash = normalized_hash(hash_input, "variant")
 
+        variants = [variant.model_copy(deep=True) for variant in flag_definition.ruleset.variants]
+        if rollout.variant_splits:
+            for variant in variants:
+                if variant.key in rollout.variant_splits:
+                    variant.split = rollout.variant_splits[variant.key]
+
         selected = variants[0]
         cumulative = 0.0
         for variant in variants:
@@ -255,7 +279,11 @@ def _get_assigned_variant(
             if variant_hash < cumulative:
                 break
 
-        return SelectedVariant(variant_key=selected.key, variant_value=selected.value)
+        return SelectedVariant(
+            variant_key=selected.key,
+            variant_value=selected.value,
+            experiment_id=flag_definition.experiment_id,
+            is_experiment_active=flag_definition.is_experiment_active)
 
     def _get_assigned_rollout(
         self,
@@ -304,15 +332,20 @@ def _get_matching_variant(
         for variant in flag.ruleset.variants:
             if variant_key.casefold() == variant.key.casefold():
                 return SelectedVariant(
-                    variant_key=variant.key, variant_value=variant.value
+                    variant_key=variant.key,
+                    variant_value=variant.value,
+                    experiment_id=flag.experiment_id,
+                    is_experiment_active=flag.is_experiment_active,
+                    is_qa_tester=True,
                 )
         return None
 
     async def _afetch_flag_definitions(self) -> None:
         try:
             start_time = datetime.now()
+            headers = {"traceparent": generate_traceparent()}
             response = await self._async_client.get(
-                self.FLAGS_DEFINITIONS_URL_PATH, params=self._request_params
+                self.FLAGS_DEFINITIONS_URL_PATH, params=self._request_params, headers=headers
             )
             end_time = datetime.now()
             self._handle_response(response, start_time, end_time)
@@ -322,8 +355,9 @@ async def _afetch_flag_definitions(self) -> None:
     def _fetch_flag_definitions(self) -> None:
         try:
             start_time = datetime.now()
+            headers = {"traceparent": generate_traceparent()}
             response = self._sync_client.get(
-                self.FLAGS_DEFINITIONS_URL_PATH, params=self._request_params
+                self.FLAGS_DEFINITIONS_URL_PATH, params=self._request_params, headers=headers
             )
             end_time = datetime.now()
             self._handle_response(response, start_time, end_time)
@@ -370,6 +404,9 @@ def _track_exposure(
                 "$experiment_type": "feature_flag",
                 "Flag evaluation mode": "local",
                 "Variant fetch latency (ms)": latency_in_seconds * 1000,
+                "$experiment_id": variant.experiment_id,
+                "$is_experiment_active": variant.is_experiment_active,
+                "$is_qa_tester": variant.is_qa_tester,
             }
 
             self._tracker(distinct_id, EXPOSURE_EVENT, properties)
diff --git a/mixpanel/flags/remote_feature_flags.py b/mixpanel/flags/remote_feature_flags.py
@@ -8,7 +8,7 @@
 from asgiref.sync import sync_to_async
 
 from .types import RemoteFlagsConfig, SelectedVariant, RemoteFlagsResponse
-from .utils import REQUEST_HEADERS, EXPOSURE_EVENT, prepare_common_query_params
+from .utils import REQUEST_HEADERS, EXPOSURE_EVENT, prepare_common_query_params, generate_traceparent
 
 logger = logging.getLogger(__name__)
 logging.getLogger("httpx").setLevel(logging.ERROR)
@@ -66,7 +66,8 @@ async def aget_variant(
         try:
             params = self._prepare_query_params(flag_key, context)
             start_time = datetime.now()
-            response = await self._async_client.get(self.FLAGS_URL_PATH, params=params)
+            headers = {"traceparent": generate_traceparent()}
+            response = await self._async_client.get(self.FLAGS_URL_PATH, params=params, headers=headers)
             end_time = datetime.now()
             self._instrument_call(start_time, end_time)
             selected_variant, is_fallback = self._handle_response(
@@ -126,7 +127,8 @@ def get_variant(
         try:
             params = self._prepare_query_params(flag_key, context)
             start_time = datetime.now()
-            response = self._sync_client.get(self.FLAGS_URL_PATH, params=params)
+            headers = {"traceparent": generate_traceparent()}
+            response = self._sync_client.get(self.FLAGS_URL_PATH, params=params, headers=headers)
             end_time = datetime.now()
             self._instrument_call(start_time, end_time)
             selected_variant, is_fallback = self._handle_response(
diff --git a/mixpanel/flags/types.py b/mixpanel/flags/types.py
@@ -32,6 +32,7 @@ class Rollout(BaseModel):
     rollout_percentage: float
     runtime_evaluation_definition: Optional[Dict[str, str]] = None
     variant_override: Optional[VariantOverride] = None
+    variant_splits: Optional[Dict[str,float]] = None
 
 class RuleSet(BaseModel):
     variants: List[Variant]
@@ -41,16 +42,23 @@ class RuleSet(BaseModel):
 class ExperimentationFlag(BaseModel):
     id: str
     name: str
-    key: str 
+    key: str
     status: str
     project_id: int
-    ruleset: RuleSet 
+    ruleset: RuleSet
     context: str
+    experiment_id: Optional[str] = None
+    is_experiment_active: Optional[bool] = None
+
 
 class SelectedVariant(BaseModel):
     # variant_key can be None if being used as a fallback
     variant_key: Optional[str] = None
     variant_value: Any
+    experiment_id: Optional[str] = None
+    is_experiment_active: Optional[bool] = None
+    is_qa_tester: Optional[bool] = None
+
 
 class ExperimentationFlags(BaseModel):
     flags: List[ExperimentationFlag] 
diff --git a/mixpanel/flags/utils.py b/mixpanel/flags/utils.py
@@ -1,4 +1,5 @@
 from typing import Dict
+from uuid import uuid
 
 EXPOSURE_EVENT = "$experiment_started"
 
@@ -47,4 +48,19 @@ def prepare_common_query_params(token: str, sdk_version: str) -> Dict[str, str]:
         'token': token
     }
 
-    return params
+    return params
+
+def generate_traceparent() -> str:
+    """ Generates a W3C traceparent header for easy interop with distributed tracing systems i.e Open Telemetry
+    https://www.w3.org/TR/trace-context/#traceparent-header
+    :return: A traceparent string
+    """
+
+    trace_id = uuid.uuid4().hex
+    span_id = uuid.uuid4().hex[:16]
+
+    # Trace flags: '01' for sampled
+    trace_flags = '01'
+
+    traceparent = f"00-{trace_id}-{span_id}-{trace_flags}"
+    return traceparent
