feat(flags): Enable local evaluation for all cohorts (#91)

* add function for active variants

* format

CHANGELOG.md

@@ -1,3 +1,13 @@
+## 2.4.0 - 2023-03-14
+
+1. Support evaluating all cohorts in feature flags for local evaluation
+
+## 2.3.1 - 2023-02-07
+
+1. Log instead of raise error on posthog personal api key errors
+2. Remove upper bound on backoff dependency
+
+
 ## 2.3.0 - 2023-01-31
 
 1. Add support for returning payloads of matched feature flags

posthog/client.py

@@ -48,7 +48,6 @@ class Client(object):
         personal_api_key=None,
         project_api_key=None,
     ):
-
         self.queue = queue.Queue(max_queue_size)
 
         # api_key: This should be the Team API Key (token), public
@@ -66,6 +65,7 @@ class Client(object):
         self.feature_flags = None
         self.feature_flags_by_key = None
         self.group_type_mapping = None
+        self.cohorts = None
         self.poll_interval = poll_interval
         self.poller = None
         self.distinct_ids_feature_flags_reported = SizeLimitedDict(MAX_DICT_SIZE, set)
@@ -131,6 +131,12 @@ class Client(object):
         resp_data = self.get_decide(distinct_id, groups, person_properties, group_properties)
         return resp_data["featureFlags"]
 
+    def _get_active_feature_variants(self, distinct_id, groups=None, person_properties=None, group_properties=None):
+        feature_variants = self.get_feature_variants(distinct_id, groups, person_properties, group_properties)
+        return {
+            k: v for (k, v) in feature_variants.items() if v is not False
+        }  # explicitly test for false to account for values that may seem falsy (ex: 0)
+
     def get_feature_payloads(self, distinct_id, groups=None, person_properties=None, group_properties=None):
         resp_data = self.get_decide(distinct_id, groups, person_properties, group_properties)
         return resp_data["featureFlagPayloads"]
@@ -150,6 +156,8 @@ class Client(object):
             "group_properties": group_properties,
         }
         resp_data = decide(self.api_key, self.host, timeout=10, **request_data)
+
+        print(resp_data)
         return resp_data
 
     def capture(
@@ -184,7 +192,7 @@ class Client(object):
 
         if send_feature_flags:
             try:
-                feature_variants = self.get_feature_variants(distinct_id, groups)
+                feature_variants = self._get_active_feature_variants(distinct_id, groups)
             except Exception as e:
                 self.log.exception(f"[FEATURE FLAGS] Unable to get feature variants: {e}")
             else:
@@ -367,10 +375,9 @@ class Client(object):
 
     def _load_feature_flags(self):
         try:
-
             response = get(
                 self.personal_api_key,
-                f"/api/feature_flag/local_evaluation/?token={self.api_key}",
+                f"/api/feature_flag/local_evaluation/?token={self.api_key}&send_cohorts",
                 self.host,
                 timeout=10,
             )
@@ -380,15 +387,20 @@ class Client(object):
                 flag["key"]: flag for flag in self.feature_flags if flag.get("key") is not None
             }
             self.group_type_mapping = response["group_type_mapping"] or {}
+            self.cohorts = response["cohorts"] or {}
 
         except APIError as e:
             if e.status == 401:
