Merge branch 'master' into Add-docstrings-for-_runs

Author: andy-fisher

src/yandex_cloud_ml_sdk/_assistants/assistant.py

@@ -32,7 +32,7 @@
 
 
 @dataclasses.dataclass(frozen=True)
-class BaseAssistant(ExpirableResource, Generic[RunTypeT, ThreadTypeT]):
+class BaseAssistant(ExpirableResource[ProtoAssistant], Generic[RunTypeT, ThreadTypeT]):
     expiration_config: ExpirationConfig
     model: BaseGPTModel
     instruction: str | None

src/yandex_cloud_ml_sdk/_datasets/dataset.py

@@ -24,6 +24,7 @@
 
 from yandex_cloud_ml_sdk._logging import get_logger
 from yandex_cloud_ml_sdk._types.misc import UNDEFINED, PathLike, UndefinedOr, coerce_path, get_defined_value
+from yandex_cloud_ml_sdk._types.proto import ProtoBased
 from yandex_cloud_ml_sdk._types.resource import BaseDeleteableResource, safe_on_delete
 from yandex_cloud_ml_sdk._utils.packages import requires_package
 from yandex_cloud_ml_sdk._utils.pyarrow import read_dataset_records
@@ -40,13 +41,14 @@
 DEFAULT_MAX_PARALLEL_DOWNLOADS: Final[int] = 16 # maximum number of files open for writing during download
 
 @dataclasses.dataclass(frozen=True)
-class ValidationErrorInfo:
+class ValidationErrorInfo(ProtoBased[ProtoValidationError]):
     error: str
     description: str
     rows: tuple[int, ...]
 
+    # pylint: disable=unused-argument
     @classmethod
