diff --git a/flagsmith/api/types.py b/flagsmith/api/types.py
index fda43fe..eb1b00e 100644
--- a/flagsmith/api/types.py
+++ b/flagsmith/api/types.py
@@ -49,7 +49,7 @@ class MultivariateFeatureStateValueModel(typing.TypedDict):
 
 class FeatureStateModel(typing.TypedDict):
     enabled: bool
-    feature_segment: NotRequired[FeatureSegmentModel]
+    feature_segment: NotRequired[typing.Optional[FeatureSegmentModel]]
     feature_state_value: object
     feature: FeatureModel
     featurestate_uuid: str
diff --git a/flagsmith/mappers.py b/flagsmith/mappers.py
index 7564a89..9d2251a 100644
--- a/flagsmith/mappers.py
+++ b/flagsmith/mappers.py
@@ -267,7 +267,7 @@ def _map_environment_document_feature_states_to_feature_contexts(
                 for multivariate_feature_state_value in multivariate_feature_state_values
             ]
 
-        if "feature_segment" in feature_state:
-            feature_context["priority"] = feature_state["feature_segment"]["priority"]
+        if feature_segment := feature_state.get("feature_segment"):
+            feature_context["priority"] = feature_segment["priority"]
 
         yield feature_context
diff --git a/tests/data/environment.json b/tests/data/environment.json
index ea63981..7703255 100644
--- a/tests/data/environment.json
+++ b/tests/data/environment.json
@@ -51,7 +51,8 @@
                 "id": 1
             },
             "segment_id": null,
-            "enabled": true
+            "enabled": true,
+            "feature_segment": null
         }
     ],
     "updated_at": "2023-07-14 16:12:00.000000",