diff --git a/packages/leann-backend-diskann/leann_backend_diskann/graph_partition.py b/packages/leann-backend-diskann/leann_backend_diskann/graph_partition.py
index 7afb9e9..4c85e28 100644
--- a/packages/leann-backend-diskann/leann_backend_diskann/graph_partition.py
+++ b/packages/leann-backend-diskann/leann_backend_diskann/graph_partition.py
@@ -263,8 +263,8 @@ class GraphPartitioner:
 
 def partition_graph(
     index_prefix_path: str,
-    output_dir: str | None = None,
-    partition_prefix: str | None = None,
+    output_dir: Optional[str] = None,
+    partition_prefix: Optional[str] = None,
     build_type: str = "release",
     **kwargs,
 ) -> tuple[str, str]: