diff --git a/engine/language_client_codegen/src/python/templates/types.py.j2 b/engine/language_client_codegen/src/python/templates/types.py.j2
index bbc7a7e9a..8c837846d 100644
--- a/engine/language_client_codegen/src/python/templates/types.py.j2
+++ b/engine/language_client_codegen/src/python/templates/types.py.j2
@@ -2,7 +2,8 @@
 import baml_py
 from enum import Enum
 from pydantic import BaseModel, ConfigDict
-from typing import Dict, Generic, List, Literal, Optional, TypeVar, Union, TypeAlias
+from typing_extensions import TypeAlias
+from typing import Dict, Generic, List, Literal, Optional, TypeVar, Union
 
 
 T = TypeVar('T')
diff --git a/integ-tests/python/baml_client/types.py b/integ-tests/python/baml_client/types.py
index 9d5644d6a..1e4f29f80 100644
--- a/integ-tests/python/baml_client/types.py
+++ b/integ-tests/python/baml_client/types.py
@@ -16,7 +16,8 @@
 import baml_py
 from enum import Enum
 from pydantic import BaseModel, ConfigDict
-from typing import Dict, Generic, List, Literal, Optional, TypeVar, Union, TypeAlias
+from typing_extensions import TypeAlias
+from typing import Dict, Generic, List, Literal, Optional, TypeVar, Union
 
 
 T = TypeVar('T')