-    def _from_proto(cls, proto: ProtoValidationError) -> ValidationErrorInfo:
+    def _from_proto(cls, *, proto: ProtoValidationError, sdk: BaseSDK) -> ValidationErrorInfo:
         return cls(
             error=proto.error,
             description=proto.error_description,
@@ -74,15 +76,15 @@ class DatasetInfo:
 
 
 @dataclasses.dataclass(frozen=True)
-class BaseDataset(DatasetInfo, BaseDeleteableResource):
+class BaseDataset(DatasetInfo, BaseDeleteableResource[ProtoDatasetInfo]):
     @classmethod
-    def _kwargs_from_message(cls, proto: ProtoDatasetInfo, sdk: BaseSDK) -> dict[str, Any]:  # type: ignore[override]
+    def _kwargs_from_message(cls, proto: ProtoDatasetInfo, sdk: BaseSDK) -> dict[str, Any]:
         kwargs = super()._kwargs_from_message(proto, sdk=sdk)
         kwargs['id'] = proto.dataset_id
         kwargs['created_by'] = proto.created_by_id
         kwargs['status'] = DatasetStatus._from_proto(proto.status)
         kwargs['validation_errors'] = tuple(
-            ValidationErrorInfo._from_proto(p) for p in proto.validation_error
+            ValidationErrorInfo._from_proto(proto=p, sdk=sdk) for p in proto.validation_error
         )
         kwargs['allow_data_logging'] = proto.allow_data_log
         return kwargs

src/yandex_cloud_ml_sdk/_datasets/validation.py

@@ -2,28 +2,28 @@
 from __future__ import annotations
 
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING
 
 from yandex.cloud.ai.dataset.v1.dataset_pb2 import ValidationError as ProtoValidationError
 from yandex.cloud.ai.dataset.v1.dataset_service_pb2 import ValidateDatasetResponse
 
-from yandex_cloud_ml_sdk._types.result import BaseResult, ProtoMessage
+from yandex_cloud_ml_sdk._types.proto import ProtoBased
+from yandex_cloud_ml_sdk._types.result import BaseProtoResult
 from yandex_cloud_ml_sdk.exceptions import DatasetValidationError
 
 if TYPE_CHECKING:
     from yandex_cloud_ml_sdk._sdk import BaseSDK
 
 
 @dataclass(frozen=True)
-class ValidationErrorInfo:
+class ValidationErrorInfo(ProtoBased[ProtoValidationError]):
     error: str
     description: str
     rows: tuple[int, ...]
 
     # pylint: disable=unused-argument
     @classmethod
-    def _from_proto(cls, *, proto: ProtoMessage, sdk: BaseSDK) -> ValidationErrorInfo:
-        proto = cast(ProtoValidationError, proto)
+    def _from_proto(cls, *, proto: ProtoValidationError, sdk: BaseSDK) -> ValidationErrorInfo:
         return cls(
             error=proto.error,
             description=proto.error_description,
@@ -32,16 +32,14 @@ def _from_proto(cls, *, proto: ProtoMessage, sdk: BaseSDK) -> ValidationErrorInf
 
 
 @dataclass(frozen=True)
-class DatasetValidationResult(BaseResult):
+class DatasetValidationResult(BaseProtoResult[ValidateDatasetResponse]):
     _sdk: BaseSDK = field(repr=False)
     dataset_id: str
     is_valid: bool
     errors: tuple[ValidationErrorInfo, ...]
 
     @classmethod
-    def _from_proto(cls, *, proto: ProtoMessage, sdk: BaseSDK) -> DatasetValidationResult:
-        proto = cast(ValidateDatasetResponse, proto)
-
+    def _from_proto(cls, *, proto: ValidateDatasetResponse, sdk: BaseSDK) -> DatasetValidationResult:
         return cls(
             dataset_id=proto.dataset_id,
             is_valid=proto.is_valid,

src/yandex_cloud_ml_sdk/_files/file.py

@@ -20,7 +20,7 @@
 
 
 @dataclasses.dataclass(frozen=True)
-class BaseFile(ExpirableResource):
+class BaseFile(ExpirableResource[ProtoFile]):
     @safe_on_delete
     async def _get_url(
         self,

src/yandex_cloud_ml_sdk/_messages/base.py

@@ -3,11 +3,11 @@
 import dataclasses
 from typing import Any
 
-from yandex_cloud_ml_sdk._types.result import BaseResult
+from yandex_cloud_ml_sdk._types.result import BaseProtoResult, ProtoMessageTypeT_contra
 
 
 @dataclasses.dataclass(frozen=True)
-class BaseMessage(BaseResult):
+class BaseMessage(BaseProtoResult[ProtoMessageTypeT_contra]):
     parts: tuple[Any, ...]
 
     @property

src/yandex_cloud_ml_sdk/_messages/citations.py

@@ -11,7 +11,7 @@
 
 from yandex_cloud_ml_sdk._files.file import BaseFile
 from yandex_cloud_ml_sdk._search_indexes.search_index import BaseSearchIndex
-from yandex_cloud_ml_sdk._types.result import BaseResult
+from yandex_cloud_ml_sdk._types.result import BaseProtoResult
 
 from .base import BaseMessage
 
@@ -20,7 +20,7 @@
 
 
 @dataclasses.dataclass(frozen=True)
-class Citation(BaseResult):
+class Citation(BaseProtoResult[ProtoCitation]):
     sources: tuple[Source, ...]
 
     @classmethod
@@ -32,7 +32,7 @@ def _from_proto(cls, proto: ProtoCitation, sdk: BaseSDK) -> Citation:  # type: i
             )
         )
 
-class Source(BaseResult):
+class Source(BaseProtoResult[ProtoSource]):
     @property
     @abc.abstractmethod
     def type(self) -> str:
@@ -47,7 +47,7 @@ def _from_proto(cls, proto: ProtoSource, sdk: BaseSDK) -> Source:  # type: ignor
 
 
 @dataclasses.dataclass(frozen=True)
-class FileChunk(Source, BaseMessage):
+class FileChunk(Source, BaseMessage[ProtoSource]):
     search_index: BaseSearchIndex
     file: BaseFile | None
 

src/yandex_cloud_ml_sdk/_messages/message.py

@@ -37,7 +37,7 @@ class Author:
 
 
 @dataclasses.dataclass(frozen=True)
-class Message(BaseMessage, BaseResource):
+class Message(BaseMessage[ProtoMessage], BaseResource[ProtoMessage]):
     thread_id: str
     created_by: str
     created_at: datetime
@@ -79,7 +79,7 @@ def _kwargs_from_message(cls, proto: ProtoMessage, sdk: BaseSDK) -> dict[str, An
 
 
 @dataclasses.dataclass(frozen=True)
-class PartialMessage(BaseMessage, BaseResource):
+class PartialMessage(BaseMessage[MessageContent], BaseResource[MessageContent]):
     @classmethod
     def _kwargs_from_message(cls, proto: MessageContent, sdk: BaseSDK) -> dict[str, Any]:  # type: ignore[override]
         kwargs = super()._kwargs_from_message(proto, sdk=sdk)

src/yandex_cloud_ml_sdk/_models/completions/result.py

@@ -13,7 +13,7 @@
 from yandex_cloud_ml_sdk._tools.tool_call_list import ProtoCompletionsToolCallList, ToolCallList
 from yandex_cloud_ml_sdk._types.message import TextMessage
 from yandex_cloud_ml_sdk._types.proto import ProtoBased, SDKType
-from yandex_cloud_ml_sdk._types.result import BaseResult
+from yandex_cloud_ml_sdk._types.result import BaseProtoResult
 
 
 @dataclass(frozen=True)
@@ -109,7 +109,7 @@ def _from_proto(cls, *, proto: ProtoAlternative, sdk: SDKType) -> Alternative:
 
 
 @dataclass(frozen=True)
-class GPTModelResult(BaseResult[CompletionResponse], Sequence, HaveToolCalls[ToolCallTypeT]):
+class GPTModelResult(BaseProtoResult[CompletionResponse], Sequence, HaveToolCalls[ToolCallTypeT]):
     """A class representing the result of a GPT model completion request."""
     #: a tuple of alternatives generated by the model
     alternatives: tuple[Alternative[ToolCallTypeT], ...]

src/yandex_cloud_ml_sdk/_models/image_generation/result.py

@@ -1,29 +1,28 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, cast
+from typing import TYPE_CHECKING
 
 from typing_extensions import Self
 # pylint: disable-next=no-name-in-module
 from yandex.cloud.ai.foundation_models.v1.image_generation.image_generation_service_pb2 import ImageGenerationResponse
 
-from yandex_cloud_ml_sdk._types.result import BaseResult, ProtoMessage
+from yandex_cloud_ml_sdk._types.result import BaseProtoResult
 
 if TYPE_CHECKING:
     from yandex_cloud_ml_sdk._sdk import BaseSDK
 
 
 @dataclass(frozen=True, repr=False)
-class ImageGenerationModelResult(BaseResult):
+class ImageGenerationModelResult(BaseProtoResult[ImageGenerationResponse]):
     """This class represents the result of an image generation model inference."""
     #: the generated image in bytes
     image_bytes: bytes
     #: the version of the model used for generation
     model_version: str
 
     @classmethod
-    def _from_proto(cls, *, proto: ProtoMessage, sdk: BaseSDK) -> Self:  # pylint: disable=unused-argument
-        proto = cast(ImageGenerationResponse, proto)
+    def _from_proto(cls, *, proto: ImageGenerationResponse, sdk: BaseSDK) -> Self:  # pylint: disable=unused-argument
         return cls(
             image_bytes=proto.image,
             model_version=proto.model_version,

src/yandex_cloud_ml_sdk/_models/text_classifiers/result.py

@@ -1,15 +1,15 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Generic, Sequence, TypeVar, cast, overload
+from typing import TYPE_CHECKING, Generic, Sequence, TypeVar, overload
 
 from typing_extensions import Self
 # pylint: disable-next=no-name-in-module
 from yandex.cloud.ai.foundation_models.v1.text_classification.text_classification_service_pb2 import (
     FewShotTextClassificationResponse, TextClassificationResponse
 )
 
-from yandex_cloud_ml_sdk._types.result import BaseResult, ProtoMessage
+from yandex_cloud_ml_sdk._types.result import BaseProtoResult
 
 from .types import TextClassificationLabel
 
@@ -25,7 +25,7 @@
 
 
 @dataclass(frozen=True)
-class TextClassifiersModelResultBase(BaseResult, Sequence, Generic[TextClassificationResponseT]):
+class TextClassifiersModelResultBase(BaseProtoResult[TextClassificationResponseT], Sequence, Generic[TextClassificationResponseT]):
     """A class for text classifiers model results.
     It represents the common structure for the results returned by text classification models.
     """
@@ -37,8 +37,7 @@ class TextClassifiersModelResultBase(BaseResult, Sequence, Generic[TextClassific
     input_tokens: int
 
     @classmethod
-    def _from_proto(cls, *, proto: ProtoMessage, sdk: BaseSDK) -> Self:  # pylint: disable=unused-argument
-        proto = cast(TextClassificationResponseT, proto)
+    def _from_proto(cls, *, proto: TextClassificationResponseT, sdk: BaseSDK) -> Self:  # pylint: disable=unused-argument
         predictions = tuple(
             TextClassificationLabel(
                 label=p.label,