From de3dd2047a34a1b122db51222be5dc43d55764b0 Mon Sep 17 00:00:00 2001 From: Asaf Gardin Date: Tue, 2 Jan 2024 14:27:14 +0200 Subject: [PATCH 1/7] refactor: moved classes to models package --- .../resources/sagemaker_paraphrase.py | 2 +- .../resources/sagemaker_summarize.py | 2 +- .../studio/resources/studio_improvements.py | 2 +- .../studio/resources/studio_library.py | 4 ++-- .../studio/resources/studio_paraphrase.py | 2 +- .../studio/resources/studio_segmentation.py | 2 +- .../studio/resources/studio_summarize.py | 2 +- .../resources/studio_summarize_by_segment.py | 2 +- ai21/models/__init__.py | 24 +++++++++++++++++++ ai21/{resources => }/models/answer_length.py | 0 ai21/{resources => }/models/document_type.py | 0 ai21/models/embed_type.py | 6 +++++ .../models/improvement_type.py | 0 ai21/models/message.py | 12 ++++++++++ ai21/{resources => }/models/mode.py | 0 .../models/paraphrase_style_type.py | 0 ai21/{resources => }/models/penalty.py | 0 ai21/{resources => }/models/role_type.py | 0 ai21/{resources => }/models/summary_method.py | 0 ai21/resources/__init__.py | 23 ------------------ ai21/resources/bases/answer_base.py | 3 +-- ai21/resources/bases/chat_base.py | 12 +--------- ai21/resources/bases/completion_base.py | 2 +- ai21/resources/bases/embed_base.py | 7 +----- ai21/resources/bases/improvements_base.py | 2 +- ai21/resources/bases/paraphrase_base.py | 2 +- ai21/resources/bases/segmentation_base.py | 2 +- ai21/resources/bases/summarize_base.py | 2 +- .../bases/summarize_by_segment_base.py | 2 +- ai21/resources/models/__init__.py | 0 ai21/resources/responses/chat_response.py | 2 +- 31 files changed, 61 insertions(+), 58 deletions(-) rename ai21/{resources => }/models/answer_length.py (100%) rename ai21/{resources => }/models/document_type.py (100%) create mode 100644 ai21/models/embed_type.py rename ai21/{resources => }/models/improvement_type.py (100%) create mode 100644 ai21/models/message.py rename ai21/{resources => }/models/mode.py (100%) rename ai21/{resources => }/models/paraphrase_style_type.py (100%) rename ai21/{resources => }/models/penalty.py (100%) rename ai21/{resources => }/models/role_type.py (100%) rename ai21/{resources => }/models/summary_method.py (100%) delete mode 100644 ai21/resources/models/__init__.py diff --git a/ai21/clients/sagemaker/resources/sagemaker_paraphrase.py b/ai21/clients/sagemaker/resources/sagemaker_paraphrase.py index 49d9ce5d..90a396b6 100644 --- a/ai21/clients/sagemaker/resources/sagemaker_paraphrase.py +++ b/ai21/clients/sagemaker/resources/sagemaker_paraphrase.py @@ -1,7 +1,7 @@ from typing import Optional from ai21.resources.bases.paraphrase_base import Paraphrase -from ai21.resources.models.paraphrase_style_type import ParaphraseStyleType +from ai21.models.paraphrase_style_type import ParaphraseStyleType from ai21.resources.responses.paraphrase_response import ParaphraseResponse from ai21.resources.sagemaker_resource import SageMakerResource diff --git a/ai21/clients/sagemaker/resources/sagemaker_summarize.py b/ai21/clients/sagemaker/resources/sagemaker_summarize.py index c7db557b..ebb9e139 100644 --- a/ai21/clients/sagemaker/resources/sagemaker_summarize.py +++ b/ai21/clients/sagemaker/resources/sagemaker_summarize.py @@ -3,7 +3,7 @@ from typing import Optional from ai21.resources.bases.summarize_base import Summarize -from ai21.resources.models.summary_method import SummaryMethod +from ai21.models.summary_method import SummaryMethod from ai21.resources.responses.summarize_response import SummarizeResponse from ai21.resources.sagemaker_resource import SageMakerResource diff --git a/ai21/clients/studio/resources/studio_improvements.py b/ai21/clients/studio/resources/studio_improvements.py index 2e17cfd9..3c503ed3 100644 --- a/ai21/clients/studio/resources/studio_improvements.py +++ b/ai21/clients/studio/resources/studio_improvements.py @@ -2,7 +2,7 @@ from ai21.errors import EmptyMandatoryListError from ai21.resources.bases.improvements_base import Improvements -from ai21.resources.models.improvement_type import ImprovementType +from ai21.models.improvement_type import ImprovementType from ai21.resources.responses.improvement_response import ImprovementsResponse from ai21.resources.studio_resource import StudioResource diff --git a/ai21/clients/studio/resources/studio_library.py b/ai21/clients/studio/resources/studio_library.py index 1962cd58..0cc306ed 100644 --- a/ai21/clients/studio/resources/studio_library.py +++ b/ai21/clients/studio/resources/studio_library.py @@ -1,8 +1,8 @@ from typing import Optional, List from ai21.ai21_http_client import AI21HTTPClient -from ai21.resources.models.answer_length import AnswerLength -from ai21.resources.models.mode import Mode +from ai21.models.answer_length import AnswerLength +from ai21.models.mode import Mode from ai21.resources.responses.file_response import FileResponse from ai21.resources.responses.library_answer_response import LibraryAnswerResponse from ai21.resources.responses.library_search_response import LibrarySearchResponse diff --git a/ai21/clients/studio/resources/studio_paraphrase.py b/ai21/clients/studio/resources/studio_paraphrase.py index 1dbf06fb..944f1bee 100644 --- a/ai21/clients/studio/resources/studio_paraphrase.py +++ b/ai21/clients/studio/resources/studio_paraphrase.py @@ -1,7 +1,7 @@ from typing import Optional from ai21.resources.bases.paraphrase_base import Paraphrase -from ai21.resources.models.paraphrase_style_type import ParaphraseStyleType +from ai21.models.paraphrase_style_type import ParaphraseStyleType from ai21.resources.responses.paraphrase_response import ParaphraseResponse from ai21.resources.studio_resource import StudioResource diff --git a/ai21/clients/studio/resources/studio_segmentation.py b/ai21/clients/studio/resources/studio_segmentation.py index 0586561f..623c5dda 100644 --- a/ai21/clients/studio/resources/studio_segmentation.py +++ b/ai21/clients/studio/resources/studio_segmentation.py @@ -1,5 +1,5 @@ from ai21.resources.bases.segmentation_base import Segmentation -from ai21.resources.models.document_type import DocumentType +from ai21.models.document_type import DocumentType from ai21.resources.responses.segmentation_response import SegmentationResponse from ai21.resources.studio_resource import StudioResource diff --git a/ai21/clients/studio/resources/studio_summarize.py b/ai21/clients/studio/resources/studio_summarize.py index 1b9c44d3..bcc2a1ba 100644 --- a/ai21/clients/studio/resources/studio_summarize.py +++ b/ai21/clients/studio/resources/studio_summarize.py @@ -1,7 +1,7 @@ from typing import Optional from ai21.resources.bases.summarize_base import Summarize -from ai21.resources.models.summary_method import SummaryMethod +from ai21.models.summary_method import SummaryMethod from ai21.resources.responses.summarize_response import SummarizeResponse from ai21.resources.studio_resource import StudioResource diff --git a/ai21/clients/studio/resources/studio_summarize_by_segment.py b/ai21/clients/studio/resources/studio_summarize_by_segment.py index 7644a5f4..3392befb 100644 --- a/ai21/clients/studio/resources/studio_summarize_by_segment.py +++ b/ai21/clients/studio/resources/studio_summarize_by_segment.py @@ -1,7 +1,7 @@ from typing import Optional from ai21.resources.bases.summarize_by_segment_base import SummarizeBySegment -from ai21.resources.models.document_type import DocumentType +from ai21.models.document_type import DocumentType from ai21.resources.responses.summarize_by_segment_response import ( SummarizeBySegmentResponse, ) diff --git a/ai21/models/__init__.py b/ai21/models/__init__.py index e69de29b..393d5c28 100644 --- a/ai21/models/__init__.py +++ b/ai21/models/__init__.py @@ -0,0 +1,24 @@ +from ai21.models.answer_length import AnswerLength +from ai21.models.document_type import DocumentType +from ai21.models.embed_type import EmbedType +from ai21.models.improvement_type import ImprovementType +from ai21.models.message import Message +from ai21.models.mode import Mode +from ai21.models.paraphrase_style_type import ParaphraseStyleType +from ai21.models.penalty import Penalty +from ai21.models.role_type import RoleType +from ai21.models.summary_method import SummaryMethod + + +__all__ = [ + "AnswerLength", + "Mode", + "Message", + "RoleType", + "Penalty", + "EmbedType", + "ImprovementType", + "ParaphraseStyleType", + "DocumentType", + "SummaryMethod", +] diff --git a/ai21/resources/models/answer_length.py b/ai21/models/answer_length.py similarity index 100% rename from ai21/resources/models/answer_length.py rename to ai21/models/answer_length.py diff --git a/ai21/resources/models/document_type.py b/ai21/models/document_type.py similarity index 100% rename from ai21/resources/models/document_type.py rename to ai21/models/document_type.py diff --git a/ai21/models/embed_type.py b/ai21/models/embed_type.py new file mode 100644 index 00000000..d1268a86 --- /dev/null +++ b/ai21/models/embed_type.py @@ -0,0 +1,6 @@ +from enum import Enum + + +class EmbedType(str, Enum): + QUERY = "query" + SEGMENT = "segment" diff --git a/ai21/resources/models/improvement_type.py b/ai21/models/improvement_type.py similarity index 100% rename from ai21/resources/models/improvement_type.py rename to ai21/models/improvement_type.py diff --git a/ai21/models/message.py b/ai21/models/message.py new file mode 100644 index 00000000..12937a3b --- /dev/null +++ b/ai21/models/message.py @@ -0,0 +1,12 @@ +from dataclasses import dataclass +from typing import Optional + +from ai21.models import RoleType +from ai21.models.ai21_base_model_mixin import AI21BaseModelMixin + + +@dataclass +class Message(AI21BaseModelMixin): + role: RoleType + text: str + name: Optional[str] = None diff --git a/ai21/resources/models/mode.py b/ai21/models/mode.py similarity index 100% rename from ai21/resources/models/mode.py rename to ai21/models/mode.py diff --git a/ai21/resources/models/paraphrase_style_type.py b/ai21/models/paraphrase_style_type.py similarity index 100% rename from ai21/resources/models/paraphrase_style_type.py rename to ai21/models/paraphrase_style_type.py diff --git a/ai21/resources/models/penalty.py b/ai21/models/penalty.py similarity index 100% rename from ai21/resources/models/penalty.py rename to ai21/models/penalty.py diff --git a/ai21/resources/models/role_type.py b/ai21/models/role_type.py similarity index 100% rename from ai21/resources/models/role_type.py rename to ai21/models/role_type.py diff --git a/ai21/resources/models/summary_method.py b/ai21/models/summary_method.py similarity index 100% rename from ai21/resources/models/summary_method.py rename to ai21/models/summary_method.py diff --git a/ai21/resources/__init__.py b/ai21/resources/__init__.py index 626c7221..e69de29b 100644 --- a/ai21/resources/__init__.py +++ b/ai21/resources/__init__.py @@ -1,23 +0,0 @@ -from ai21.resources.bases.chat_base import Message -from ai21.resources.bases.embed_base import EmbedType -from ai21.resources.models.answer_length import AnswerLength -from ai21.resources.models.document_type import DocumentType -from ai21.resources.models.improvement_type import ImprovementType -from ai21.resources.models.mode import Mode -from ai21.resources.models.paraphrase_style_type import ParaphraseStyleType -from ai21.resources.models.penalty import Penalty -from ai21.resources.models.role_type import RoleType -from ai21.resources.models.summary_method import SummaryMethod - -__all__ = [ - "AnswerLength", - "Mode", - "Message", - "RoleType", - "Penalty", - "EmbedType", - "ImprovementType", - "ParaphraseStyleType", - "DocumentType", - "SummaryMethod", -] diff --git a/ai21/resources/bases/answer_base.py b/ai21/resources/bases/answer_base.py index 96663084..c64a6046 100644 --- a/ai21/resources/bases/answer_base.py +++ b/ai21/resources/bases/answer_base.py @@ -1,8 +1,7 @@ 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.models import Mode, AnswerLength from ai21.resources.responses.answer_response import AnswerResponse diff --git a/ai21/resources/bases/chat_base.py b/ai21/resources/bases/chat_base.py index 54854bb7..c3eeb73c 100644 --- a/ai21/resources/bases/chat_base.py +++ b/ai21/resources/bases/chat_base.py @@ -1,20 +1,10 @@ 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.models import Message, Penalty from ai21.resources.responses.chat_response import ChatResponse -@dataclass -class Message(AI21BaseModelMixin): - role: RoleType - text: str - name: Optional[str] - - class Chat(ABC): _module_name = "chat" diff --git a/ai21/resources/bases/completion_base.py b/ai21/resources/bases/completion_base.py index df9f2352..693e1c6d 100644 --- a/ai21/resources/bases/completion_base.py +++ b/ai21/resources/bases/completion_base.py @@ -1,7 +1,7 @@ from abc import ABC, abstractmethod from typing import Optional, List, Dict, Any -from ai21.resources import Penalty +from ai21.models import Penalty from ai21.resources.responses.completion_response import CompletionsResponse diff --git a/ai21/resources/bases/embed_base.py b/ai21/resources/bases/embed_base.py index ee07187b..534b01b8 100644 --- a/ai21/resources/bases/embed_base.py +++ b/ai21/resources/bases/embed_base.py @@ -1,15 +1,10 @@ from abc import ABC, abstractmethod -from enum import Enum from typing import List, Any, Dict, Optional +from ai21.models import EmbedType from ai21.resources.responses.embed_response import EmbedResponse -class EmbedType(str, Enum): - QUERY = "query" - SEGMENT = "segment" - - class Embed(ABC): _module_name = "embed" diff --git a/ai21/resources/bases/improvements_base.py b/ai21/resources/bases/improvements_base.py index 69ac49e7..f406406a 100644 --- a/ai21/resources/bases/improvements_base.py +++ b/ai21/resources/bases/improvements_base.py @@ -1,7 +1,7 @@ from abc import ABC, abstractmethod from typing import Any, Dict, List -from ai21.resources.models.improvement_type import ImprovementType +from ai21.models import ImprovementType from ai21.resources.responses.improvement_response import ImprovementsResponse diff --git a/ai21/resources/bases/paraphrase_base.py b/ai21/resources/bases/paraphrase_base.py index aa2e2e7d..a9d250a0 100644 --- a/ai21/resources/bases/paraphrase_base.py +++ b/ai21/resources/bases/paraphrase_base.py @@ -1,7 +1,7 @@ from abc import ABC, abstractmethod from typing import Optional, Any, Dict -from ai21.resources.models.paraphrase_style_type import ParaphraseStyleType +from ai21.models import ParaphraseStyleType from ai21.resources.responses.paraphrase_response import ParaphraseResponse diff --git a/ai21/resources/bases/segmentation_base.py b/ai21/resources/bases/segmentation_base.py index f20cecb3..aad1f4c5 100644 --- a/ai21/resources/bases/segmentation_base.py +++ b/ai21/resources/bases/segmentation_base.py @@ -1,7 +1,7 @@ from abc import ABC, abstractmethod from typing import Any, Dict -from ai21.resources.models.document_type import DocumentType +from ai21.models import DocumentType from ai21.resources.responses.segmentation_response import SegmentationResponse diff --git a/ai21/resources/bases/summarize_base.py b/ai21/resources/bases/summarize_base.py index 435549a5..3a4754fd 100644 --- a/ai21/resources/bases/summarize_base.py +++ b/ai21/resources/bases/summarize_base.py @@ -1,7 +1,7 @@ from abc import ABC, abstractmethod from typing import Optional, Any, Dict -from ai21.resources.models.summary_method import SummaryMethod +from ai21.models import SummaryMethod from ai21.resources.responses.summarize_response import SummarizeResponse diff --git a/ai21/resources/bases/summarize_by_segment_base.py b/ai21/resources/bases/summarize_by_segment_base.py index ac30d1ae..542bd886 100644 --- a/ai21/resources/bases/summarize_by_segment_base.py +++ b/ai21/resources/bases/summarize_by_segment_base.py @@ -1,7 +1,7 @@ from abc import ABC, abstractmethod from typing import Optional, Any, Dict -from ai21.resources.models.document_type import DocumentType +from ai21.models import DocumentType from ai21.resources.responses.summarize_by_segment_response import ( SummarizeBySegmentResponse, ) diff --git a/ai21/resources/models/__init__.py b/ai21/resources/models/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/ai21/resources/responses/chat_response.py b/ai21/resources/responses/chat_response.py index 81ce3062..e1a15a9f 100644 --- a/ai21/resources/responses/chat_response.py +++ b/ai21/resources/responses/chat_response.py @@ -2,7 +2,7 @@ from typing import Optional, List from ai21.models.ai21_base_model_mixin import AI21BaseModelMixin -from ai21.resources.models.role_type import RoleType +from ai21.models.role_type import RoleType @dataclass From ff0bbd88677c1718d790520dbc0d14eeae52c9d9 Mon Sep 17 00:00:00 2001 From: Asaf Gardin Date: Tue, 2 Jan 2024 14:27:55 +0200 Subject: [PATCH 2/7] refactor: moved responses = to models package --- ai21/__init__.py | 15 --------------- .../bedrock/resources/bedrock_completion.py | 2 +- .../sagemaker/resources/sagemaker_answer.py | 2 +- .../sagemaker/resources/sagemaker_completion.py | 2 +- ai21/clients/sagemaker/resources/sagemaker_gec.py | 2 +- .../sagemaker/resources/sagemaker_paraphrase.py | 2 +- .../sagemaker/resources/sagemaker_summarize.py | 2 +- ai21/clients/studio/resources/studio_answer.py | 2 +- ai21/clients/studio/resources/studio_chat.py | 2 +- .../clients/studio/resources/studio_completion.py | 2 +- .../studio/resources/studio_custom_model.py | 2 +- ai21/clients/studio/resources/studio_dataset.py | 2 +- ai21/clients/studio/resources/studio_embed.py | 2 +- ai21/clients/studio/resources/studio_gec.py | 2 +- .../studio/resources/studio_improvements.py | 2 +- ai21/clients/studio/resources/studio_library.py | 6 +++--- .../clients/studio/resources/studio_paraphrase.py | 2 +- .../studio/resources/studio_segmentation.py | 2 +- ai21/clients/studio/resources/studio_summarize.py | 2 +- .../resources/studio_summarize_by_segment.py | 2 +- ai21/{resources => models}/responses/__init__.py | 0 .../responses/answer_response.py | 0 .../responses/chat_response.py | 0 .../responses/completion_response.py | 0 .../responses/custom_model_response.py | 0 .../responses/dataset_response.py | 0 .../responses/embed_response.py | 0 .../responses/file_response.py | 0 .../responses/gec_response.py | 0 .../responses/improvement_response.py | 0 .../responses/library_answer_response.py | 0 .../responses/library_search_response.py | 0 .../responses/paraphrase_response.py | 0 .../responses/segmentation_response.py | 0 .../responses/summarize_by_segment_response.py | 0 .../responses/summarize_response.py | 0 ai21/resources/bases/answer_base.py | 2 +- ai21/resources/bases/chat_base.py | 2 +- ai21/resources/bases/completion_base.py | 2 +- ai21/resources/bases/custom_model_base.py | 2 +- ai21/resources/bases/dataset_base.py | 2 +- ai21/resources/bases/embed_base.py | 2 +- ai21/resources/bases/gec_base.py | 2 +- ai21/resources/bases/improvements_base.py | 2 +- ai21/resources/bases/paraphrase_base.py | 2 +- ai21/resources/bases/segmentation_base.py | 2 +- ai21/resources/bases/summarize_base.py | 2 +- ai21/resources/bases/summarize_by_segment_base.py | 2 +- .../clients/studio/resources/conftest.py | 4 ++-- 49 files changed, 35 insertions(+), 50 deletions(-) rename ai21/{resources => models}/responses/__init__.py (100%) rename ai21/{resources => models}/responses/answer_response.py (100%) rename ai21/{resources => models}/responses/chat_response.py (100%) rename ai21/{resources => models}/responses/completion_response.py (100%) rename ai21/{resources => models}/responses/custom_model_response.py (100%) rename ai21/{resources => models}/responses/dataset_response.py (100%) rename ai21/{resources => models}/responses/embed_response.py (100%) rename ai21/{resources => models}/responses/file_response.py (100%) rename ai21/{resources => models}/responses/gec_response.py (100%) rename ai21/{resources => models}/responses/improvement_response.py (100%) rename ai21/{resources => models}/responses/library_answer_response.py (100%) rename ai21/{resources => models}/responses/library_search_response.py (100%) rename ai21/{resources => models}/responses/paraphrase_response.py (100%) rename ai21/{resources => models}/responses/segmentation_response.py (100%) rename ai21/{resources => models}/responses/summarize_by_segment_response.py (100%) rename ai21/{resources => models}/responses/summarize_response.py (100%) diff --git a/ai21/__init__.py b/ai21/__init__.py index 6c5fb3e9..5fa7b336 100644 --- a/ai21/__init__.py +++ b/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 diff --git a/ai21/clients/bedrock/resources/bedrock_completion.py b/ai21/clients/bedrock/resources/bedrock_completion.py index 150cf381..578cf215 100644 --- a/ai21/clients/bedrock/resources/bedrock_completion.py +++ b/ai21/clients/bedrock/resources/bedrock_completion.py @@ -2,7 +2,7 @@ from ai21.resources import Penalty from ai21.resources.bedrock_resource import BedrockResource -from ai21.resources.responses.completion_response import CompletionsResponse +from ai21.models.responses.completion_response import CompletionsResponse class BedrockCompletion(BedrockResource): diff --git a/ai21/clients/sagemaker/resources/sagemaker_answer.py b/ai21/clients/sagemaker/resources/sagemaker_answer.py index 1584abc9..1ed83464 100644 --- a/ai21/clients/sagemaker/resources/sagemaker_answer.py +++ b/ai21/clients/sagemaker/resources/sagemaker_answer.py @@ -1,7 +1,7 @@ from typing import Optional from ai21.resources.bases.answer_base import Answer, AnswerLength, Mode -from ai21.resources.responses.answer_response import AnswerResponse +from ai21.models.responses.answer_response import AnswerResponse from ai21.resources.sagemaker_resource import SageMakerResource diff --git a/ai21/clients/sagemaker/resources/sagemaker_completion.py b/ai21/clients/sagemaker/resources/sagemaker_completion.py index 687f9b94..55408607 100644 --- a/ai21/clients/sagemaker/resources/sagemaker_completion.py +++ b/ai21/clients/sagemaker/resources/sagemaker_completion.py @@ -1,7 +1,7 @@ from typing import Optional, List from ai21.resources import Penalty -from ai21.resources.responses.completion_response import CompletionsResponse +from ai21.models.responses.completion_response import CompletionsResponse from ai21.resources.sagemaker_resource import SageMakerResource diff --git a/ai21/clients/sagemaker/resources/sagemaker_gec.py b/ai21/clients/sagemaker/resources/sagemaker_gec.py index db8717e8..5b593622 100644 --- a/ai21/clients/sagemaker/resources/sagemaker_gec.py +++ b/ai21/clients/sagemaker/resources/sagemaker_gec.py @@ -1,5 +1,5 @@ from ai21.resources.bases.gec_base import GEC -from ai21.resources.responses.gec_response import GECResponse +from ai21.models.responses.gec_response import GECResponse from ai21.resources.sagemaker_resource import SageMakerResource diff --git a/ai21/clients/sagemaker/resources/sagemaker_paraphrase.py b/ai21/clients/sagemaker/resources/sagemaker_paraphrase.py index 90a396b6..54c31eb4 100644 --- a/ai21/clients/sagemaker/resources/sagemaker_paraphrase.py +++ b/ai21/clients/sagemaker/resources/sagemaker_paraphrase.py @@ -2,7 +2,7 @@ from ai21.resources.bases.paraphrase_base import Paraphrase from ai21.models.paraphrase_style_type import ParaphraseStyleType -from ai21.resources.responses.paraphrase_response import ParaphraseResponse +from ai21.models.responses.paraphrase_response import ParaphraseResponse from ai21.resources.sagemaker_resource import SageMakerResource diff --git a/ai21/clients/sagemaker/resources/sagemaker_summarize.py b/ai21/clients/sagemaker/resources/sagemaker_summarize.py index ebb9e139..18491534 100644 --- a/ai21/clients/sagemaker/resources/sagemaker_summarize.py +++ b/ai21/clients/sagemaker/resources/sagemaker_summarize.py @@ -4,7 +4,7 @@ from ai21.resources.bases.summarize_base import Summarize from ai21.models.summary_method import SummaryMethod -from ai21.resources.responses.summarize_response import SummarizeResponse +from ai21.models.responses import SummarizeResponse from ai21.resources.sagemaker_resource import SageMakerResource diff --git a/ai21/clients/studio/resources/studio_answer.py b/ai21/clients/studio/resources/studio_answer.py index 3962353c..cbc2ca17 100644 --- a/ai21/clients/studio/resources/studio_answer.py +++ b/ai21/clients/studio/resources/studio_answer.py @@ -1,7 +1,7 @@ from typing import Optional from ai21.resources.bases.answer_base import Answer, AnswerLength, Mode -from ai21.resources.responses.answer_response import AnswerResponse +from ai21.models.responses.answer_response import AnswerResponse from ai21.resources.studio_resource import StudioResource diff --git a/ai21/clients/studio/resources/studio_chat.py b/ai21/clients/studio/resources/studio_chat.py index f1dab12b..fb757ae0 100644 --- a/ai21/clients/studio/resources/studio_chat.py +++ b/ai21/clients/studio/resources/studio_chat.py @@ -1,7 +1,7 @@ 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.models.responses import ChatResponse from ai21.resources.studio_resource import StudioResource diff --git a/ai21/clients/studio/resources/studio_completion.py b/ai21/clients/studio/resources/studio_completion.py index bcac84f8..1d1318d8 100644 --- a/ai21/clients/studio/resources/studio_completion.py +++ b/ai21/clients/studio/resources/studio_completion.py @@ -2,7 +2,7 @@ from ai21.resources import Penalty from ai21.resources.bases.completion_base import Completion -from ai21.resources.responses.completion_response import CompletionsResponse +from ai21.models.responses.completion_response import CompletionsResponse from ai21.resources.studio_resource import StudioResource diff --git a/ai21/clients/studio/resources/studio_custom_model.py b/ai21/clients/studio/resources/studio_custom_model.py index 32166e53..82718138 100644 --- a/ai21/clients/studio/resources/studio_custom_model.py +++ b/ai21/clients/studio/resources/studio_custom_model.py @@ -1,7 +1,7 @@ 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.models.responses.custom_model_response import CustomBaseModelResponse from ai21.resources.studio_resource import StudioResource diff --git a/ai21/clients/studio/resources/studio_dataset.py b/ai21/clients/studio/resources/studio_dataset.py index 8626d71b..351391ab 100644 --- a/ai21/clients/studio/resources/studio_dataset.py +++ b/ai21/clients/studio/resources/studio_dataset.py @@ -1,7 +1,7 @@ from typing import Optional, List from ai21.resources.bases.dataset_base import Dataset -from ai21.resources.responses.dataset_response import DatasetResponse +from ai21.models.responses.dataset_response import DatasetResponse from ai21.resources.studio_resource import StudioResource diff --git a/ai21/clients/studio/resources/studio_embed.py b/ai21/clients/studio/resources/studio_embed.py index 80cad1c7..e2cde9d9 100644 --- a/ai21/clients/studio/resources/studio_embed.py +++ b/ai21/clients/studio/resources/studio_embed.py @@ -1,7 +1,7 @@ from typing import List, Optional from ai21.resources.bases.embed_base import Embed -from ai21.resources.responses.embed_response import EmbedResponse +from ai21.models.responses import EmbedResponse from ai21.resources.studio_resource import StudioResource diff --git a/ai21/clients/studio/resources/studio_gec.py b/ai21/clients/studio/resources/studio_gec.py index 3ce45a6c..1ee5b8a3 100644 --- a/ai21/clients/studio/resources/studio_gec.py +++ b/ai21/clients/studio/resources/studio_gec.py @@ -1,5 +1,5 @@ from ai21.resources.bases.gec_base import GEC -from ai21.resources.responses.gec_response import GECResponse +from ai21.models.responses.gec_response import GECResponse from ai21.resources.studio_resource import StudioResource diff --git a/ai21/clients/studio/resources/studio_improvements.py b/ai21/clients/studio/resources/studio_improvements.py index 3c503ed3..1b3f81c0 100644 --- a/ai21/clients/studio/resources/studio_improvements.py +++ b/ai21/clients/studio/resources/studio_improvements.py @@ -3,7 +3,7 @@ from ai21.errors import EmptyMandatoryListError from ai21.resources.bases.improvements_base import Improvements from ai21.models.improvement_type import ImprovementType -from ai21.resources.responses.improvement_response import ImprovementsResponse +from ai21.models.responses.improvement_response import ImprovementsResponse from ai21.resources.studio_resource import StudioResource diff --git a/ai21/clients/studio/resources/studio_library.py b/ai21/clients/studio/resources/studio_library.py index 0cc306ed..dece0cfe 100644 --- a/ai21/clients/studio/resources/studio_library.py +++ b/ai21/clients/studio/resources/studio_library.py @@ -3,9 +3,9 @@ from ai21.ai21_http_client import AI21HTTPClient from ai21.models.answer_length import AnswerLength from ai21.models.mode import Mode -from ai21.resources.responses.file_response import FileResponse -from ai21.resources.responses.library_answer_response import LibraryAnswerResponse -from ai21.resources.responses.library_search_response import LibrarySearchResponse +from ai21.models.responses.file_response import FileResponse +from ai21.models.responses.library_answer_response import LibraryAnswerResponse +from ai21.models.responses.library_search_response import LibrarySearchResponse from ai21.resources.studio_resource import StudioResource diff --git a/ai21/clients/studio/resources/studio_paraphrase.py b/ai21/clients/studio/resources/studio_paraphrase.py index 944f1bee..9202a164 100644 --- a/ai21/clients/studio/resources/studio_paraphrase.py +++ b/ai21/clients/studio/resources/studio_paraphrase.py @@ -2,7 +2,7 @@ from ai21.resources.bases.paraphrase_base import Paraphrase from ai21.models.paraphrase_style_type import ParaphraseStyleType -from ai21.resources.responses.paraphrase_response import ParaphraseResponse +from ai21.models.responses.paraphrase_response import ParaphraseResponse from ai21.resources.studio_resource import StudioResource diff --git a/ai21/clients/studio/resources/studio_segmentation.py b/ai21/clients/studio/resources/studio_segmentation.py index 623c5dda..8e49f5b5 100644 --- a/ai21/clients/studio/resources/studio_segmentation.py +++ b/ai21/clients/studio/resources/studio_segmentation.py @@ -1,6 +1,6 @@ from ai21.resources.bases.segmentation_base import Segmentation from ai21.models.document_type import DocumentType -from ai21.resources.responses.segmentation_response import SegmentationResponse +from ai21.models.responses import SegmentationResponse from ai21.resources.studio_resource import StudioResource diff --git a/ai21/clients/studio/resources/studio_summarize.py b/ai21/clients/studio/resources/studio_summarize.py index bcc2a1ba..ce44ad5d 100644 --- a/ai21/clients/studio/resources/studio_summarize.py +++ b/ai21/clients/studio/resources/studio_summarize.py @@ -2,7 +2,7 @@ from ai21.resources.bases.summarize_base import Summarize from ai21.models.summary_method import SummaryMethod -from ai21.resources.responses.summarize_response import SummarizeResponse +from ai21.models.responses import SummarizeResponse from ai21.resources.studio_resource import StudioResource diff --git a/ai21/clients/studio/resources/studio_summarize_by_segment.py b/ai21/clients/studio/resources/studio_summarize_by_segment.py index 3392befb..1dec9336 100644 --- a/ai21/clients/studio/resources/studio_summarize_by_segment.py +++ b/ai21/clients/studio/resources/studio_summarize_by_segment.py @@ -2,7 +2,7 @@ from ai21.resources.bases.summarize_by_segment_base import SummarizeBySegment from ai21.models.document_type import DocumentType -from ai21.resources.responses.summarize_by_segment_response import ( +from ai21.models.responses import ( SummarizeBySegmentResponse, ) from ai21.resources.studio_resource import StudioResource diff --git a/ai21/resources/responses/__init__.py b/ai21/models/responses/__init__.py similarity index 100% rename from ai21/resources/responses/__init__.py rename to ai21/models/responses/__init__.py diff --git a/ai21/resources/responses/answer_response.py b/ai21/models/responses/answer_response.py similarity index 100% rename from ai21/resources/responses/answer_response.py rename to ai21/models/responses/answer_response.py diff --git a/ai21/resources/responses/chat_response.py b/ai21/models/responses/chat_response.py similarity index 100% rename from ai21/resources/responses/chat_response.py rename to ai21/models/responses/chat_response.py diff --git a/ai21/resources/responses/completion_response.py b/ai21/models/responses/completion_response.py similarity index 100% rename from ai21/resources/responses/completion_response.py rename to ai21/models/responses/completion_response.py diff --git a/ai21/resources/responses/custom_model_response.py b/ai21/models/responses/custom_model_response.py similarity index 100% rename from ai21/resources/responses/custom_model_response.py rename to ai21/models/responses/custom_model_response.py diff --git a/ai21/resources/responses/dataset_response.py b/ai21/models/responses/dataset_response.py similarity index 100% rename from ai21/resources/responses/dataset_response.py rename to ai21/models/responses/dataset_response.py diff --git a/ai21/resources/responses/embed_response.py b/ai21/models/responses/embed_response.py similarity index 100% rename from ai21/resources/responses/embed_response.py rename to ai21/models/responses/embed_response.py diff --git a/ai21/resources/responses/file_response.py b/ai21/models/responses/file_response.py similarity index 100% rename from ai21/resources/responses/file_response.py rename to ai21/models/responses/file_response.py diff --git a/ai21/resources/responses/gec_response.py b/ai21/models/responses/gec_response.py similarity index 100% rename from ai21/resources/responses/gec_response.py rename to ai21/models/responses/gec_response.py diff --git a/ai21/resources/responses/improvement_response.py b/ai21/models/responses/improvement_response.py similarity index 100% rename from ai21/resources/responses/improvement_response.py rename to ai21/models/responses/improvement_response.py diff --git a/ai21/resources/responses/library_answer_response.py b/ai21/models/responses/library_answer_response.py similarity index 100% rename from ai21/resources/responses/library_answer_response.py rename to ai21/models/responses/library_answer_response.py diff --git a/ai21/resources/responses/library_search_response.py b/ai21/models/responses/library_search_response.py similarity index 100% rename from ai21/resources/responses/library_search_response.py rename to ai21/models/responses/library_search_response.py diff --git a/ai21/resources/responses/paraphrase_response.py b/ai21/models/responses/paraphrase_response.py similarity index 100% rename from ai21/resources/responses/paraphrase_response.py rename to ai21/models/responses/paraphrase_response.py diff --git a/ai21/resources/responses/segmentation_response.py b/ai21/models/responses/segmentation_response.py similarity index 100% rename from ai21/resources/responses/segmentation_response.py rename to ai21/models/responses/segmentation_response.py diff --git a/ai21/resources/responses/summarize_by_segment_response.py b/ai21/models/responses/summarize_by_segment_response.py similarity index 100% rename from ai21/resources/responses/summarize_by_segment_response.py rename to ai21/models/responses/summarize_by_segment_response.py diff --git a/ai21/resources/responses/summarize_response.py b/ai21/models/responses/summarize_response.py similarity index 100% rename from ai21/resources/responses/summarize_response.py rename to ai21/models/responses/summarize_response.py diff --git a/ai21/resources/bases/answer_base.py b/ai21/resources/bases/answer_base.py index c64a6046..a1543646 100644 --- a/ai21/resources/bases/answer_base.py +++ b/ai21/resources/bases/answer_base.py @@ -2,7 +2,7 @@ from typing import Optional, Any, Dict from ai21.models import Mode, AnswerLength -from ai21.resources.responses.answer_response import AnswerResponse +from ai21.models.responses.answer_response import AnswerResponse class Answer(ABC): diff --git a/ai21/resources/bases/chat_base.py b/ai21/resources/bases/chat_base.py index c3eeb73c..bf898942 100644 --- a/ai21/resources/bases/chat_base.py +++ b/ai21/resources/bases/chat_base.py @@ -2,7 +2,7 @@ from typing import List, Any, Dict, Optional from ai21.models import Message, Penalty -from ai21.resources.responses.chat_response import ChatResponse +from ai21.models.responses import ChatResponse class Chat(ABC): diff --git a/ai21/resources/bases/completion_base.py b/ai21/resources/bases/completion_base.py index 693e1c6d..a2bc8c3d 100644 --- a/ai21/resources/bases/completion_base.py +++ b/ai21/resources/bases/completion_base.py @@ -2,7 +2,7 @@ from typing import Optional, List, Dict, Any from ai21.models import Penalty -from ai21.resources.responses.completion_response import CompletionsResponse +from ai21.models.responses.completion_response import CompletionsResponse class Completion(ABC): diff --git a/ai21/resources/bases/custom_model_base.py b/ai21/resources/bases/custom_model_base.py index 6e360a17..7d3b55ae 100644 --- a/ai21/resources/bases/custom_model_base.py +++ b/ai21/resources/bases/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): diff --git a/ai21/resources/bases/dataset_base.py b/ai21/resources/bases/dataset_base.py index 2be49fc7..9fa57f85 100644 --- a/ai21/resources/bases/dataset_base.py +++ b/ai21/resources/bases/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): diff --git a/ai21/resources/bases/embed_base.py b/ai21/resources/bases/embed_base.py index 534b01b8..674ebfe0 100644 --- a/ai21/resources/bases/embed_base.py +++ b/ai21/resources/bases/embed_base.py @@ -2,7 +2,7 @@ from typing import List, Any, Dict, Optional from ai21.models import EmbedType -from ai21.resources.responses.embed_response import EmbedResponse +from ai21.models.responses import EmbedResponse class Embed(ABC): diff --git a/ai21/resources/bases/gec_base.py b/ai21/resources/bases/gec_base.py index 0d623dfd..8de743e2 100644 --- a/ai21/resources/bases/gec_base.py +++ b/ai21/resources/bases/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): diff --git a/ai21/resources/bases/improvements_base.py b/ai21/resources/bases/improvements_base.py index f406406a..df912e1d 100644 --- a/ai21/resources/bases/improvements_base.py +++ b/ai21/resources/bases/improvements_base.py @@ -2,7 +2,7 @@ from typing import Any, Dict, List from ai21.models import ImprovementType -from ai21.resources.responses.improvement_response import ImprovementsResponse +from ai21.models.responses.improvement_response import ImprovementsResponse class Improvements(ABC): diff --git a/ai21/resources/bases/paraphrase_base.py b/ai21/resources/bases/paraphrase_base.py index a9d250a0..917cdd75 100644 --- a/ai21/resources/bases/paraphrase_base.py +++ b/ai21/resources/bases/paraphrase_base.py @@ -2,7 +2,7 @@ from typing import Optional, Any, Dict from ai21.models import ParaphraseStyleType -from ai21.resources.responses.paraphrase_response import ParaphraseResponse +from ai21.models.responses.paraphrase_response import ParaphraseResponse class Paraphrase(ABC): diff --git a/ai21/resources/bases/segmentation_base.py b/ai21/resources/bases/segmentation_base.py index aad1f4c5..dfbb3e70 100644 --- a/ai21/resources/bases/segmentation_base.py +++ b/ai21/resources/bases/segmentation_base.py @@ -2,7 +2,7 @@ from typing import Any, Dict from ai21.models import DocumentType -from ai21.resources.responses.segmentation_response import SegmentationResponse +from ai21.models.responses import SegmentationResponse class Segmentation(ABC): diff --git a/ai21/resources/bases/summarize_base.py b/ai21/resources/bases/summarize_base.py index 3a4754fd..ba2c5683 100644 --- a/ai21/resources/bases/summarize_base.py +++ b/ai21/resources/bases/summarize_base.py @@ -2,7 +2,7 @@ from typing import Optional, Any, Dict from ai21.models import SummaryMethod -from ai21.resources.responses.summarize_response import SummarizeResponse +from ai21.models.responses import SummarizeResponse class Summarize(ABC): diff --git a/ai21/resources/bases/summarize_by_segment_base.py b/ai21/resources/bases/summarize_by_segment_base.py index 542bd886..fa16c251 100644 --- a/ai21/resources/bases/summarize_by_segment_base.py +++ b/ai21/resources/bases/summarize_by_segment_base.py @@ -2,7 +2,7 @@ from typing import Optional, Any, Dict from ai21.models import DocumentType -from ai21.resources.responses.summarize_by_segment_response import ( +from ai21.models.responses import ( SummarizeBySegmentResponse, ) diff --git a/tests/unittests/clients/studio/resources/conftest.py b/tests/unittests/clients/studio/resources/conftest.py index 8c50a082..1009a3dc 100644 --- a/tests/unittests/clients/studio/resources/conftest.py +++ b/tests/unittests/clients/studio/resources/conftest.py @@ -7,8 +7,8 @@ from ai21.clients.studio.resources.studio_chat import StudioChat from ai21.clients.studio.resources.studio_completion import StudioCompletion from ai21.resources import Message, RoleType -from ai21.resources.responses.chat_response import ChatOutput, FinishReason -from ai21.resources.responses.completion_response import Prompt, Completion, CompletionData, CompletionFinishReason +from ai21.models.responses import ChatOutput, FinishReason +from ai21.models.responses.completion_response import Prompt, Completion, CompletionData, CompletionFinishReason @pytest.fixture From 3881c96e3292c4f75ec2ece69a82bfee87c83d0e Mon Sep 17 00:00:00 2001 From: Asaf Gardin Date: Tue, 2 Jan 2024 14:31:37 +0200 Subject: [PATCH 3/7] refactor: moved resources to common package --- ai21/clients/bedrock/resources/bedrock_completion.py | 2 +- ai21/{ => clients/bedrock}/resources/bedrock_resource.py | 0 ai21/{resources/bases => clients/common}/__init__.py | 0 ai21/{resources/bases => clients/common}/answer_base.py | 0 ai21/{resources/bases => clients/common}/chat_base.py | 0 ai21/{resources/bases => clients/common}/completion_base.py | 0 ai21/{resources/bases => clients/common}/custom_model_base.py | 0 ai21/{resources/bases => clients/common}/dataset_base.py | 0 ai21/{resources/bases => clients/common}/embed_base.py | 0 ai21/{resources/bases => clients/common}/gec_base.py | 0 ai21/{resources/bases => clients/common}/improvements_base.py | 0 ai21/{resources/bases => clients/common}/paraphrase_base.py | 0 ai21/{resources/bases => clients/common}/segmentation_base.py | 0 ai21/{resources/bases => clients/common}/summarize_base.py | 0 .../bases => clients/common}/summarize_by_segment_base.py | 0 ai21/clients/sagemaker/resources/sagemaker_answer.py | 4 ++-- ai21/clients/sagemaker/resources/sagemaker_completion.py | 2 +- ai21/clients/sagemaker/resources/sagemaker_gec.py | 4 ++-- ai21/clients/sagemaker/resources/sagemaker_paraphrase.py | 4 ++-- ai21/{ => clients/sagemaker}/resources/sagemaker_resource.py | 0 ai21/clients/sagemaker/resources/sagemaker_summarize.py | 4 ++-- ai21/clients/studio/resources/studio_answer.py | 4 ++-- ai21/clients/studio/resources/studio_chat.py | 4 ++-- ai21/clients/studio/resources/studio_completion.py | 4 ++-- ai21/clients/studio/resources/studio_custom_model.py | 4 ++-- ai21/clients/studio/resources/studio_dataset.py | 4 ++-- ai21/clients/studio/resources/studio_embed.py | 4 ++-- ai21/clients/studio/resources/studio_gec.py | 4 ++-- ai21/clients/studio/resources/studio_improvements.py | 4 ++-- ai21/clients/studio/resources/studio_library.py | 2 +- ai21/clients/studio/resources/studio_paraphrase.py | 4 ++-- ai21/{ => clients/studio}/resources/studio_resource.py | 0 ai21/clients/studio/resources/studio_segmentation.py | 4 ++-- ai21/clients/studio/resources/studio_summarize.py | 4 ++-- ai21/clients/studio/resources/studio_summarize_by_segment.py | 4 ++-- .../clients/studio/resources/test_studio_resources.py | 2 +- 36 files changed, 36 insertions(+), 36 deletions(-) rename ai21/{ => clients/bedrock}/resources/bedrock_resource.py (100%) rename ai21/{resources/bases => clients/common}/__init__.py (100%) rename ai21/{resources/bases => clients/common}/answer_base.py (100%) rename ai21/{resources/bases => clients/common}/chat_base.py (100%) rename ai21/{resources/bases => clients/common}/completion_base.py (100%) rename ai21/{resources/bases => clients/common}/custom_model_base.py (100%) rename ai21/{resources/bases => clients/common}/dataset_base.py (100%) rename ai21/{resources/bases => clients/common}/embed_base.py (100%) rename ai21/{resources/bases => clients/common}/gec_base.py (100%) rename ai21/{resources/bases => clients/common}/improvements_base.py (100%) rename ai21/{resources/bases => clients/common}/paraphrase_base.py (100%) rename ai21/{resources/bases => clients/common}/segmentation_base.py (100%) rename ai21/{resources/bases => clients/common}/summarize_base.py (100%) rename ai21/{resources/bases => clients/common}/summarize_by_segment_base.py (100%) rename ai21/{ => clients/sagemaker}/resources/sagemaker_resource.py (100%) rename ai21/{ => clients/studio}/resources/studio_resource.py (100%) diff --git a/ai21/clients/bedrock/resources/bedrock_completion.py b/ai21/clients/bedrock/resources/bedrock_completion.py index 578cf215..f9ff4646 100644 --- a/ai21/clients/bedrock/resources/bedrock_completion.py +++ b/ai21/clients/bedrock/resources/bedrock_completion.py @@ -1,7 +1,7 @@ from typing import Optional, List from ai21.resources import Penalty -from ai21.resources.bedrock_resource import BedrockResource +from ai21.clients.bedrock.resources.bedrock_resource import BedrockResource from ai21.models.responses.completion_response import CompletionsResponse diff --git a/ai21/resources/bedrock_resource.py b/ai21/clients/bedrock/resources/bedrock_resource.py similarity index 100% rename from ai21/resources/bedrock_resource.py rename to ai21/clients/bedrock/resources/bedrock_resource.py diff --git a/ai21/resources/bases/__init__.py b/ai21/clients/common/__init__.py similarity index 100% rename from ai21/resources/bases/__init__.py rename to ai21/clients/common/__init__.py diff --git a/ai21/resources/bases/answer_base.py b/ai21/clients/common/answer_base.py similarity index 100% rename from ai21/resources/bases/answer_base.py rename to ai21/clients/common/answer_base.py diff --git a/ai21/resources/bases/chat_base.py b/ai21/clients/common/chat_base.py similarity index 100% rename from ai21/resources/bases/chat_base.py rename to ai21/clients/common/chat_base.py diff --git a/ai21/resources/bases/completion_base.py b/ai21/clients/common/completion_base.py similarity index 100% rename from ai21/resources/bases/completion_base.py rename to ai21/clients/common/completion_base.py diff --git a/ai21/resources/bases/custom_model_base.py b/ai21/clients/common/custom_model_base.py similarity index 100% rename from ai21/resources/bases/custom_model_base.py rename to ai21/clients/common/custom_model_base.py diff --git a/ai21/resources/bases/dataset_base.py b/ai21/clients/common/dataset_base.py similarity index 100% rename from ai21/resources/bases/dataset_base.py rename to ai21/clients/common/dataset_base.py diff --git a/ai21/resources/bases/embed_base.py b/ai21/clients/common/embed_base.py similarity index 100% rename from ai21/resources/bases/embed_base.py rename to ai21/clients/common/embed_base.py diff --git a/ai21/resources/bases/gec_base.py b/ai21/clients/common/gec_base.py similarity index 100% rename from ai21/resources/bases/gec_base.py rename to ai21/clients/common/gec_base.py diff --git a/ai21/resources/bases/improvements_base.py b/ai21/clients/common/improvements_base.py similarity index 100% rename from ai21/resources/bases/improvements_base.py rename to ai21/clients/common/improvements_base.py diff --git a/ai21/resources/bases/paraphrase_base.py b/ai21/clients/common/paraphrase_base.py similarity index 100% rename from ai21/resources/bases/paraphrase_base.py rename to ai21/clients/common/paraphrase_base.py diff --git a/ai21/resources/bases/segmentation_base.py b/ai21/clients/common/segmentation_base.py similarity index 100% rename from ai21/resources/bases/segmentation_base.py rename to ai21/clients/common/segmentation_base.py diff --git a/ai21/resources/bases/summarize_base.py b/ai21/clients/common/summarize_base.py similarity index 100% rename from ai21/resources/bases/summarize_base.py rename to ai21/clients/common/summarize_base.py diff --git a/ai21/resources/bases/summarize_by_segment_base.py b/ai21/clients/common/summarize_by_segment_base.py similarity index 100% rename from ai21/resources/bases/summarize_by_segment_base.py rename to ai21/clients/common/summarize_by_segment_base.py diff --git a/ai21/clients/sagemaker/resources/sagemaker_answer.py b/ai21/clients/sagemaker/resources/sagemaker_answer.py index 1ed83464..f344d6b0 100644 --- a/ai21/clients/sagemaker/resources/sagemaker_answer.py +++ b/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.clients.common.answer_base import Answer, AnswerLength, Mode from ai21.models.responses.answer_response import AnswerResponse -from ai21.resources.sagemaker_resource import SageMakerResource +from ai21.clients.sagemaker.resources.sagemaker_resource import SageMakerResource class SageMakerAnswer(SageMakerResource, Answer): diff --git a/ai21/clients/sagemaker/resources/sagemaker_completion.py b/ai21/clients/sagemaker/resources/sagemaker_completion.py index 55408607..373fdcb8 100644 --- a/ai21/clients/sagemaker/resources/sagemaker_completion.py +++ b/ai21/clients/sagemaker/resources/sagemaker_completion.py @@ -2,7 +2,7 @@ from ai21.resources import Penalty from ai21.models.responses.completion_response import CompletionsResponse -from ai21.resources.sagemaker_resource import SageMakerResource +from ai21.clients.sagemaker.resources.sagemaker_resource import SageMakerResource class SageMakerCompletion(SageMakerResource): diff --git a/ai21/clients/sagemaker/resources/sagemaker_gec.py b/ai21/clients/sagemaker/resources/sagemaker_gec.py index 5b593622..0750a7ea 100644 --- a/ai21/clients/sagemaker/resources/sagemaker_gec.py +++ b/ai21/clients/sagemaker/resources/sagemaker_gec.py @@ -1,6 +1,6 @@ -from ai21.resources.bases.gec_base import GEC +from ai21.clients.common.gec_base import GEC from ai21.models.responses.gec_response import GECResponse -from ai21.resources.sagemaker_resource import SageMakerResource +from ai21.clients.sagemaker.resources.sagemaker_resource import SageMakerResource class SageMakerGEC(SageMakerResource, GEC): diff --git a/ai21/clients/sagemaker/resources/sagemaker_paraphrase.py b/ai21/clients/sagemaker/resources/sagemaker_paraphrase.py index 54c31eb4..b2588019 100644 --- a/ai21/clients/sagemaker/resources/sagemaker_paraphrase.py +++ b/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.clients.common.paraphrase_base import Paraphrase from ai21.models.paraphrase_style_type import ParaphraseStyleType from ai21.models.responses.paraphrase_response import ParaphraseResponse -from ai21.resources.sagemaker_resource import SageMakerResource +from ai21.clients.sagemaker.resources.sagemaker_resource import SageMakerResource class SageMakerParaphrase(SageMakerResource, Paraphrase): diff --git a/ai21/resources/sagemaker_resource.py b/ai21/clients/sagemaker/resources/sagemaker_resource.py similarity index 100% rename from ai21/resources/sagemaker_resource.py rename to ai21/clients/sagemaker/resources/sagemaker_resource.py diff --git a/ai21/clients/sagemaker/resources/sagemaker_summarize.py b/ai21/clients/sagemaker/resources/sagemaker_summarize.py index 18491534..1d5a7bc2 100644 --- a/ai21/clients/sagemaker/resources/sagemaker_summarize.py +++ b/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.clients.common.summarize_base import Summarize from ai21.models.summary_method import SummaryMethod from ai21.models.responses import SummarizeResponse -from ai21.resources.sagemaker_resource import SageMakerResource +from ai21.clients.sagemaker.resources.sagemaker_resource import SageMakerResource class SageMakerSummarize(SageMakerResource, Summarize): diff --git a/ai21/clients/studio/resources/studio_answer.py b/ai21/clients/studio/resources/studio_answer.py index cbc2ca17..b2ceb03d 100644 --- a/ai21/clients/studio/resources/studio_answer.py +++ b/ai21/clients/studio/resources/studio_answer.py @@ -1,8 +1,8 @@ from typing import Optional -from ai21.resources.bases.answer_base import Answer, AnswerLength, Mode +from ai21.clients.common.answer_base import Answer, AnswerLength, Mode from ai21.models.responses.answer_response import AnswerResponse -from ai21.resources.studio_resource import StudioResource +from ai21.clients.studio.resources.studio_resource import StudioResource class StudioAnswer(StudioResource, Answer): diff --git a/ai21/clients/studio/resources/studio_chat.py b/ai21/clients/studio/resources/studio_chat.py index fb757ae0..435b9f21 100644 --- a/ai21/clients/studio/resources/studio_chat.py +++ b/ai21/clients/studio/resources/studio_chat.py @@ -1,8 +1,8 @@ from typing import List, Any, Optional, Dict -from ai21.resources.bases.chat_base import Chat, Message +from ai21.clients.common.chat_base import Chat, Message from ai21.models.responses import ChatResponse -from ai21.resources.studio_resource import StudioResource +from ai21.clients.studio.resources.studio_resource import StudioResource class StudioChat(StudioResource, Chat): diff --git a/ai21/clients/studio/resources/studio_completion.py b/ai21/clients/studio/resources/studio_completion.py index 1d1318d8..9f9a0ceb 100644 --- a/ai21/clients/studio/resources/studio_completion.py +++ b/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.clients.common.completion_base import Completion from ai21.models.responses.completion_response import CompletionsResponse -from ai21.resources.studio_resource import StudioResource +from ai21.clients.studio.resources.studio_resource import StudioResource class StudioCompletion(StudioResource, Completion): diff --git a/ai21/clients/studio/resources/studio_custom_model.py b/ai21/clients/studio/resources/studio_custom_model.py index 82718138..831b4575 100644 --- a/ai21/clients/studio/resources/studio_custom_model.py +++ b/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.clients.common.custom_model_base import CustomModel from ai21.models.responses.custom_model_response import CustomBaseModelResponse -from ai21.resources.studio_resource import StudioResource +from ai21.clients.studio.resources.studio_resource import StudioResource class StudioCustomModel(StudioResource, CustomModel): diff --git a/ai21/clients/studio/resources/studio_dataset.py b/ai21/clients/studio/resources/studio_dataset.py index 351391ab..ccfb4bac 100644 --- a/ai21/clients/studio/resources/studio_dataset.py +++ b/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.clients.common.dataset_base import Dataset from ai21.models.responses.dataset_response import DatasetResponse -from ai21.resources.studio_resource import StudioResource +from ai21.clients.studio.resources.studio_resource import StudioResource class StudioDataset(StudioResource, Dataset): diff --git a/ai21/clients/studio/resources/studio_embed.py b/ai21/clients/studio/resources/studio_embed.py index e2cde9d9..48e22de7 100644 --- a/ai21/clients/studio/resources/studio_embed.py +++ b/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.clients.common.embed_base import Embed from ai21.models.responses import EmbedResponse -from ai21.resources.studio_resource import StudioResource +from ai21.clients.studio.resources.studio_resource import StudioResource class StudioEmbed(StudioResource, Embed): diff --git a/ai21/clients/studio/resources/studio_gec.py b/ai21/clients/studio/resources/studio_gec.py index 1ee5b8a3..3e716e7a 100644 --- a/ai21/clients/studio/resources/studio_gec.py +++ b/ai21/clients/studio/resources/studio_gec.py @@ -1,6 +1,6 @@ -from ai21.resources.bases.gec_base import GEC +from ai21.clients.common.gec_base import GEC from ai21.models.responses.gec_response import GECResponse -from ai21.resources.studio_resource import StudioResource +from ai21.clients.studio.resources.studio_resource import StudioResource class StudioGEC(StudioResource, GEC): diff --git a/ai21/clients/studio/resources/studio_improvements.py b/ai21/clients/studio/resources/studio_improvements.py index 1b3f81c0..a305c190 100644 --- a/ai21/clients/studio/resources/studio_improvements.py +++ b/ai21/clients/studio/resources/studio_improvements.py @@ -1,10 +1,10 @@ from typing import List from ai21.errors import EmptyMandatoryListError -from ai21.resources.bases.improvements_base import Improvements +from ai21.clients.common.improvements_base import Improvements from ai21.models.improvement_type import ImprovementType from ai21.models.responses.improvement_response import ImprovementsResponse -from ai21.resources.studio_resource import StudioResource +from ai21.clients.studio.resources.studio_resource import StudioResource class StudioImprovements(StudioResource, Improvements): diff --git a/ai21/clients/studio/resources/studio_library.py b/ai21/clients/studio/resources/studio_library.py index dece0cfe..f51514e4 100644 --- a/ai21/clients/studio/resources/studio_library.py +++ b/ai21/clients/studio/resources/studio_library.py @@ -6,7 +6,7 @@ from ai21.models.responses.file_response import FileResponse from ai21.models.responses.library_answer_response import LibraryAnswerResponse from ai21.models.responses.library_search_response import LibrarySearchResponse -from ai21.resources.studio_resource import StudioResource +from ai21.clients.studio.resources.studio_resource import StudioResource class StudioLibrary(StudioResource): diff --git a/ai21/clients/studio/resources/studio_paraphrase.py b/ai21/clients/studio/resources/studio_paraphrase.py index 9202a164..8a7dc1cc 100644 --- a/ai21/clients/studio/resources/studio_paraphrase.py +++ b/ai21/clients/studio/resources/studio_paraphrase.py @@ -1,9 +1,9 @@ from typing import Optional -from ai21.resources.bases.paraphrase_base import Paraphrase +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.resources.studio_resource import StudioResource +from ai21.clients.studio.resources.studio_resource import StudioResource class StudioParaphrase(StudioResource, Paraphrase): diff --git a/ai21/resources/studio_resource.py b/ai21/clients/studio/resources/studio_resource.py similarity index 100% rename from ai21/resources/studio_resource.py rename to ai21/clients/studio/resources/studio_resource.py diff --git a/ai21/clients/studio/resources/studio_segmentation.py b/ai21/clients/studio/resources/studio_segmentation.py index 8e49f5b5..b2c9fdf0 100644 --- a/ai21/clients/studio/resources/studio_segmentation.py +++ b/ai21/clients/studio/resources/studio_segmentation.py @@ -1,7 +1,7 @@ -from ai21.resources.bases.segmentation_base import Segmentation +from ai21.clients.common.segmentation_base import Segmentation from ai21.models.document_type import DocumentType from ai21.models.responses import SegmentationResponse -from ai21.resources.studio_resource import StudioResource +from ai21.clients.studio.resources.studio_resource import StudioResource class StudioSegmentation(StudioResource, Segmentation): diff --git a/ai21/clients/studio/resources/studio_summarize.py b/ai21/clients/studio/resources/studio_summarize.py index ce44ad5d..a2882f2b 100644 --- a/ai21/clients/studio/resources/studio_summarize.py +++ b/ai21/clients/studio/resources/studio_summarize.py @@ -1,9 +1,9 @@ from typing import Optional -from ai21.resources.bases.summarize_base import Summarize +from ai21.clients.common.summarize_base import Summarize from ai21.models.summary_method import SummaryMethod from ai21.models.responses import SummarizeResponse -from ai21.resources.studio_resource import StudioResource +from ai21.clients.studio.resources.studio_resource import StudioResource class StudioSummarize(StudioResource, Summarize): diff --git a/ai21/clients/studio/resources/studio_summarize_by_segment.py b/ai21/clients/studio/resources/studio_summarize_by_segment.py index 1dec9336..4066588e 100644 --- a/ai21/clients/studio/resources/studio_summarize_by_segment.py +++ b/ai21/clients/studio/resources/studio_summarize_by_segment.py @@ -1,11 +1,11 @@ from typing import Optional -from ai21.resources.bases.summarize_by_segment_base import SummarizeBySegment +from ai21.clients.common.summarize_by_segment_base import SummarizeBySegment from ai21.models.document_type import DocumentType from ai21.models.responses import ( SummarizeBySegmentResponse, ) -from ai21.resources.studio_resource import StudioResource +from ai21.clients.studio.resources.studio_resource import StudioResource class StudioSummarizeBySegment(StudioResource, SummarizeBySegment): diff --git a/tests/unittests/clients/studio/resources/test_studio_resources.py b/tests/unittests/clients/studio/resources/test_studio_resources.py index 0e4de3af..feb1f14e 100644 --- a/tests/unittests/clients/studio/resources/test_studio_resources.py +++ b/tests/unittests/clients/studio/resources/test_studio_resources.py @@ -5,7 +5,7 @@ from ai21 import AnswerResponse from ai21.ai21_http_client import AI21HTTPClient from ai21.clients.studio.resources.studio_answer import StudioAnswer -from ai21.resources.studio_resource import StudioResource +from ai21.clients.studio.resources.studio_resource import StudioResource from tests.unittests.clients.studio.resources.conftest import get_studio_answer, get_studio_chat, get_studio_completion _BASE_URL = "https://test.api.ai21.com/studio/v1" From 5c92441f0c5def58a7c6a42726e65c06eaebfa15 Mon Sep 17 00:00:00 2001 From: Asaf Gardin Date: Tue, 2 Jan 2024 15:07:30 +0200 Subject: [PATCH 4/7] refactor: chat message rename --- ai21/clients/common/chat_base.py | 6 +++--- ai21/clients/studio/resources/studio_chat.py | 4 ++-- ai21/models/__init__.py | 5 ++--- ai21/models/{message.py => chat_message.py} | 2 +- ai21/resources/__init__.py | 0 5 files changed, 8 insertions(+), 9 deletions(-) rename ai21/models/{message.py => chat_message.py} (85%) delete mode 100644 ai21/resources/__init__.py diff --git a/ai21/clients/common/chat_base.py b/ai21/clients/common/chat_base.py index bf898942..f970eef9 100644 --- a/ai21/clients/common/chat_base.py +++ b/ai21/clients/common/chat_base.py @@ -1,7 +1,7 @@ from abc import ABC, abstractmethod from typing import List, Any, Dict, Optional -from ai21.models import Message, Penalty +from ai21.models import ChatMessage, Penalty from ai21.models.responses import ChatResponse @@ -12,7 +12,7 @@ class Chat(ABC): def create( self, model: str, - messages: List[Message], + messages: List[ChatMessage], system: str, *, num_results: Optional[int] = 1, @@ -35,7 +35,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, diff --git a/ai21/clients/studio/resources/studio_chat.py b/ai21/clients/studio/resources/studio_chat.py index 435b9f21..9204023a 100644 --- a/ai21/clients/studio/resources/studio_chat.py +++ b/ai21/clients/studio/resources/studio_chat.py @@ -1,6 +1,6 @@ from typing import List, Any, Optional, Dict -from ai21.clients.common.chat_base import Chat, Message +from ai21.clients.common.chat_base import Chat, ChatMessage from ai21.models.responses import ChatResponse from ai21.clients.studio.resources.studio_resource import StudioResource @@ -9,7 +9,7 @@ class StudioChat(StudioResource, Chat): def create( self, model: str, - messages: List[Message], + messages: List[ChatMessage], system: str, *, num_results: Optional[int] = 1, diff --git a/ai21/models/__init__.py b/ai21/models/__init__.py index 393d5c28..ba0ff318 100644 --- a/ai21/models/__init__.py +++ b/ai21/models/__init__.py @@ -1,19 +1,18 @@ from ai21.models.answer_length import AnswerLength +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.message import Message from ai21.models.mode import Mode from ai21.models.paraphrase_style_type import ParaphraseStyleType from ai21.models.penalty import Penalty from ai21.models.role_type import RoleType from ai21.models.summary_method import SummaryMethod - __all__ = [ "AnswerLength", "Mode", - "Message", + "ChatMessage", "RoleType", "Penalty", "EmbedType", diff --git a/ai21/models/message.py b/ai21/models/chat_message.py similarity index 85% rename from ai21/models/message.py rename to ai21/models/chat_message.py index 12937a3b..ec4cdc95 100644 --- a/ai21/models/message.py +++ b/ai21/models/chat_message.py @@ -6,7 +6,7 @@ @dataclass -class Message(AI21BaseModelMixin): +class ChatMessage(AI21BaseModelMixin): role: RoleType text: str name: Optional[str] = None diff --git a/ai21/resources/__init__.py b/ai21/resources/__init__.py deleted file mode 100644 index e69de29b..00000000 From b60e5fc1471b2403596f62cc797091600bfcac4b Mon Sep 17 00:00:00 2001 From: Asaf Gardin Date: Tue, 2 Jan 2024 15:50:12 +0200 Subject: [PATCH 5/7] refactor: init --- ai21/__init__.py | 15 --------- ai21/clients/studio/resources/studio_chat.py | 2 +- .../studio/resources/studio_completion.py | 4 +-- ai21/clients/studio/resources/studio_embed.py | 2 +- .../studio/resources/studio_improvements.py | 6 ++-- .../studio/resources/studio_library.py | 3 +- .../studio/resources/studio_paraphrase.py | 4 +-- .../studio/resources/studio_segmentation.py | 4 +-- .../studio/resources/studio_summarize.py | 4 +-- .../resources/studio_summarize_by_segment.py | 2 +- ai21/models/__init__.py | 31 +++++++++++++++++++ 11 files changed, 46 insertions(+), 31 deletions(-) diff --git a/ai21/__init__.py b/ai21/__init__.py index 5fa7b336..bfc95fa8 100644 --- a/ai21/__init__.py +++ b/ai21/__init__.py @@ -60,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", ] diff --git a/ai21/clients/studio/resources/studio_chat.py b/ai21/clients/studio/resources/studio_chat.py index 9204023a..0d4fd502 100644 --- a/ai21/clients/studio/resources/studio_chat.py +++ b/ai21/clients/studio/resources/studio_chat.py @@ -1,8 +1,8 @@ from typing import List, Any, Optional, Dict from ai21.clients.common.chat_base import Chat, ChatMessage -from ai21.models.responses import ChatResponse from ai21.clients.studio.resources.studio_resource import StudioResource +from ai21.models.responses.chat_response import ChatResponse class StudioChat(StudioResource, Chat): diff --git a/ai21/clients/studio/resources/studio_completion.py b/ai21/clients/studio/resources/studio_completion.py index 9f9a0ceb..75364113 100644 --- a/ai21/clients/studio/resources/studio_completion.py +++ b/ai21/clients/studio/resources/studio_completion.py @@ -1,9 +1,9 @@ from typing import Optional, List -from ai21.resources import Penalty from ai21.clients.common.completion_base import Completion -from ai21.models.responses.completion_response import CompletionsResponse 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): diff --git a/ai21/clients/studio/resources/studio_embed.py b/ai21/clients/studio/resources/studio_embed.py index 48e22de7..7495ea67 100644 --- a/ai21/clients/studio/resources/studio_embed.py +++ b/ai21/clients/studio/resources/studio_embed.py @@ -1,8 +1,8 @@ from typing import List, Optional from ai21.clients.common.embed_base import Embed -from ai21.models.responses import EmbedResponse from ai21.clients.studio.resources.studio_resource import StudioResource +from ai21.models.responses.embed_response import EmbedResponse class StudioEmbed(StudioResource, Embed): diff --git a/ai21/clients/studio/resources/studio_improvements.py b/ai21/clients/studio/resources/studio_improvements.py index a305c190..513a413c 100644 --- a/ai21/clients/studio/resources/studio_improvements.py +++ b/ai21/clients/studio/resources/studio_improvements.py @@ -1,10 +1,10 @@ from typing import List -from ai21.errors import EmptyMandatoryListError from ai21.clients.common.improvements_base import Improvements -from ai21.models.improvement_type import ImprovementType -from ai21.models.responses.improvement_response import ImprovementsResponse from ai21.clients.studio.resources.studio_resource import StudioResource +from ai21.errors import EmptyMandatoryListError +from ai21.models import ImprovementType +from ai21.models.responses.improvement_response import ImprovementsResponse class StudioImprovements(StudioResource, Improvements): diff --git a/ai21/clients/studio/resources/studio_library.py b/ai21/clients/studio/resources/studio_library.py index f51514e4..fa73b123 100644 --- a/ai21/clients/studio/resources/studio_library.py +++ b/ai21/clients/studio/resources/studio_library.py @@ -1,8 +1,7 @@ from typing import Optional, List from ai21.ai21_http_client import AI21HTTPClient -from ai21.models.answer_length import AnswerLength -from ai21.models.mode import Mode +from ai21.models import Mode, AnswerLength from ai21.models.responses.file_response import FileResponse from ai21.models.responses.library_answer_response import LibraryAnswerResponse from ai21.models.responses.library_search_response import LibrarySearchResponse diff --git a/ai21/clients/studio/resources/studio_paraphrase.py b/ai21/clients/studio/resources/studio_paraphrase.py index 8a7dc1cc..686841db 100644 --- a/ai21/clients/studio/resources/studio_paraphrase.py +++ b/ai21/clients/studio/resources/studio_paraphrase.py @@ -1,9 +1,9 @@ from typing import Optional 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.studio.resources.studio_resource import StudioResource +from ai21.models import ParaphraseStyleType +from ai21.models.responses.paraphrase_response import ParaphraseResponse class StudioParaphrase(StudioResource, Paraphrase): diff --git a/ai21/clients/studio/resources/studio_segmentation.py b/ai21/clients/studio/resources/studio_segmentation.py index b2c9fdf0..dbda4225 100644 --- a/ai21/clients/studio/resources/studio_segmentation.py +++ b/ai21/clients/studio/resources/studio_segmentation.py @@ -1,7 +1,7 @@ from ai21.clients.common.segmentation_base import Segmentation -from ai21.models.document_type import DocumentType -from ai21.models.responses import SegmentationResponse from ai21.clients.studio.resources.studio_resource import StudioResource +from ai21.models.document_type import DocumentType +from ai21.models.responses.segmentation_response import SegmentationResponse class StudioSegmentation(StudioResource, Segmentation): diff --git a/ai21/clients/studio/resources/studio_summarize.py b/ai21/clients/studio/resources/studio_summarize.py index a2882f2b..b2b5f860 100644 --- a/ai21/clients/studio/resources/studio_summarize.py +++ b/ai21/clients/studio/resources/studio_summarize.py @@ -1,9 +1,9 @@ from typing import Optional from ai21.clients.common.summarize_base import Summarize -from ai21.models.summary_method import SummaryMethod -from ai21.models.responses import SummarizeResponse from ai21.clients.studio.resources.studio_resource import StudioResource +from ai21.models.responses.summarize_response import SummarizeResponse +from ai21.models.summary_method import SummaryMethod class StudioSummarize(StudioResource, Summarize): diff --git a/ai21/clients/studio/resources/studio_summarize_by_segment.py b/ai21/clients/studio/resources/studio_summarize_by_segment.py index 4066588e..8292a1f9 100644 --- a/ai21/clients/studio/resources/studio_summarize_by_segment.py +++ b/ai21/clients/studio/resources/studio_summarize_by_segment.py @@ -2,7 +2,7 @@ from ai21.clients.common.summarize_by_segment_base import SummarizeBySegment from ai21.models.document_type import DocumentType -from ai21.models.responses import ( +from ai21.models.responses.summarize_by_segment_response import ( SummarizeBySegmentResponse, ) from ai21.clients.studio.resources.studio_resource import StudioResource diff --git a/ai21/models/__init__.py b/ai21/models/__init__.py index ba0ff318..6c14c3ea 100644 --- a/ai21/models/__init__.py +++ b/ai21/models/__init__.py @@ -6,9 +6,25 @@ from ai21.models.mode import Mode 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 +from ai21.models.responses.completion_response import CompletionsResponse +from ai21.models.responses.custom_model_response import CustomBaseModelResponse +from ai21.models.responses.dataset_response import DatasetResponse +from ai21.models.responses.embed_response import EmbedResponse +from ai21.models.responses.file_response import FileResponse +from ai21.models.responses.gec_response import GECResponse +from ai21.models.responses.improvement_response import ImprovementsResponse +from ai21.models.responses.library_answer_response import LibraryAnswerResponse +from ai21.models.responses.library_search_response import LibrarySearchResponse +from ai21.models.responses.paraphrase_response import ParaphraseResponse +from ai21.models.responses.segmentation_response import SegmentationResponse +from ai21.models.responses.summarize_by_segment_response import SummarizeBySegmentResponse +from ai21.models.responses.summarize_response import SummarizeResponse from ai21.models.role_type import RoleType from ai21.models.summary_method import SummaryMethod + __all__ = [ "AnswerLength", "Mode", @@ -20,4 +36,19 @@ "ParaphraseStyleType", "DocumentType", "SummaryMethod", + "AnswerResponse", + "ChatResponse", + "CompletionsResponse", + "CustomBaseModelResponse", + "DatasetResponse", + "EmbedResponse", + "FileResponse", + "GECResponse", + "ImprovementsResponse", + "LibraryAnswerResponse", + "LibrarySearchResponse", + "ParaphraseResponse", + "SegmentationResponse", + "SummarizeBySegmentResponse", + "SummarizeResponse", ] From c047f4f6e142b95fa2ce62d0cd640e7fc3b84b74 Mon Sep 17 00:00:00 2001 From: Asaf Gardin Date: Tue, 2 Jan 2024 16:00:56 +0200 Subject: [PATCH 6/7] fix: imports --- ai21/ai21_http_client.py | 2 +- ai21/clients/common/chat_base.py | 5 +++-- ai21/clients/common/embed_base.py | 4 ++-- ai21/clients/common/segmentation_base.py | 4 ++-- ai21/clients/common/summarize_base.py | 4 ++-- .../clients/common/summarize_by_segment_base.py | 4 ++-- ai21/clients/studio/resources/studio_answer.py | 3 ++- ai21/clients/studio/resources/studio_chat.py | 3 ++- ai21/models/chat_message.py | 2 +- ai21/models/responses/answer_response.py | 4 ++-- .../clients/studio/resources/conftest.py | 17 +++++++++++------ .../studio/resources/test_studio_resources.py | 2 +- 12 files changed, 31 insertions(+), 23 deletions(-) diff --git a/ai21/ai21_http_client.py b/ai21/ai21_http_client.py index 68007654..995f1f19 100644 --- a/ai21/ai21_http_client.py +++ b/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 diff --git a/ai21/clients/common/chat_base.py b/ai21/clients/common/chat_base.py index f970eef9..869ddfaa 100644 --- a/ai21/clients/common/chat_base.py +++ b/ai21/clients/common/chat_base.py @@ -1,8 +1,9 @@ from abc import ABC, abstractmethod from typing import List, Any, Dict, Optional -from ai21.models import ChatMessage, Penalty -from ai21.models.responses import ChatResponse +from ai21.models.chat_message import ChatMessage +from ai21.models.penalty import Penalty +from ai21.models.responses.chat_response import ChatResponse class Chat(ABC): diff --git a/ai21/clients/common/embed_base.py b/ai21/clients/common/embed_base.py index 674ebfe0..baadd4ec 100644 --- a/ai21/clients/common/embed_base.py +++ b/ai21/clients/common/embed_base.py @@ -1,8 +1,8 @@ from abc import ABC, abstractmethod from typing import List, Any, Dict, Optional -from ai21.models import EmbedType -from ai21.models.responses import EmbedResponse +from ai21.models.embed_type import EmbedType +from ai21.models.responses.embed_response import EmbedResponse class Embed(ABC): diff --git a/ai21/clients/common/segmentation_base.py b/ai21/clients/common/segmentation_base.py index dfbb3e70..c4f658a9 100644 --- a/ai21/clients/common/segmentation_base.py +++ b/ai21/clients/common/segmentation_base.py @@ -1,8 +1,8 @@ from abc import ABC, abstractmethod from typing import Any, Dict -from ai21.models import DocumentType -from ai21.models.responses import SegmentationResponse +from ai21.models.document_type import DocumentType +from ai21.models.responses.segmentation_response import SegmentationResponse class Segmentation(ABC): diff --git a/ai21/clients/common/summarize_base.py b/ai21/clients/common/summarize_base.py index ba2c5683..85cec2f5 100644 --- a/ai21/clients/common/summarize_base.py +++ b/ai21/clients/common/summarize_base.py @@ -1,8 +1,8 @@ from abc import ABC, abstractmethod from typing import Optional, Any, Dict -from ai21.models import SummaryMethod -from ai21.models.responses import SummarizeResponse +from ai21.models.responses.summarize_response import SummarizeResponse +from ai21.models.summary_method import SummaryMethod class Summarize(ABC): diff --git a/ai21/clients/common/summarize_by_segment_base.py b/ai21/clients/common/summarize_by_segment_base.py index fa16c251..40a4abfa 100644 --- a/ai21/clients/common/summarize_by_segment_base.py +++ b/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.models import DocumentType -from ai21.models.responses import ( +from ai21.models.document_type import DocumentType +from ai21.models.responses.summarize_by_segment_response import ( SummarizeBySegmentResponse, ) diff --git a/ai21/clients/studio/resources/studio_answer.py b/ai21/clients/studio/resources/studio_answer.py index b2ceb03d..a8cfc13e 100644 --- a/ai21/clients/studio/resources/studio_answer.py +++ b/ai21/clients/studio/resources/studio_answer.py @@ -1,6 +1,7 @@ from typing import Optional -from ai21.clients.common.answer_base import Answer, AnswerLength, Mode +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 diff --git a/ai21/clients/studio/resources/studio_chat.py b/ai21/clients/studio/resources/studio_chat.py index 0d4fd502..710ed308 100644 --- a/ai21/clients/studio/resources/studio_chat.py +++ b/ai21/clients/studio/resources/studio_chat.py @@ -1,7 +1,8 @@ from typing import List, Any, Optional, Dict -from ai21.clients.common.chat_base import Chat, ChatMessage +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 diff --git a/ai21/models/chat_message.py b/ai21/models/chat_message.py index ec4cdc95..ebf719e3 100644 --- a/ai21/models/chat_message.py +++ b/ai21/models/chat_message.py @@ -1,7 +1,7 @@ from dataclasses import dataclass from typing import Optional -from ai21.models import RoleType +from ai21.models.role_type import RoleType from ai21.models.ai21_base_model_mixin import AI21BaseModelMixin diff --git a/ai21/models/responses/answer_response.py b/ai21/models/responses/answer_response.py index 3056b6f2..437881b4 100644 --- a/ai21/models/responses/answer_response.py +++ b/ai21/models/responses/answer_response.py @@ -7,5 +7,5 @@ @dataclass class AnswerResponse(AI21BaseModelMixin): id: str - answer_in_context: Optional[bool] - answer: Optional[str] + answer_in_context: Optional[bool] = None + answer: Optional[str] = None diff --git a/tests/unittests/clients/studio/resources/conftest.py b/tests/unittests/clients/studio/resources/conftest.py index 1009a3dc..a28e559d 100644 --- a/tests/unittests/clients/studio/resources/conftest.py +++ b/tests/unittests/clients/studio/resources/conftest.py @@ -1,14 +1,19 @@ import pytest from pytest_mock import MockerFixture -from ai21 import AnswerResponse, ChatResponse, CompletionsResponse from ai21.ai21_http_client import AI21HTTPClient from ai21.clients.studio.resources.studio_answer import StudioAnswer from ai21.clients.studio.resources.studio_chat import StudioChat from ai21.clients.studio.resources.studio_completion import StudioCompletion -from ai21.resources import Message, RoleType -from ai21.models.responses import ChatOutput, FinishReason -from ai21.models.responses.completion_response import Prompt, Completion, CompletionData, CompletionFinishReason +from ai21.models import AnswerResponse, ChatMessage, RoleType, ChatResponse +from ai21.models.responses.chat_response import ChatOutput, FinishReason +from ai21.models.responses.completion_response import ( + Prompt, + Completion, + CompletionData, + CompletionFinishReason, + CompletionsResponse, +) @pytest.fixture @@ -37,8 +42,8 @@ def get_studio_answer(): def get_studio_chat(): _DUMMY_MODEL = "dummy-chat-model" _DUMMY_MESSAGES = [ - Message(text="Hello, I need help with a signup process.", role=RoleType.USER, name="Alice"), - Message( + ChatMessage(text="Hello, I need help with a signup process.", role=RoleType.USER, name="Alice"), + ChatMessage( text="Hi Alice, I can help you with that. What seems to be the problem?", role=RoleType.ASSISTANT, name="Bob", diff --git a/tests/unittests/clients/studio/resources/test_studio_resources.py b/tests/unittests/clients/studio/resources/test_studio_resources.py index feb1f14e..cc534fe2 100644 --- a/tests/unittests/clients/studio/resources/test_studio_resources.py +++ b/tests/unittests/clients/studio/resources/test_studio_resources.py @@ -2,10 +2,10 @@ import pytest -from ai21 import AnswerResponse from ai21.ai21_http_client import AI21HTTPClient from ai21.clients.studio.resources.studio_answer import StudioAnswer from ai21.clients.studio.resources.studio_resource import StudioResource +from ai21.models.responses.answer_response import AnswerResponse from tests.unittests.clients.studio.resources.conftest import get_studio_answer, get_studio_chat, get_studio_completion _BASE_URL = "https://test.api.ai21.com/studio/v1" From 1f367fdf27acb020c3979a34f56acf5efccb86bc Mon Sep 17 00:00:00 2001 From: Asaf Gardin Date: Tue, 2 Jan 2024 16:12:21 +0200 Subject: [PATCH 7/7] refactor: added more to imports --- ai21/models/__init__.py | 42 +++++++++++++++++++++++++++++++---------- 1 file changed, 32 insertions(+), 10 deletions(-) diff --git a/ai21/models/__init__.py b/ai21/models/__init__.py index 6c14c3ea..5e5877d8 100644 --- a/ai21/models/__init__.py +++ b/ai21/models/__init__.py @@ -7,19 +7,25 @@ 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 -from ai21.models.responses.completion_response import CompletionsResponse -from ai21.models.responses.custom_model_response import CustomBaseModelResponse +from ai21.models.responses.chat_response import ChatResponse, ChatOutput, FinishReason +from ai21.models.responses.completion_response import ( + CompletionsResponse, + Completion, + CompletionFinishReason, + CompletionData, + Prompt, +) +from ai21.models.responses.custom_model_response import CustomBaseModelResponse, BaseModelMetadata from ai21.models.responses.dataset_response import DatasetResponse -from ai21.models.responses.embed_response import EmbedResponse +from ai21.models.responses.embed_response import EmbedResponse, EmbedResult from ai21.models.responses.file_response import FileResponse -from ai21.models.responses.gec_response import GECResponse -from ai21.models.responses.improvement_response import ImprovementsResponse -from ai21.models.responses.library_answer_response import LibraryAnswerResponse -from ai21.models.responses.library_search_response import LibrarySearchResponse -from ai21.models.responses.paraphrase_response import ParaphraseResponse +from ai21.models.responses.gec_response import GECResponse, Correction, CorrectionType +from ai21.models.responses.improvement_response import ImprovementsResponse, Improvement +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 +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 @@ -38,17 +44,33 @@ "SummaryMethod", "AnswerResponse", "ChatResponse", + "ChatOutput", + "FinishReason", "CompletionsResponse", + "Completion", + "CompletionFinishReason", + "CompletionData", + "Prompt", "CustomBaseModelResponse", + "BaseModelMetadata", "DatasetResponse", "EmbedResponse", + "EmbedResult", "FileResponse", "GECResponse", + "Correction", + "CorrectionType", "ImprovementsResponse", + "Improvement", "LibraryAnswerResponse", + "SourceDocument", "LibrarySearchResponse", + "LibrarySearchResult", "ParaphraseResponse", + "Suggestion", "SegmentationResponse", + "SegmentSummary", + "Highlight", "SummarizeBySegmentResponse", "SummarizeResponse", ]