fix: Restructure packages

ai21/__init__.py

@@ -10,21 +10,6 @@
     TooManyRequestsError,
 )
 from ai21.logger import setup_logger
-from ai21.resources.responses.answer_response import AnswerResponse
-from ai21.resources.responses.chat_response import ChatResponse
-from ai21.resources.responses.completion_response import CompletionsResponse
-from ai21.resources.responses.custom_model_response import CustomBaseModelResponse
-from ai21.resources.responses.dataset_response import DatasetResponse
-from ai21.resources.responses.embed_response import EmbedResponse
-from ai21.resources.responses.file_response import FileResponse
-from ai21.resources.responses.gec_response import GECResponse
-from ai21.resources.responses.improvement_response import ImprovementsResponse
-from ai21.resources.responses.library_answer_response import LibraryAnswerResponse
-from ai21.resources.responses.library_search_response import LibrarySearchResponse
-from ai21.resources.responses.paraphrase_response import ParaphraseResponse
-from ai21.resources.responses.segmentation_response import SegmentationResponse
-from ai21.resources.responses.summarize_by_segment_response import SummarizeBySegmentResponse
-from ai21.resources.responses.summarize_response import SummarizeResponse
 from ai21.services.sagemaker import SageMaker
 from ai21.version import VERSION
 
@@ -75,20 +60,5 @@ def __getattr__(name: str) -> Any:
     "AI21BedrockClient",
     "AI21SageMakerClient",
     "BedrockModelID",
-    "AnswerResponse",
-    "ChatResponse",
-    "CompletionsResponse",
-    "CustomBaseModelResponse",
-    "DatasetResponse",
-    "EmbedResponse",
-    "FileResponse",
-    "GECResponse",
-    "ImprovementsResponse",
-    "LibraryAnswerResponse",
-    "LibrarySearchResponse",
-    "ParaphraseResponse",
     "SageMaker",
-    "SegmentationResponse",
-    "SummarizeBySegmentResponse",
-    "SummarizeResponse",
 ]

ai21/ai21_http_client.py

