diff --git a/graphrag/llm/types/llm_io.py b/graphrag/llm/types/llm_io.py
index be80a54d96..52bf4b9e82 100644
--- a/graphrag/llm/types/llm_io.py
+++ b/graphrag/llm/types/llm_io.py
@@ -3,9 +3,9 @@
 """LLM Types."""
 
 from dataclasses import dataclass, field
-from typing import Generic, TypedDict, TypeVar
+from typing import Generic, TypeVar
 
-from typing_extensions import NotRequired
+from typing_extensions import NotRequired, TypedDict
 
 from .llm_callbacks import IsResponseValidFn