diff --git a/reflex/components/layout/flex.py b/reflex/components/layout/flex.py
index ab1f24d67..73cd01e32 100644
--- a/reflex/components/layout/flex.py
+++ b/reflex/components/layout/flex.py
@@ -1,5 +1,7 @@
 """A reflexive container component."""
 
+from typing import List, Union
+
 from reflex.components.libs.chakra import ChakraComponent
 from reflex.vars import Var
 
@@ -16,7 +18,7 @@ class Flex(ChakraComponent):
     basis: Var[str]
 
     # Shorthand for flexDirection style prop
-    direction: Var[str]
+    direction: Var[Union[str, List[str]]]
 
     # Shorthand for flexGrow style prop
     grow: Var[str]
@@ -25,7 +27,7 @@ class Flex(ChakraComponent):
     justify: Var[str]
 
     # Shorthand for flexWrap style prop
-    wrap: Var[str]
+    wrap: Var[Union[str, List[str]]]
 
     # Shorthand for flexShrink style prop
     shrink: Var[str]