@@ -25,7 +25,7 @@ def __init__(
         self._env_config = env_config
         self._api_key = api_key or self._env_config.api_key
 
-        if self._api_key is None:
+        if not self._api_key:
             raise MissingApiKeyError()
 
         self._api_host = api_host or self._env_config.api_host

ai21/clients/bedrock/resources/bedrock_completion.py

@@ -1,8 +1,8 @@
 from typing import Optional, List
 
 from ai21.resources import Penalty
-from ai21.resources.bedrock_resource import BedrockResource
-from ai21.resources.responses.completion_response import CompletionsResponse
+from ai21.clients.bedrock.resources.bedrock_resource import BedrockResource
+from ai21.models.responses.completion_response import CompletionsResponse
 
 
 class BedrockCompletion(BedrockResource):

ai21/resources/bases/answer_base.py → ai21/clients/common/answer_base.py

@@ -1,9 +1,8 @@
 from abc import ABC
 from typing import Optional, Any, Dict
 
-from ai21.resources.models.answer_length import AnswerLength
-from ai21.resources.models.mode import Mode
-from ai21.resources.responses.answer_response import AnswerResponse
+from ai21.models import Mode, AnswerLength
+from ai21.models.responses.answer_response import AnswerResponse
 
 
 class Answer(ABC):

ai21/resources/bases/chat_base.py → ai21/clients/common/chat_base.py

@@ -1,18 +1,9 @@
 from abc import ABC, abstractmethod
-from dataclasses import dataclass
 from typing import List, Any, Dict, Optional
 
-from ai21.models.ai21_base_model_mixin import AI21BaseModelMixin
-from ai21.resources.models.penalty import Penalty
-from ai21.resources.models.role_type import RoleType
-from ai21.resources.responses.chat_response import ChatResponse
-
-
-@dataclass
-class Message(AI21BaseModelMixin):
-    role: RoleType
-    text: str
-    name: Optional[str]
+from ai21.models.chat_message import ChatMessage
+from ai21.models.penalty import Penalty
+from ai21.models.responses.chat_response import ChatResponse
 
 
 class Chat(ABC):
@@ -22,7 +13,7 @@ class Chat(ABC):
     def create(
         self,
         model: str,
-        messages: List[Message],
+        messages: List[ChatMessage],
         system: str,
         *,
         num_results: Optional[int] = 1,
@@ -45,7 +36,7 @@ def _json_to_response(self, json: Dict[str, Any]) -> ChatResponse:
     def _create_body(
         self,
         model: str,
-        messages: List[Message],
+        messages: List[ChatMessage],
         system: str,
         num_results: Optional[int] = 1,
         temperature: Optional[float] = 0.7,

ai21/resources/bases/completion_base.py → ai21/clients/common/completion_base.py

@@ -1,8 +1,8 @@
 from abc import ABC, abstractmethod
 from typing import Optional, List, Dict, Any
 
-from ai21.resources import Penalty
-from ai21.resources.responses.completion_response import CompletionsResponse
+from ai21.models import Penalty
+from ai21.models.responses.completion_response import CompletionsResponse
 
 
 class Completion(ABC):

ai21/resources/bases/custom_model_base.py → ai21/clients/common/custom_model_base.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 from typing import Optional, List, Any, Dict
 
-from ai21.resources.responses.custom_model_response import CustomBaseModelResponse
+from ai21.models.responses.custom_model_response import CustomBaseModelResponse
 
 
 class CustomModel(ABC):

ai21/resources/bases/dataset_base.py → ai21/clients/common/dataset_base.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 from typing import Optional, Any, Dict
 
-from ai21.resources.responses.dataset_response import DatasetResponse
+from ai21.models.responses.dataset_response import DatasetResponse
 
 
 class Dataset(ABC):

ai21/resources/bases/embed_base.py → ai21/clients/common/embed_base.py

@@ -1,13 +1,8 @@
 from abc import ABC, abstractmethod
-from enum import Enum
 from typing import List, Any, Dict, Optional
 
-from ai21.resources.responses.embed_response import EmbedResponse
-
-
-class EmbedType(str, Enum):
-    QUERY = "query"
-    SEGMENT = "segment"
+from ai21.models.embed_type import EmbedType
+from ai21.models.responses.embed_response import EmbedResponse
 
 
 class Embed(ABC):

ai21/resources/bases/gec_base.py → ai21/clients/common/gec_base.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 from typing import Dict, Any
 
-from ai21.resources.responses.gec_response import GECResponse
+from ai21.models.responses.gec_response import GECResponse
 
 
 class GEC(ABC):

ai21/resources/bases/improvements_base.py → ai21/clients/common/improvements_base.py

@@ -1,8 +1,8 @@
 from abc import ABC, abstractmethod
 from typing import Any, Dict, List
 
-from ai21.resources.models.improvement_type import ImprovementType
-from ai21.resources.responses.improvement_response import ImprovementsResponse
+from ai21.models import ImprovementType
+from ai21.models.responses.improvement_response import ImprovementsResponse
 
 
 class Improvements(ABC):

ai21/resources/bases/paraphrase_base.py → ai21/clients/common/paraphrase_base.py

@@ -1,8 +1,8 @@
 from abc import ABC, abstractmethod
 from typing import Optional, Any, Dict
 
-from ai21.resources.models.paraphrase_style_type import ParaphraseStyleType
-from ai21.resources.responses.paraphrase_response import ParaphraseResponse
+from ai21.models import ParaphraseStyleType
+from ai21.models.responses.paraphrase_response import ParaphraseResponse
 
 
 class Paraphrase(ABC):

ai21/resources/bases/segmentation_base.py → ai21/clients/common/segmentation_base.py

@@ -1,8 +1,8 @@
 from abc import ABC, abstractmethod
 from typing import Any, Dict
 
-from ai21.resources.models.document_type import DocumentType
-from ai21.resources.responses.segmentation_response import SegmentationResponse
+from ai21.models.document_type import DocumentType
+from ai21.models.responses.segmentation_response import SegmentationResponse
 
 
 class Segmentation(ABC):

ai21/resources/bases/summarize_base.py → ai21/clients/common/summarize_base.py

@@ -1,8 +1,8 @@
 from abc import ABC, abstractmethod
 from typing import Optional, Any, Dict
 
-from ai21.resources.models.summary_method import SummaryMethod
-from ai21.resources.responses.summarize_response import SummarizeResponse
+from ai21.models.responses.summarize_response import SummarizeResponse
+from ai21.models.summary_method import SummaryMethod
 
 
 class Summarize(ABC):

ai21/resources/bases/summarize_by_segment_base.py → ai21/clients/common/summarize_by_segment_base.py

@@ -1,8 +1,8 @@
 from abc import ABC, abstractmethod
 from typing import Optional, Any, Dict
 
-from ai21.resources.models.document_type import DocumentType
-from ai21.resources.responses.summarize_by_segment_response import (
+from ai21.models.document_type import DocumentType
+from ai21.models.responses.summarize_by_segment_response import (
     SummarizeBySegmentResponse,
 )
 

ai21/clients/sagemaker/resources/sagemaker_answer.py

@@ -1,8 +1,8 @@
 from typing import Optional
 
-from ai21.resources.bases.answer_base import Answer, AnswerLength, Mode
-from ai21.resources.responses.answer_response import AnswerResponse
-from ai21.resources.sagemaker_resource import SageMakerResource
+from ai21.clients.common.answer_base import Answer, AnswerLength, Mode
+from ai21.models.responses.answer_response import AnswerResponse
+from ai21.clients.sagemaker.resources.sagemaker_resource import SageMakerResource
 
 
 class SageMakerAnswer(SageMakerResource, Answer):

ai21/clients/sagemaker/resources/sagemaker_completion.py

@@ -1,8 +1,8 @@
 from typing import Optional, List
 
 from ai21.resources import Penalty
-from ai21.resources.responses.completion_response import CompletionsResponse
-from ai21.resources.sagemaker_resource import SageMakerResource
+from ai21.models.responses.completion_response import CompletionsResponse
+from ai21.clients.sagemaker.resources.sagemaker_resource import SageMakerResource
 
 
 class SageMakerCompletion(SageMakerResource):

ai21/clients/sagemaker/resources/sagemaker_gec.py

@@ -1,6 +1,6 @@
-from ai21.resources.bases.gec_base import GEC
-from ai21.resources.responses.gec_response import GECResponse
-from ai21.resources.sagemaker_resource import SageMakerResource
+from ai21.clients.common.gec_base import GEC
+from ai21.models.responses.gec_response import GECResponse
+from ai21.clients.sagemaker.resources.sagemaker_resource import SageMakerResource
 
 
 class SageMakerGEC(SageMakerResource, GEC):

ai21/clients/sagemaker/resources/sagemaker_paraphrase.py

@@ -1,9 +1,9 @@
 from typing import Optional
 
-from ai21.resources.bases.paraphrase_base import Paraphrase
-from ai21.resources.models.paraphrase_style_type import ParaphraseStyleType
-from ai21.resources.responses.paraphrase_response import ParaphraseResponse
-from ai21.resources.sagemaker_resource import SageMakerResource
+from ai21.clients.common.paraphrase_base import Paraphrase
+from ai21.models.paraphrase_style_type import ParaphraseStyleType
+from ai21.models.responses.paraphrase_response import ParaphraseResponse
+from ai21.clients.sagemaker.resources.sagemaker_resource import SageMakerResource
 
 
 class SageMakerParaphrase(SageMakerResource, Paraphrase):

ai21/clients/sagemaker/resources/sagemaker_summarize.py

@@ -2,10 +2,10 @@
 
 from typing import Optional
 
-from ai21.resources.bases.summarize_base import Summarize
-from ai21.resources.models.summary_method import SummaryMethod
-from ai21.resources.responses.summarize_response import SummarizeResponse
-from ai21.resources.sagemaker_resource import SageMakerResource
+from ai21.clients.common.summarize_base import Summarize
+from ai21.models.summary_method import SummaryMethod
+from ai21.models.responses import SummarizeResponse
+from ai21.clients.sagemaker.resources.sagemaker_resource import SageMakerResource
 
 
 class SageMakerSummarize(SageMakerResource, Summarize):

ai21/clients/studio/resources/studio_answer.py

@@ -1,8 +1,9 @@
 from typing import Optional
 
-from ai21.resources.bases.answer_base import Answer, AnswerLength, Mode
-from ai21.resources.responses.answer_response import AnswerResponse
-from ai21.resources.studio_resource import StudioResource
+from ai21.clients.common.answer_base import Answer
+from ai21.models import AnswerLength, Mode
+from ai21.models.responses.answer_response import AnswerResponse
+from ai21.clients.studio.resources.studio_resource import StudioResource
 
 
 class StudioAnswer(StudioResource, Answer):

ai21/clients/studio/resources/studio_chat.py

@@ -1,15 +1,16 @@
 from typing import List, Any, Optional, Dict
 
-from ai21.resources.bases.chat_base import Chat, Message
-from ai21.resources.responses.chat_response import ChatResponse
-from ai21.resources.studio_resource import StudioResource
+from ai21.clients.common.chat_base import Chat
+from ai21.clients.studio.resources.studio_resource import StudioResource
+from ai21.models.chat_message import ChatMessage
+from ai21.models.responses.chat_response import ChatResponse
 
 
 class StudioChat(StudioResource, Chat):
     def create(
         self,
         model: str,
-        messages: List[Message],
+        messages: List[ChatMessage],
         system: str,
         *,
         num_results: Optional[int] = 1,

ai21/clients/studio/resources/studio_completion.py

@@ -1,9 +1,9 @@
 from typing import Optional, List
 
-from ai21.resources import Penalty
-from ai21.resources.bases.completion_base import Completion
-from ai21.resources.responses.completion_response import CompletionsResponse
-from ai21.resources.studio_resource import StudioResource
+from ai21.clients.common.completion_base import Completion
+from ai21.clients.studio.resources.studio_resource import StudioResource
+from ai21.models import Penalty
+from ai21.models.responses.completion_response import CompletionsResponse
 
 
 class StudioCompletion(StudioResource, Completion):

ai21/clients/studio/resources/studio_custom_model.py

@@ -1,8 +1,8 @@
 from typing import List, Optional
 
-from ai21.resources.bases.custom_model_base import CustomModel
-from ai21.resources.responses.custom_model_response import CustomBaseModelResponse
-from ai21.resources.studio_resource import StudioResource
+from ai21.clients.common.custom_model_base import CustomModel
+from ai21.models.responses.custom_model_response import CustomBaseModelResponse
+from ai21.clients.studio.resources.studio_resource import StudioResource
 
 
 class StudioCustomModel(StudioResource, CustomModel):

ai21/clients/studio/resources/studio_dataset.py

@@ -1,8 +1,8 @@
 from typing import Optional, List
 
-from ai21.resources.bases.dataset_base import Dataset
-from ai21.resources.responses.dataset_response import DatasetResponse
-from ai21.resources.studio_resource import StudioResource
+from ai21.clients.common.dataset_base import Dataset
+from ai21.models.responses.dataset_response import DatasetResponse
+from ai21.clients.studio.resources.studio_resource import StudioResource
 
 
 class StudioDataset(StudioResource, Dataset):

ai21/clients/studio/resources/studio_embed.py

@@ -1,8 +1,8 @@
 from typing import List, Optional
 
-from ai21.resources.bases.embed_base import Embed
-from ai21.resources.responses.embed_response import EmbedResponse
-from ai21.resources.studio_resource import StudioResource
+from ai21.clients.common.embed_base import Embed
+from ai21.clients.studio.resources.studio_resource import StudioResource
+from ai21.models.responses.embed_response import EmbedResponse
 
 
 class StudioEmbed(StudioResource, Embed):

ai21/clients/studio/resources/studio_gec.py

@@ -1,6 +1,6 @@
-from ai21.resources.bases.gec_base import GEC
-from ai21.resources.responses.gec_response import GECResponse
-from ai21.resources.studio_resource import StudioResource
+from ai21.clients.common.gec_base import GEC
+from ai21.models.responses.gec_response import GECResponse
+from ai21.clients.studio.resources.studio_resource import StudioResource
 
 
 class StudioGEC(StudioResource, GEC):

ai21/clients/studio/resources/studio_improvements.py

@@ -1,10 +1,10 @@
 from typing import List
 
+from ai21.clients.common.improvements_base import Improvements
+from ai21.clients.studio.resources.studio_resource import StudioResource
 from ai21.errors import EmptyMandatoryListError
-from ai21.resources.bases.improvements_base import Improvements
-from ai21.resources.models.improvement_type import ImprovementType
-from ai21.resources.responses.improvement_response import ImprovementsResponse
-from ai21.resources.studio_resource import StudioResource
+from ai21.models import ImprovementType
+from ai21.models.responses.improvement_response import ImprovementsResponse
 
 
 class StudioImprovements(StudioResource, Improvements):