-                raise APIError(
-                    status=401,
-                    message="You are using a write-only key with feature flags. "
-                    "To use feature flags, please set a personal_api_key "
-                    "More information: https://posthog.com/docs/api/overview",
+                self.log.error(
+                    f"[FEATURE FLAGS] Error loading feature flags: To use feature flags, please set a valid personal_api_key. More information: https://posthog.com/docs/api/overview"
                 )
+                if self.debug:
+                    raise APIError(
+                        status=401,
+                        message="You are using a write-only key with feature flags. "
+                        "To use feature flags, please set a personal_api_key "
+                        "More information: https://posthog.com/docs/api/overview",
+                    )
             else:
                 self.log.error(f"[FEATURE FLAGS] Error loading feature flags: {e}")
         except Exception as e:
@@ -412,7 +424,6 @@ class Client(object):
             self.poller.start()
 
     def _compute_flag_locally(self, feature_flag, distinct_id, *, groups={}, person_properties={}, group_properties={}):
-
         if feature_flag.get("ensure_experience_continuity", False):
             raise InconclusiveMatchError("Flag has experience continuity enabled")
 
@@ -442,7 +453,7 @@ class Client(object):
             focused_group_properties = group_properties[group_name]
             return match_feature_flag_properties(feature_flag, groups[group_name], focused_group_properties)
         else:
-            return match_feature_flag_properties(feature_flag, distinct_id, person_properties)
+            return match_feature_flag_properties(feature_flag, distinct_id, person_properties, self.cohorts)
 
     def feature_enabled(
         self,

posthog/feature_flags.py

@@ -1,5 +1,6 @@
 import datetime
 import hashlib
+import logging
 import re
 
 from dateutil import parser
@@ -8,6 +9,8 @@ from posthog.utils import convert_to_datetime_aware, is_valid_regex
 
 __LONG_SCALE__ = float(0xFFFFFFFFFFFFFFF)
 
+log = logging.getLogger("posthog")
+
 
 class InconclusiveMatchError(Exception):
     pass
@@ -42,9 +45,10 @@ def variant_lookup_table(feature_flag):
     return lookup_table
 
 
-def match_feature_flag_properties(flag, distinct_id, properties):
+def match_feature_flag_properties(flag, distinct_id, properties, cohort_properties=None):
     flag_conditions = (flag.get("filters") or {}).get("groups") or []
     is_inconclusive = False
+    cohort_properties = cohort_properties or {}
 
     # Stable sort conditions with variant overrides to the top. This ensures that if overrides are present, they are
     # evaluated first, and the variant override is applied to the first matching condition.
@@ -57,7 +61,7 @@ def match_feature_flag_properties(flag, distinct_id, properties):
         try:
             # if any one condition resolves to True, we can shortcircuit and return
             # the matching variant
-            if is_condition_match(flag, distinct_id, condition, properties):
+            if is_condition_match(flag, distinct_id, condition, properties, cohort_properties):
                 variant_override = condition.get("variant")
                 # Some filters can be explicitly set to null, which require accessing variants like so
                 flag_variants = ((flag.get("filters") or {}).get("multivariate") or {}).get("variants") or []
@@ -77,12 +81,19 @@ def match_feature_flag_properties(flag, distinct_id, properties):
     return False
 
 
-def is_condition_match(feature_flag, distinct_id, condition, properties):
+def is_condition_match(feature_flag, distinct_id, condition, properties, cohort_properties):
     rollout_percentage = condition.get("rollout_percentage")
     if len(condition.get("properties") or []) > 0:
-        if not all(match_property(prop, properties) for prop in condition.get("properties")):
-            return False
-        elif rollout_percentage is None:
+        for prop in condition.get("properties"):
+            property_type = prop.get("type")
+            if property_type == "cohort":
+                matches = match_cohort(prop, properties, cohort_properties)
+            else:
+                matches = match_property(prop, properties)
+            if not matches:
+                return False
+
+        if rollout_percentage is None:
             return True
 
     if rollout_percentage is not None and _hash(feature_flag["key"], distinct_id) > (rollout_percentage / 100):
@@ -175,3 +186,88 @@ def match_property(property, property_values) -> bool:
             raise InconclusiveMatchError("The date provided must be a string or date object")
 
     return False
+
+
+def match_cohort(property, property_values, cohort_properties) -> bool:
+    # Cohort properties are in the form of property groups like this:
+    # {
+    #     "cohort_id": {
+    #         "type": "AND|OR",
+    #         "values": [{
+    #            "key": "property_name", "value": "property_value"
+    #        }]
+    #     }
+    # }
+    cohort_id = str(property.get("value"))
+    if cohort_id not in cohort_properties:
+        raise InconclusiveMatchError("can't match cohort without a given cohort property value")
+
+    property_group = cohort_properties[cohort_id]
+    return match_property_group(property_group, property_values, cohort_properties)
+
+
+def match_property_group(property_group, property_values, cohort_properties) -> bool:
+    if not property_group:
+        return True
+
+    property_group_type = property_group.get("type")
+    properties = property_group.get("values")
+
+    if not properties or len(properties) == 0:
+        # empty groups are no-ops, always match
+        return True
+
+    error_matching_locally = False
+
+    if "values" in properties[0]:
+        # a nested property group
+        for prop in properties:
+            try:
+                matches = match_property_group(prop, property_values, cohort_properties)
+                if property_group_type == "AND":
+                    if not matches:
+                        return False
+                else:
+                    # OR group
+                    if matches:
+                        return True
+            except InconclusiveMatchError as e:
+                log.debug(f"Failed to compute property {prop} locally: {e}")
+                error_matching_locally = True
+
+        if error_matching_locally:
+            raise InconclusiveMatchError("Can't match cohort without a given cohort property value")
+        # if we get here, all matched in AND case, or none matched in OR case
+        return property_group_type == "AND"
+
+    else:
+        for prop in properties:
+            try:
+                if prop.get("type") == "cohort":
+                    matches = match_cohort(prop, property_values, cohort_properties)
+                else:
+                    matches = match_property(prop, property_values)
+
+                negation = prop.get("negation", False)
+
+                if property_group_type == "AND":
+                    # if negated property, do the inverse
+                    if not matches and not negation:
+                        return False
+                    if matches and negation:
+                        return False
+                else:
+                    # OR group
+                    if matches and not negation:
+                        return True
+                    if not matches and negation:
+                        return True
+            except InconclusiveMatchError as e:
+                log.debug(f"Failed to compute property {prop} locally: {e}")
+                error_matching_locally = True
+
+        if error_matching_locally:
+            raise InconclusiveMatchError("can't match cohort without a given cohort property value")
+
+        # if we get here, all matched in AND case, or none matched in OR case
+        return property_group_type == "AND"

posthog/test/test_client.py

@@ -71,7 +71,6 @@ class TestClient(unittest.TestCase):
         self.assertEqual(msg["properties"]["$lib_version"], VERSION)
 
     def test_basic_capture_with_project_api_key(self):
-
         client = Client(project_api_key=FAKE_TEST_API_KEY, on_error=self.set_fail)
 
         success, msg = client.capture("distinct_id", "python test event")
@@ -107,6 +106,40 @@ class TestClient(unittest.TestCase):
 
         self.assertEqual(patch_decide.call_count, 1)
 
+    @mock.patch("posthog.client.decide")
+    def test_get_active_feature_flags(self, patch_decide):
+        patch_decide.return_value = {
+            "featureFlags": {"beta-feature": "random-variant", "alpha-feature": True, "off-feature": False}
+        }
+
+        client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, personal_api_key=FAKE_TEST_API_KEY)
+        variants = client._get_active_feature_variants("some_id", None, None, None)
+        self.assertEqual(variants, {"beta-feature": "random-variant", "alpha-feature": True})
+
+    @mock.patch("posthog.client.decide")
+    def test_basic_capture_with_feature_flags_returns_active_only(self, patch_decide):
+        patch_decide.return_value = {
+            "featureFlags": {"beta-feature": "random-variant", "alpha-feature": True, "off-feature": False}
+        }
+
+        client = Client(FAKE_TEST_API_KEY, on_error=self.set_fail, personal_api_key=FAKE_TEST_API_KEY)
+        success, msg = client.capture("distinct_id", "python test event", send_feature_flags=True)
+        client.flush()
+        self.assertTrue(success)
+        self.assertFalse(self.failed)
+
+        self.assertEqual(msg["event"], "python test event")
+        self.assertTrue(isinstance(msg["timestamp"], str))
+        self.assertIsNone(msg.get("uuid"))
+        self.assertEqual(msg["distinct_id"], "distinct_id")
+        self.assertEqual(msg["properties"]["$lib"], "posthog-python")
+        self.assertEqual(msg["properties"]["$lib_version"], VERSION)
+        self.assertEqual(msg["properties"]["$feature/beta-feature"], "random-variant")
+        self.assertEqual(msg["properties"]["$feature/alpha-feature"], True)
+        self.assertEqual(msg["properties"]["$active_feature_flags"], ["beta-feature", "alpha-feature"])
+
+        self.assertEqual(patch_decide.call_count, 1)
+
     @mock.patch("posthog.client.decide")
     def test_basic_capture_with_feature_flags_switched_off_doesnt_send_them(self, patch_decide):
         patch_decide.return_value = {"featureFlags": {"beta-feature": "random-variant"}}

posthog/test/test_feature_flags.py

@@ -960,6 +960,159 @@ class TestLocalEvaluation(unittest.TestCase):
         self.assertEqual(patch_decide.call_count, 0)
         self.assertEqual(patch_capture.call_count, 0)
 
+    @mock.patch("posthog.client.decide")
+    @mock.patch("posthog.client.get")
+    def test_feature_flags_local_evaluation_for_cohorts(self, patch_get, patch_decide):
+        client = Client(FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
+        client.feature_flags = [
+            {
+                "id": 2,
+                "name": "Beta Feature",
+                "key": "beta-feature",
+                "is_simple_flag": False,
+                "active": True,
+                "filters": {
+                    "groups": [
+                        {
+                            "properties": [
+                                {
+                                    "key": "region",
+                                    "operator": "exact",
+                                    "value": ["USA"],
+                                    "type": "person",
+                                },
+                                {"key": "id", "value": 98, "operator": None, "type": "cohort"},
+                            ],
+                            "rollout_percentage": 100,
+                        }
+                    ],
+                },
+            },
+        ]
+        client.cohorts = {
+            "98": {
+                "type": "OR",
+                "values": [
+                    {"key": "id", "value": 1, "type": "cohort"},
+                    {
+                        "key": "nation",
+                        "operator": "exact",
+                        "value": ["UK"],
+                        "type": "person",
+                    },
+                ],
+            },
+            "1": {
+                "type": "AND",
+                "values": [{"key": "other", "operator": "exact", "value": ["thing"], "type": "person"}],
+            },
+        }
+
+        feature_flag_match = client.get_feature_flag(
+            "beta-feature", "some-distinct-id", person_properties={"region": "UK"}
+        )
+
+        self.assertEqual(feature_flag_match, False)
+        self.assertEqual(patch_decide.call_count, 0)
+        self.assertEqual(patch_get.call_count, 0)
+
+        feature_flag_match = client.get_feature_flag(
+            "beta-feature", "some-distinct-id", person_properties={"region": "USA", "nation": "UK"}
+        )
+        # even though 'other' property is not present, the cohort should still match since it's an OR condition
+        self.assertEqual(feature_flag_match, True)
+        self.assertEqual(patch_decide.call_count, 0)
+        self.assertEqual(patch_get.call_count, 0)
+
+        feature_flag_match = client.get_feature_flag(
+            "beta-feature", "some-distinct-id", person_properties={"region": "USA", "other": "thing"}
+        )
+        self.assertEqual(feature_flag_match, True)
+        self.assertEqual(patch_decide.call_count, 0)
+        self.assertEqual(patch_get.call_count, 0)
+
+    @mock.patch("posthog.client.decide")
+    @mock.patch("posthog.client.get")
+    def test_feature_flags_local_evaluation_for_negated_cohorts(self, patch_get, patch_decide):
+        client = Client(FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
+        client.feature_flags = [
+            {
+                "id": 2,
+                "name": "Beta Feature",
+                "key": "beta-feature",
+                "is_simple_flag": False,
+                "active": True,
+                "filters": {
+                    "groups": [
+                        {
+                            "properties": [
+                                {
+                                    "key": "region",
+                                    "operator": "exact",
+                                    "value": ["USA"],
+                                    "type": "person",
+                                },
+                                {"key": "id", "value": 98, "operator": None, "type": "cohort"},
+                            ],
+                            "rollout_percentage": 100,
+                        }
+                    ],
+                },
+            },
+        ]
+        client.cohorts = {
+            "98": {
+                "type": "OR",
+                "values": [
+                    {"key": "id", "value": 1, "type": "cohort"},
+                    {
+                        "key": "nation",
+                        "operator": "exact",
+                        "value": ["UK"],
+                        "type": "person",
+                    },
+                ],
+            },
+            "1": {
+                "type": "AND",
+                "values": [
+                    {"key": "other", "operator": "exact", "value": ["thing"], "type": "person", "negation": True}
+                ],
+            },
+        }
+
+        feature_flag_match = client.get_feature_flag(
+            "beta-feature", "some-distinct-id", person_properties={"region": "UK"}
+        )
+
+        self.assertEqual(feature_flag_match, False)
+        self.assertEqual(patch_decide.call_count, 0)
+        self.assertEqual(patch_get.call_count, 0)
+
+        feature_flag_match = client.get_feature_flag(
+            "beta-feature", "some-distinct-id", person_properties={"region": "USA", "nation": "UK"}
+        )
+        # even though 'other' property is not present, the cohort should still match since it's an OR condition
+        self.assertEqual(feature_flag_match, True)
+        self.assertEqual(patch_decide.call_count, 0)
+        self.assertEqual(patch_get.call_count, 0)
+
+        feature_flag_match = client.get_feature_flag(
+            "beta-feature", "some-distinct-id", person_properties={"region": "USA", "other": "thing"}
+        )
+        # since 'other' is negated, we return False. Since 'nation' is not present, we can't tell whether the flag should be true or false, so go to decide
+        self.assertEqual(patch_decide.call_count, 1)
+        self.assertEqual(patch_get.call_count, 0)
+
+        patch_decide.reset_mock()
+
+        feature_flag_match = client.get_feature_flag(
+            "beta-feature", "some-distinct-id", person_properties={"region": "USA", "other": "thing2"}
+        )
+        self.assertEqual(feature_flag_match, True)
+        self.assertEqual(patch_decide.call_count, 0)
+        self.assertEqual(patch_get.call_count, 0)
+
     @mock.patch("posthog.client.Poller")
     @mock.patch("posthog.client.get")
     def test_load_feature_flags(self, patch_get, patch_poll):
@@ -981,8 +1134,15 @@ class TestLocalEvaluation(unittest.TestCase):
 
     def test_load_feature_flags_wrong_key(self):
         client = Client(FAKE_TEST_API_KEY, personal_api_key=FAKE_TEST_API_KEY)
-        with freeze_time("2020-01-01T12:01:00.0000Z"):
-            self.assertRaises(APIError, client.load_feature_flags)
+
+        with self.assertLogs("posthog", level="ERROR") as logs:
+            client.load_feature_flags()
+            self.assertEqual(
+                logs.output[0],
+                "ERROR:posthog:[FEATURE FLAGS] Error loading feature flags: To use feature flags, please set a valid personal_api_key. More information: https://posthog.com/docs/api/overview",
+            )
+        client.debug = True
+        self.assertRaises(APIError, client.load_feature_flags)
 
     @mock.patch("posthog.client.decide")
     @mock.patch("posthog.client.get")

posthog/version.py

@@ -1,4 +1,4 @@
-VERSION = "2.3.0"
+VERSION = "2.4.0"
 
 if __name__ == "__main__":
     print(VERSION, end="")

setup.py

@@ -14,7 +14,7 @@ long_description = """
 PostHog is developer-friendly, self-hosted product analytics. posthog-python is the python package.
 """
 
-install_requires = ["requests>=2.7,<3.0", "six>=1.5", "monotonic>=1.5", "backoff>=1.10.0,<2.0.0", "python-dateutil>2.1"]
+install_requires = ["requests>=2.7,<3.0", "six>=1.5", "monotonic>=1.5", "backoff>=1.10.0", "python-dateutil>2.1"]
 
 extras_require = {
     "dev": [