diff --git a/semantic_router/schema.py b/semantic_router/schema.py
index 5dc4f5e87fba43d15e5e1ec39b578e3f37a54cfe..07b460b15c65088eeca6bffe394da8751ccb92ff 100644
--- a/semantic_router/schema.py
+++ b/semantic_router/schema.py
@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import Dict, List, Literal, Optional
+from typing import List, Literal, Optional
 
 from pydantic.v1 import BaseModel
 from pydantic.v1.dataclasses import dataclass
@@ -10,7 +10,7 @@ from semantic_router.encoders import (
     FastEmbedEncoder,
     OpenAIEncoder,
 )
-from semantic_router.utils.splitters import semantic_splitter
+from semantic_router.utils.splitters import semantic_splitter, DocumentSplit
 
 
 class EncoderType(Enum):
@@ -77,7 +77,7 @@ class Conversation(BaseModel):
         split_method: Literal[
             "consecutive_similarity_drop", "cumulative_similarity_drop"
         ] = "consecutive_similarity_drop",
-    ) -> Dict[str, List[str]]:
+    ) -> list[DocumentSplit]:
         docs = [f"{m.role}: {m.content}" for m in self.messages]
         return semantic_splitter(
             encoder=encoder, docs=docs, threshold=threshold, split_method=split_method