refactor: imports and file structure

Author: asafgardin

ai21/clients/studio/resources/studio_chat_completion.py

@@ -3,8 +3,7 @@
 from typing import List, Optional, Union, Any, Dict
 
 from ai21.clients.studio.resources.studio_resource import StudioResource
-from ai21.models.chat import ChatMessage
-from ai21.models.responses.chat_completion_response import ChatCompletionResponse
+from ai21.models.chat import ChatMessage, ChatCompletionResponse
 from ai21.types import NotGiven, NOT_GIVEN
 from ai21.utils.typing import remove_not_given
 

ai21/models/__init__.py

@@ -18,7 +18,7 @@
 from ai21.models.responses.file_response import FileResponse
 from ai21.models.responses.gec_response import GECResponse, Correction, CorrectionType
 from ai21.models.responses.improvement_response import ImprovementsResponse, Improvement
-from ai21.models.responses.legacy.chat_response import ChatResponse, ChatOutput, FinishReason
+from ai21.models.responses.chat_response import ChatResponse, ChatOutput, FinishReason
 from ai21.models.responses.library_answer_response import LibraryAnswerResponse, SourceDocument
 from ai21.models.responses.library_search_response import LibrarySearchResponse, LibrarySearchResult
 from ai21.models.responses.paraphrase_response import ParaphraseResponse, Suggestion

ai21/models/chat/__init__.py

@@ -1,3 +1,5 @@
 from __future__ import annotations
 
 from .chat_message import ChatMessage as ChatMessage
+from .chat_completion_response import ChatCompletionResponse as ChatCompletionResponse
+from .chat_completion_response import ChatCompletionResponseChoice as ChatCompletionResponseChoice

ai21/models/responses/chat_completion_response.py renamed to ai21/models/chat/chat_completion_response.py

@@ -1,11 +1,13 @@
 from dataclasses import dataclass
 from typing import Optional, List
 
-from ai21.models import ChatMessage
 from ai21.models.ai21_base_model_mixin import AI21BaseModelMixin
+from ai21.models.chat import ChatMessage
 from ai21.models.logprobs import Logprobs
 from ai21.models.usage_info import UsageInfo
 
+__all__ = ["ChatCompletionResponse", "ChatCompletionResponseChoice"]
+
 
 @dataclass
 class ChatCompletionResponseChoice(AI21BaseModelMixin):

ai21/models/responses/legacy/chat_response.py renamed to ai21/models/responses/chat_response.py

@@ -3,7 +3,7 @@
 from ai21 import AI21Client
 from ai21.models.chat import ChatMessage
 from ai21.models import RoleType
-from ai21.models.responses.chat_completion_response import ChatCompletionResponse
+from ai21.models.chat.chat_completion_response import ChatCompletionResponse
 
 
 _MODEL = "gaia-small"