update trait type

Author: khvn26

flag_engine/context/types.py

@@ -4,11 +4,11 @@
 
 from __future__ import annotations
 
-from typing import Any, Dict, List, Optional, TypedDict, Union
+from typing import Any, Dict, List, Optional, TypedDict
 
 from typing_extensions import NotRequired
 
-from flag_engine.segments.types import ConditionOperator, RuleType
+from flag_engine.segments.types import ConditionOperator, ContextValue, RuleType
 
 
 class EnvironmentContext(TypedDict):
@@ -24,7 +24,7 @@ class FeatureValue(TypedDict):
 class IdentityContext(TypedDict):
     identifier: str
     key: str
-    traits: NotRequired[Dict[str, Optional[Union[str, float, bool]]]]
+    traits: NotRequired[Dict[str, ContextValue]]
 
 
 class SegmentCondition(TypedDict):