diff --git a/atroposlib/envs/base.py b/atroposlib/envs/base.py
index 1afbc929..e4894aac 100644
--- a/atroposlib/envs/base.py
+++ b/atroposlib/envs/base.py
@@ -11,7 +11,8 @@ from abc import ABC, abstractmethod
 from datetime import datetime
 from enum import Enum
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Tuple, TypedDict, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
+from typing_extensions import TypedDict
 
 import aiohttp
 import jsonlines
diff --git a/atroposlib/type_definitions.py b/atroposlib/type_definitions.py
index 444af95d..c81fdf62 100644
--- a/atroposlib/type_definitions.py
+++ b/atroposlib/type_definitions.py
@@ -1,4 +1,5 @@
-from typing import Any, Dict, List, Literal, Optional, TypedDict
+from typing import Any, Dict, List, Literal, Optional
+from typing_extensions import TypedDict
 
 from openai.types.chat import ChatCompletionContentPartParam