diff --git a/ai21/models/chat/__init__.py b/ai21/models/chat/__init__.py
index 84093937..dca5a0b5 100644
--- a/ai21/models/chat/__init__.py
+++ b/ai21/models/chat/__init__.py
@@ -1,19 +1,18 @@
 from __future__ import annotations
 
+from .chat_completion_chunk import ChatCompletionChunk, ChoicesChunk, ChoiceDelta
 from .chat_completion_response import ChatCompletionResponse
 from .chat_completion_response import ChatCompletionResponseChoice
-from .chat_message import ChatMessage, AssistantMessage, ToolMessage, UserMessage, SystemMessage
+from .chat_message import ChatMessage, AssistantMessage, ToolMessage, UserMessage, SystemMessage, ChatMessageParam
 from .document_schema import DocumentSchema
 from .function_tool_definition import FunctionToolDefinition
 from .response_format import ResponseFormat
 from .role_type import RoleType as RoleType
-from .chat_completion_chunk import ChatCompletionChunk, ChoicesChunk, ChoiceDelta
 from .tool_call import ToolCall
 from .tool_defintions import ToolDefinition
 from .tool_function import ToolFunction
 from .tool_parameters import ToolParameters
 
-
 __all__ = [
     "ChatCompletionResponse",
     "ChatCompletionResponseChoice",
@@ -26,6 +25,7 @@
     "ToolMessage",
     "UserMessage",
     "SystemMessage",
+    "ChatMessageParam",
     "DocumentSchema",
     "FunctionToolDefinition",
     "ResponseFormat",