init commit

Author: andy-fisher

src/yandex_cloud_ml_sdk/_tools/domain.py

@@ -16,17 +16,30 @@
 
 
 class BaseTools(BaseDomain, Generic[FunctionToolsTypeT]):
+    """
+    Сlass for tools functionality in Yandex Cloud ML SDK.
+
+    Provides common functionality for both synchronous and asynchronous tools.
+    """
     _functions_impl: type[FunctionToolsTypeT]
 
     @cached_property
     def function(self) -> FunctionToolsTypeT:
+        """
+        Get the function tools instance.
+        """
         return self._functions_impl(
             name='tools.function',
             sdk=self._sdk
         )
 
     @cached_property
     def rephraser(self) -> RephraserFunction:
+        """
+        Get the rephraser function instance.
+
+        The rephraser is used to modify user queries for better search results.
+        """
         return RephraserFunction(
             name='tools.rehraser',
             sdk=self._sdk,
@@ -41,7 +54,8 @@ def search_index(
         rephraser: UndefinedOr[RephraserInputType] = UNDEFINED,
         call_strategy: UndefinedOr[CallStrategyInputType] = UNDEFINED,
     ) -> SearchIndexTool:
-        """Creates SearchIndexTool (not to be confused with :py:class:`~.SearchIndex`/:py:class:`~.AsyncSearchIndex`).
+        """
+        Creates SearchIndexTool (not to be confused with :py:class:`~.SearchIndex`/:py:class:`~.AsyncSearchIndex`).
 
         :param indexes: parameter takes :py:class:`~.BaseSearchIndex`, string with search index id,
             or a list of this values in any combination.
@@ -73,8 +87,17 @@ def search_index(
 
 
 class AsyncTools(BaseTools[AsyncFunctionTools]):
-    _functions_impl = AsyncFunctionTools
+    """
+    Asynchronous implementation of tools functionality.
 
+    Provides async versions of all tools methods.
+    """
+    _functions_impl = AsyncFunctionTools
 
 class Tools(BaseTools[FunctionTools]):
+    """
+    Synchronous implementation of tools functionality.
+
+    Provides sync versions of all tools methods.
+    """
     _functions_impl = FunctionTools

src/yandex_cloud_ml_sdk/_tools/function.py

@@ -13,6 +13,9 @@
 
 
 class BaseFunctionTools(BaseDomain, HaveToolCalls[ToolCallTypeT]):
+    """
+    Base class for function tools in Yandex Cloud ML SDK.
+    """
     _call_impl: type[ToolCallTypeT]
 
     def __call__(
@@ -23,6 +26,14 @@ def __call__(
         description: UndefinedOr[str] = UNDEFINED,
         strict: UndefinedOr[bool] = UNDEFINED,
     ) -> FunctionTool:
+        """
+        Create a function tool with given parameters.
+        
+        :param parameters: Function parameters schema
+        :param name: Optional function name (default: inferred from parameters)
+        :param description: Optional function description
+        :param strict: Whether to enforce strict parameter validation
+        """
         schema = schema_from_parameters(parameters)
         description_ = (
             get_defined_value(description, None) or
@@ -49,11 +60,20 @@ def __call__(
 
 
 class AsyncFunctionTools(BaseFunctionTools[AsyncToolCall]):
+    """
+    Asynchronous version of function tools.
+    """
     _call_impl = AsyncToolCall
 
 
 class FunctionTools(BaseFunctionTools[ToolCall]):
+    """
+    Synchronous version of function tools.
+    """
     _call_impl = ToolCall
 
 
 FunctionToolsTypeT = TypeVar('FunctionToolsTypeT', bound=BaseFunctionTools)
+"""
+Type variable representing any function tools type.
+"""

src/yandex_cloud_ml_sdk/_tools/function_call.py

@@ -17,6 +17,12 @@
 
 @dataclass(frozen=True)
 class BaseFunctionCall(ProtoBased[ProtoFunctionCall]):
+    """
+    Base class representing a function call in Yandex Cloud ML SDK.
+    
+    :param name: Name of the function being called
+    :param arguments: Arguments passed to the function
+    """
     name: str
     arguments: JsonObject
     _proto_origin: ProtoFunctionCall = field(repr=False)
@@ -28,6 +34,11 @@ def _from_proto(
         proto: ProtoFunctionCall,
         **_,
     ) -> Self:
+        """
+        Create BaseFunctionCall instance from protobuf message.
+        
+        :param proto: Protobuf message to convert
+        """
         return cls(
             name=proto.name,
             arguments=MessageToDict(proto.arguments),
@@ -36,11 +47,18 @@ def _from_proto(
 
 
 class AsyncFunctionCall(BaseFunctionCall):
-    pass
+    """
+    Asynchronous version of function call representation.
+    """
 
 
 class FunctionCall(BaseFunctionCall):
-    pass
+    """
+    Synchronous version of function call representation.
+    """
 
 
 FunctionCallTypeT = TypeVar('FunctionCallTypeT', bound=BaseFunctionCall)
+"""
+Type variable representing any function call type.
+"""

src/yandex_cloud_ml_sdk/_tools/search_index/call_strategy.py

@@ -10,12 +10,36 @@
 from yandex_cloud_ml_sdk._types.tools.function import FunctionDictType, validate_function_dict
 
 CallStrategyStringType: TypeAlias = Literal['always']
+"""
+Type alias for string-based call strategy. Currently only supports 'always' value.
+"""
 
 CallStrategyType: TypeAlias = Union[CallStrategyStringType, FunctionDictType]
+"""
+Type alias for all supported call strategy types.
+
+Can be either:
+- A string literal ('always')
+- A function dictionary with instruction
+"""
 CallStrategyInputType: TypeAlias = Union[CallStrategyType, 'CallStrategy']
+"""
+Type alias for call strategy input types.
+
+Can be either:
+- A CallStrategyType (string or function dict)
+- An existing CallStrategy instance
+"""
 
 
 class CallStrategy(ProtoBased[ProtoCallStrategy]):
+    """
+    Implements call strategy for search index tools.
+
+    The call strategy determines when a tool should be called:
+    - 'always': call the tool on every request
+    - function dict: call based on function instruction
+    """
     _call_strategy: CallStrategyType
 
     def __init__(self, call_strategy: CallStrategyType):
@@ -24,6 +48,9 @@ def __init__(self, call_strategy: CallStrategyType):
 
     @property
     def value(self) -> CallStrategyType:
+        """
+        Get the current call strategy value.
+        """
         return self._call_strategy
 
     def _validate(self):

src/yandex_cloud_ml_sdk/_tools/search_index/rephraser/function.py

@@ -12,8 +12,8 @@
 
 
 class RephraserFunction(BaseCompletions[Rephraser]):
-    """Function for creating Rephraser object, which incapsulating
-    rephrasing settings.
+    """
+    Function for creating Rephraser object, which incapsulating rephrasing settings.
     """
 
     _model_type = Rephraser
@@ -25,7 +25,8 @@ def __call__(
         *,
         model_version: str = 'latest',
     ) -> Rephraser:
-        """Creates a Rephraser object, which incapsulating rephrasing settings.
+        """
+        Creates a Rephraser object, which incapsulating rephrasing settings.
 
         :param model_name:
             Model ID used for model uri definition in a resulting Rephraser object.
@@ -45,7 +46,6 @@ def __call__(
             refer to model_name parameter documentation for details.
 
         :returns: Rephraser object, which incapsulating rephrasing settings
-
         """
 
         name: str

src/yandex_cloud_ml_sdk/_tools/search_index/rephraser/model.py

@@ -19,10 +19,12 @@ class _RephraserPseudoResult(BaseResult):
 
 
 class Rephraser(BaseModel[RephraserConfig, _RephraserPseudoResult], ProtoBased[ProtoRephraserOptions]):
-    """Class for incapsulating rephraser settings.
+    """
+    Class for incapsulating rephraser settings.
+
+    Used to rewrite the last user message for search, incorporating context from the previous conversation.
 
-    Used to rewrite the last user message for search,
-    incorporating context from the previous conversation.
+    For usage search index tool with and without rephraser example see `rephraser example <https://github.com/yandex-cloud/yandex-cloud-ml-sdk/blob/master/examples/sync/assistants/rephraser.py>`_.
     """
 
     _config_type = RephraserConfig

src/yandex_cloud_ml_sdk/_tools/search_index/tool.py

@@ -15,14 +15,25 @@
 
 @dataclass(frozen=True)
 class SearchIndexTool(BaseTool[ProtoSearchIndexTool]):
+    """
+    Tool for working with search indexes in Yandex Cloud ML SDK.
+    """
+    #: Tuple of search index IDs to use with this tool
     search_index_ids: tuple[str, ...]
-
+    #: Maximum number of results to return from search, optional
     max_num_results: int | None = None
+    #: Rephraser instance for query rephrasing, optional
     rephraser: Rephraser | None = None
+    #: Strategy for calling the search index, optional
     call_strategy: CallStrategy | None = None
 
     @classmethod
     def _from_proto(cls, *, proto: ProtoSearchIndexTool, sdk: SDKType) -> SearchIndexTool:
+        """Create SearchIndexTool instance from protocol buffer message.
+
+        :param proto: Protocol buffer message to convert from
+        :param sdk: SDK instance for type conversion
+        """
         max_num_results: int | None = None
         if proto.HasField("max_num_results"):
             max_num_results = proto.max_num_results.value
@@ -43,6 +54,11 @@ def _from_proto(cls, *, proto: ProtoSearchIndexTool, sdk: SDKType) -> SearchInde
         )
 
     def _to_proto(self, proto_type: type[ProtoToolTypeT]) -> ProtoToolTypeT:
+        """Convert SearchIndexTool instance to protocol buffer message.
+
+        :param proto_type: Protocol buffer message type to convert to
+        :raises AssertionError: If proto_type is not subclass of ProtoAssistantsTool
+        """
         assert issubclass(proto_type, ProtoAssistantsTool)
 
         max_num_results: None | Int64Value = None

src/yandex_cloud_ml_sdk/_tools/tool.py

@@ -16,20 +16,42 @@
 from yandex_cloud_ml_sdk._types.schemas import JsonSchemaType
 
 ProtoToolTypeT = TypeVar('ProtoToolTypeT', ProtoAssistantsTool, ProtoCompletionsTool)
+"""
+Type variable representing protobuf tool types.
+"""
 
 
 class BaseTool(ProtoBased[ProtoMessageTypeT]):
+    """
+    Base class for all tools in Yandex Cloud ML SDK.
+    """
+    
     @classmethod
     @abc.abstractmethod
     def _from_proto(cls, *, proto: ProtoMessageTypeT, sdk: SDKType) -> BaseTool:
+        """
+        Create tool instance from protobuf message.
+        
+        :param proto: Protobuf message to convert
+        :param sdk: SDK instance
+        """
         pass
 
     @abc.abstractmethod
     def _to_proto(self, proto_type: type[ProtoToolTypeT]) -> ProtoToolTypeT:
+        """Convert tool to protobuf message.
+        
+        :param proto_type: Protobuf message type to create
+        """
         pass
 
     @classmethod
     def _from_upper_proto(cls, proto: ProtoToolTypeT, sdk: SDKType) -> BaseTool:
+        """Create tool instance from upper-level protobuf message.
+        
+        :param proto: Protobuf message to convert
+        :param sdk: SDK instance
+        """
         if proto.HasField('function'):
             return FunctionTool._from_proto(
                 proto=proto.function,
@@ -52,13 +74,23 @@ def _from_upper_proto(cls, proto: ProtoToolTypeT, sdk: SDKType) -> BaseTool:
 
 
 ProtoFunctionTool = Union[ProtoCompletionsFunctionTool, ProtoAssistantsFunctionTool]
+"""
+Union type for function tool protobuf messages.
+"""
 
 
 @dataclass(frozen=True)
 class FunctionTool(BaseTool[ProtoFunctionTool]):
+    """
+    Function tool representation in Yandex Cloud ML SDK.
+    """
+    #: Name of the function
     name: str
+    #: Optional function description
     description: str | None
+    #: Function parameters schema
     parameters: JsonSchemaType
+    #: Whether to enforce strict parameter validation
     strict: bool | None
 
     # pylint: disable=unused-argument
@@ -69,6 +101,12 @@ def _from_proto(
         proto: ProtoFunctionTool,
         sdk:SDKType,
     ) -> FunctionTool:
+        """
+        Create FunctionTool from protobuf message.
+        
+        :param proto: Protobuf message to convert
+        :param sdk: SDK instance
+        """
         parameters = MessageToDict(proto.parameters)
 
         strict: bool | None = None
@@ -83,6 +121,11 @@ def _from_proto(
         )
 
     def _to_proto(self, proto_type: type[ProtoToolTypeT]) -> ProtoToolTypeT:
+        """Convert FunctionTool to protobuf message.
+        
+        :param proto_type: Protobuf message type to create
+        :raises ValueError: If strict validation is not supported for the proto type
+        """
         parameters = Struct()
         parameters.update(self.parameters)