Skip to content

Commit 868e5ee

Browse files
committed
init commit
1 parent 202e5d2 commit 868e5ee

File tree

11 files changed

+270
-13
lines changed

11 files changed

+270
-13
lines changed

src/yandex_cloud_ml_sdk/_tools/domain.py

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,17 +16,30 @@
1616

1717

1818
class BaseTools(BaseDomain, Generic[FunctionToolsTypeT]):
19+
"""
20+
Сlass for tools functionality in Yandex Cloud ML SDK.
21+
22+
Provides common functionality for both synchronous and asynchronous tools.
23+
"""
1924
_functions_impl: type[FunctionToolsTypeT]
2025

2126
@cached_property
2227
def function(self) -> FunctionToolsTypeT:
28+
"""
29+
Get the function tools instance.
30+
"""
2331
return self._functions_impl(
2432
name='tools.function',
2533
sdk=self._sdk
2634
)
2735

2836
@cached_property
2937
def rephraser(self) -> RephraserFunction:
38+
"""
39+
Get the rephraser function instance.
40+
41+
The rephraser is used to modify user queries for better search results.
42+
"""
3043
return RephraserFunction(
3144
name='tools.rehraser',
3245
sdk=self._sdk,
@@ -41,7 +54,8 @@ def search_index(
4154
rephraser: UndefinedOr[RephraserInputType] = UNDEFINED,
4255
call_strategy: UndefinedOr[CallStrategyInputType] = UNDEFINED,
4356
) -> SearchIndexTool:
44-
"""Creates SearchIndexTool (not to be confused with :py:class:`~.SearchIndex`/:py:class:`~.AsyncSearchIndex`).
57+
"""
58+
Creates SearchIndexTool (not to be confused with :py:class:`~.SearchIndex`/:py:class:`~.AsyncSearchIndex`).
4559
4660
:param indexes: parameter takes :py:class:`~.BaseSearchIndex`, string with search index id,
4761
or a list of this values in any combination.
@@ -73,8 +87,17 @@ def search_index(
7387

7488

7589
class AsyncTools(BaseTools[AsyncFunctionTools]):
76-
_functions_impl = AsyncFunctionTools
90+
"""
91+
Asynchronous implementation of tools functionality.
7792
93+
Provides async versions of all tools methods.
94+
"""
95+
_functions_impl = AsyncFunctionTools
7896

7997
class Tools(BaseTools[FunctionTools]):
98+
"""
99+
Synchronous implementation of tools functionality.
100+
101+
Provides sync versions of all tools methods.
102+
"""
80103
_functions_impl = FunctionTools

src/yandex_cloud_ml_sdk/_tools/function.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,9 @@
1313

1414

1515
class BaseFunctionTools(BaseDomain, HaveToolCalls[ToolCallTypeT]):
16+
"""
17+
Base class for function tools in Yandex Cloud ML SDK.
18+
"""
1619
_call_impl: type[ToolCallTypeT]
1720

1821
def __call__(
@@ -23,6 +26,14 @@ def __call__(
2326
description: UndefinedOr[str] = UNDEFINED,
2427
strict: UndefinedOr[bool] = UNDEFINED,
2528
) -> FunctionTool:
29+
"""
30+
Create a function tool with given parameters.
31+
32+
:param parameters: Function parameters schema
33+
:param name: Optional function name (default: inferred from parameters)
34+
:param description: Optional function description
35+
:param strict: Whether to enforce strict parameter validation
36+
"""
2637
schema = schema_from_parameters(parameters)
2738
description_ = (
2839
get_defined_value(description, None) or
@@ -49,11 +60,20 @@ def __call__(
4960

5061

5162
class AsyncFunctionTools(BaseFunctionTools[AsyncToolCall]):
63+
"""
64+
Asynchronous version of function tools.
65+
"""
5266
_call_impl = AsyncToolCall
5367

5468

5569
class FunctionTools(BaseFunctionTools[ToolCall]):
70+
"""
71+
Synchronous version of function tools.
72+
"""
5673
_call_impl = ToolCall
5774

5875

5976
FunctionToolsTypeT = TypeVar('FunctionToolsTypeT', bound=BaseFunctionTools)
77+
"""
78+
Type variable representing any function tools type.
79+
"""

src/yandex_cloud_ml_sdk/_tools/function_call.py

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,12 @@
1717

1818
@dataclass(frozen=True)
1919
class BaseFunctionCall(ProtoBased[ProtoFunctionCall]):
20+
"""
21+
Base class representing a function call in Yandex Cloud ML SDK.
22+
23+
:param name: Name of the function being called
24+
:param arguments: Arguments passed to the function
25+
"""
2026
name: str
2127
arguments: JsonObject
2228
_proto_origin: ProtoFunctionCall = field(repr=False)
@@ -28,6 +34,11 @@ def _from_proto(
2834
proto: ProtoFunctionCall,
2935
**_,
3036
) -> Self:
37+
"""
38+
Create BaseFunctionCall instance from protobuf message.
39+
40+
:param proto: Protobuf message to convert
41+
"""
3142
return cls(
3243
name=proto.name,
3344
arguments=MessageToDict(proto.arguments),
@@ -36,11 +47,18 @@ def _from_proto(
3647

3748

3849
class AsyncFunctionCall(BaseFunctionCall):
39-
pass
50+
"""
51+
Asynchronous version of function call representation.
52+
"""
4053

4154

4255
class FunctionCall(BaseFunctionCall):
43-
pass
56+
"""
57+
Synchronous version of function call representation.
58+
"""
4459

4560

4661
FunctionCallTypeT = TypeVar('FunctionCallTypeT', bound=BaseFunctionCall)
62+
"""
63+
Type variable representing any function call type.
64+
"""

src/yandex_cloud_ml_sdk/_tools/search_index/call_strategy.py

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,12 +10,36 @@
1010
from yandex_cloud_ml_sdk._types.tools.function import FunctionDictType, validate_function_dict
1111

1212
CallStrategyStringType: TypeAlias = Literal['always']
13+
"""
14+
Type alias for string-based call strategy. Currently only supports 'always' value.
15+
"""
1316

1417
CallStrategyType: TypeAlias = Union[CallStrategyStringType, FunctionDictType]
18+
"""
19+
Type alias for all supported call strategy types.
20+
21+
Can be either:
22+
- A string literal ('always')
23+
- A function dictionary with instruction
24+
"""
1525
CallStrategyInputType: TypeAlias = Union[CallStrategyType, 'CallStrategy']
26+
"""
27+
Type alias for call strategy input types.
28+
29+
Can be either:
30+
- A CallStrategyType (string or function dict)
31+
- An existing CallStrategy instance
32+
"""
1633

1734

1835
class CallStrategy(ProtoBased[ProtoCallStrategy]):
36+
"""
37+
Implements call strategy for search index tools.
38+
39+
The call strategy determines when a tool should be called:
40+
- 'always': call the tool on every request
41+
- function dict: call based on function instruction
42+
"""
1943
_call_strategy: CallStrategyType
2044

2145
def __init__(self, call_strategy: CallStrategyType):
@@ -24,6 +48,9 @@ def __init__(self, call_strategy: CallStrategyType):
2448

2549
@property
2650
def value(self) -> CallStrategyType:
51+
"""
52+
Get the current call strategy value.
53+
"""
2754
return self._call_strategy
2855

2956
def _validate(self):

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

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,8 @@
1212

1313

1414
class RephraserFunction(BaseCompletions[Rephraser]):
15-
"""Function for creating Rephraser object, which incapsulating
16-
rephrasing settings.
15+
"""
16+
Function for creating Rephraser object, which incapsulating rephrasing settings.
1717
"""
1818

1919
_model_type = Rephraser
@@ -25,7 +25,8 @@ def __call__(
2525
*,
2626
model_version: str = 'latest',
2727
) -> Rephraser:
28-
"""Creates a Rephraser object, which incapsulating rephrasing settings.
28+
"""
29+
Creates a Rephraser object, which incapsulating rephrasing settings.
2930
3031
:param model_name:
3132
Model ID used for model uri definition in a resulting Rephraser object.
@@ -45,7 +46,6 @@ def __call__(
4546
refer to model_name parameter documentation for details.
4647
4748
:returns: Rephraser object, which incapsulating rephrasing settings
48-
4949
"""
5050

5151
name: str

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

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,12 @@ class _RephraserPseudoResult(BaseResult):
1919

2020

2121
class Rephraser(BaseModel[RephraserConfig, _RephraserPseudoResult], ProtoBased[ProtoRephraserOptions]):
22-
"""Class for incapsulating rephraser settings.
22+
"""
23+
Class for incapsulating rephraser settings.
24+
25+
Used to rewrite the last user message for search, incorporating context from the previous conversation.
2326
24-
Used to rewrite the last user message for search,
25-
incorporating context from the previous conversation.
27+
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>`_.
2628
"""
2729

2830
_config_type = RephraserConfig

src/yandex_cloud_ml_sdk/_tools/search_index/tool.py

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,14 +15,25 @@
1515

1616
@dataclass(frozen=True)
1717
class SearchIndexTool(BaseTool[ProtoSearchIndexTool]):
18+
"""
19+
Tool for working with search indexes in Yandex Cloud ML SDK.
20+
"""
21+
#: Tuple of search index IDs to use with this tool
1822
search_index_ids: tuple[str, ...]
19-
23+
#: Maximum number of results to return from search, optional
2024
max_num_results: int | None = None
25+
#: Rephraser instance for query rephrasing, optional
2126
rephraser: Rephraser | None = None
27+
#: Strategy for calling the search index, optional
2228
call_strategy: CallStrategy | None = None
2329

2430
@classmethod
2531
def _from_proto(cls, *, proto: ProtoSearchIndexTool, sdk: SDKType) -> SearchIndexTool:
32+
"""Create SearchIndexTool instance from protocol buffer message.
33+
34+
:param proto: Protocol buffer message to convert from
35+
:param sdk: SDK instance for type conversion
36+
"""
2637
max_num_results: int | None = None
2738
if proto.HasField("max_num_results"):
2839
max_num_results = proto.max_num_results.value
@@ -43,6 +54,11 @@ def _from_proto(cls, *, proto: ProtoSearchIndexTool, sdk: SDKType) -> SearchInde
4354
)
4455

4556
def _to_proto(self, proto_type: type[ProtoToolTypeT]) -> ProtoToolTypeT:
57+
"""Convert SearchIndexTool instance to protocol buffer message.
58+
59+
:param proto_type: Protocol buffer message type to convert to
60+
:raises AssertionError: If proto_type is not subclass of ProtoAssistantsTool
61+
"""
4662
assert issubclass(proto_type, ProtoAssistantsTool)
4763

4864
max_num_results: None | Int64Value = None

src/yandex_cloud_ml_sdk/_tools/tool.py

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,20 +16,42 @@
1616
from yandex_cloud_ml_sdk._types.schemas import JsonSchemaType
1717

1818
ProtoToolTypeT = TypeVar('ProtoToolTypeT', ProtoAssistantsTool, ProtoCompletionsTool)
19+
"""
20+
Type variable representing protobuf tool types.
21+
"""
1922

2023

2124
class BaseTool(ProtoBased[ProtoMessageTypeT]):
25+
"""
26+
Base class for all tools in Yandex Cloud ML SDK.
27+
"""
28+
2229
@classmethod
2330
@abc.abstractmethod
2431
def _from_proto(cls, *, proto: ProtoMessageTypeT, sdk: SDKType) -> BaseTool:
32+
"""
33+
Create tool instance from protobuf message.
34+
35+
:param proto: Protobuf message to convert
36+
:param sdk: SDK instance
37+
"""
2538
pass
2639

2740
@abc.abstractmethod
2841
def _to_proto(self, proto_type: type[ProtoToolTypeT]) -> ProtoToolTypeT:
42+
"""Convert tool to protobuf message.
43+
44+
:param proto_type: Protobuf message type to create
45+
"""
2946
pass
3047

3148
@classmethod
3249
def _from_upper_proto(cls, proto: ProtoToolTypeT, sdk: SDKType) -> BaseTool:
50+
"""Create tool instance from upper-level protobuf message.
51+
52+
:param proto: Protobuf message to convert
53+
:param sdk: SDK instance
54+
"""
3355
if proto.HasField('function'):
3456
return FunctionTool._from_proto(
3557
proto=proto.function,
@@ -52,13 +74,23 @@ def _from_upper_proto(cls, proto: ProtoToolTypeT, sdk: SDKType) -> BaseTool:
5274

5375

5476
ProtoFunctionTool = Union[ProtoCompletionsFunctionTool, ProtoAssistantsFunctionTool]
77+
"""
78+
Union type for function tool protobuf messages.
79+
"""
5580

5681

5782
@dataclass(frozen=True)
5883
class FunctionTool(BaseTool[ProtoFunctionTool]):
84+
"""
85+
Function tool representation in Yandex Cloud ML SDK.
86+
"""
87+
#: Name of the function
5988
name: str
89+
#: Optional function description
6090
description: str | None
91+
#: Function parameters schema
6192
parameters: JsonSchemaType
93+
#: Whether to enforce strict parameter validation
6294
strict: bool | None
6395

6496
# pylint: disable=unused-argument
@@ -69,6 +101,12 @@ def _from_proto(
69101
proto: ProtoFunctionTool,
70102
sdk:SDKType,
71103
) -> FunctionTool:
104+
"""
105+
Create FunctionTool from protobuf message.
106+
107+
:param proto: Protobuf message to convert
108+
:param sdk: SDK instance
109+
"""
72110
parameters = MessageToDict(proto.parameters)
73111

74112
strict: bool | None = None
@@ -83,6 +121,11 @@ def _from_proto(
83121
)
84122

85123
def _to_proto(self, proto_type: type[ProtoToolTypeT]) -> ProtoToolTypeT:
124+
"""Convert FunctionTool to protobuf message.
125+
126+
:param proto_type: Protobuf message type to create
127+
:raises ValueError: If strict validation is not supported for the proto type
128+
"""
86129
parameters = Struct()
87130
parameters.update(self.parameters)
88131

0 commit comments

Comments
 (0)