fix: Added __all__

Author: asafgardin

ai21/models/chat/__init__.py

@@ -0,0 +1,3 @@
+from __future__ import annotations
+
+from .chat_message import ChatMessage as ChatMessage

ai21/models/chat/chat_message.py

@@ -3,6 +3,8 @@
 from ai21.models import RoleType
 from ai21.models.ai21_base_model_mixin import AI21BaseModelMixin
 
+__all__ = ["ChatMessage"]
+
 
 @dataclass
 class ChatMessage(AI21BaseModelMixin):

tests/integration_tests/clients/studio/test_chat_completions.py

@@ -1,7 +1,8 @@
 import pytest
 
 from ai21 import AI21Client
-from ai21.models import ChatMessage, RoleType
+from ai21.models.chat import ChatMessage
+from ai21.models import RoleType
 from ai21.models.responses.chat_completion_response import ChatCompletionResponse