diff --git a/ai21/models/responses/assistant_response.py b/ai21/models/responses/assistant_response.py
index 453d1c46..d7b0f186 100644
--- a/ai21/models/responses/assistant_response.py
+++ b/ai21/models/responses/assistant_response.py
@@ -1,5 +1,7 @@
 from datetime import datetime
-from typing import Optional, List, Literal, TypedDict
+from typing import Optional, List, Literal
+
+from typing_extensions import TypedDict
 
 from ai21.models.ai21_base_model import AI21BaseModel