fix: alias

Author: asafgardin

.pre-commit-config.yaml

@@ -108,8 +108,6 @@ repos:
     rev: v0.0.280
     hooks:
       - id: ruff
-        args:
-          - --fix
   - repo: local
     hooks:
       - id: hadolint

ai21/clients/studio/resources/chat/__init__.py

@@ -1,3 +1,3 @@
 from __future__ import annotations
 
-from .chat_completion import ChatCompletions as ChatCompletions
+from .chat_completions import ChatCompletions as ChatCompletions

ai21/clients/studio/resources/chat/chat_completion.py renamed to ai21/clients/studio/resources/chat/chat_completions.py

@@ -1,4 +1,3 @@
-import warnings
 from typing import List, Optional
 
 from ai21.clients.common.chat_base import Chat
@@ -26,11 +25,6 @@ def create(
         count_penalty: Optional[Penalty] = None,
         **kwargs,
     ) -> ChatResponse:
-        warnings.warn(
-            "This method is deprecated. Please use the `chat.completions.create` method in the client instead.",
-            DeprecationWarning,
-            stacklevel=2,
-        )
         body = self._create_body(
             model=model,
             messages=messages,

ai21/clients/studio/resources/studio_chat.py

@@ -1,10 +1,12 @@
+from ai21.models.chat.role_type import RoleType
+from ai21.models.chat_message import ChatMessage
 from ai21.models.document_type import DocumentType
 from ai21.models.embed_type import EmbedType
 from ai21.models.improvement_type import ImprovementType
-from ai21.models.chat_message import ChatMessage
 from ai21.models.paraphrase_style_type import ParaphraseStyleType
 from ai21.models.penalty import Penalty
 from ai21.models.responses.answer_response import AnswerResponse
+from ai21.models.responses.chat_response import ChatResponse, ChatOutput, FinishReason
 from ai21.models.responses.completion_response import (
     CompletionsResponse,
     Completion,
@@ -18,14 +20,12 @@
 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.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
 from ai21.models.responses.segmentation_response import SegmentationResponse
 from ai21.models.responses.summarize_by_segment_response import SummarizeBySegmentResponse, SegmentSummary, Highlight
 from ai21.models.responses.summarize_response import SummarizeResponse
-from ai21.models.role_type import RoleType
 from ai21.models.summary_method import SummaryMethod
 
 __all__ = [

ai21/models/__init__.py

@@ -1,5 +1,6 @@
 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
+from .chat_completion_response import ChatCompletionResponse
+from .chat_completion_response import ChatCompletionResponseChoice
+from .chat_message import ChatMessage
+from .role_type import RoleType as RoleType

ai21/models/chat/__init__.py

@@ -1,5 +1,7 @@
 from enum import Enum
 
+__all__ = ["RoleType"]
+
 
 class RoleType(str, Enum):
     USER = "user"

ai21/models/role_type.py renamed to ai21/models/chat/role_type.py

@@ -1,7 +1,7 @@
 from dataclasses import dataclass
 
 from ai21.models.ai21_base_model_mixin import AI21BaseModelMixin
-from ai21.models.role_type import RoleType
+from ai21.models.chat.role_type import RoleType
 
 
 @dataclass

ai21/models/chat_message.py

@@ -2,7 +2,7 @@
 from typing import Optional, List
 
 from ai21.models.ai21_base_model_mixin import AI21BaseModelMixin
-from ai21.models.role_type import RoleType
+from ai21.models.chat.role_type import RoleType
 
 
 @dataclass