Add tool choice api (#125)

Author: vhaldemar

examples/async/function_calling/completions/tool_choice.py

@@ -0,0 +1,77 @@
+#!/usr/bin/env python3
+
+from __future__ import annotations
+
+import asyncio
+
+from yandex_cloud_ml_sdk import AsyncYCloudML
+
+SCHEMA = {
+    "type": "object",
+    "properties": {
+        "expression": {
+            "type": "string",
+            "description": "The mathematical expression to evaluate (e.g., '2 + 3 * 4').",
+        }
+    },
+    "required": ["expression"],
+}
+
+
+async def main() -> None:
+    sdk = AsyncYCloudML(folder_id='b1ghsjum2v37c2un8h64')
+    sdk.setup_default_logging()
+
+    calculator_tool = sdk.tools.function(
+        name="calculator_tool",
+        description="A simple calculator that performs basic arithmetic and @ operations.",
+        parameters=SCHEMA  # type: ignore[arg-type]
+    )
+    another_calculator = sdk.tools.function(
+        name="another_calculator",
+        description="A simple calculator that performs basic arithmetic and % operations.",
+        parameters=SCHEMA  # type: ignore[arg-type]
+    )
+
+    model = sdk.models.completions('yandexgpt', model_version='rc').configure(
+        tools=[calculator_tool, another_calculator],
+        temperature=0,
+        # auto is equivalent to default
+        # tool_choice='auto'
+    )
+
+    request = "How much it would be 7@8?"
+    result = await model.run(request)
+
+    # Model could call the tool, but it depends on many things, for example - model version.
+    # Right now I writing this example it does not calling the tool
+    assert result.status.name == 'FINAL'
+
+    # You could configure that you don't want to call any tool
+    model = model.configure(tool_choice='none')
+    result = await model.run(request)
+    assert result.status.name == 'FINAL'
+
+    # You could configure the model to always call some tool
+    model = model.configure(tool_choice='required')
+    result = await model.run(request)
+    assert result.status.name =='TOOL_CALLS'
+    assert result.tool_calls
+    assert len(result.tool_calls) == 1
+    assert result.tool_calls[0].function
+    assert result.tool_calls[0].function.name == 'calculator_tool'
+
+    # Or configure to call specific tool
+    model = model.configure(tool_choice={'type': 'function', 'function': {'name': 'another_calculator'}})
+    # You could pass just a function tool object instead of this big dict
+    model = model.configure(tool_choice=another_calculator)
+    result = await model.run(request)
+    assert result.status.name =='TOOL_CALLS'
+    assert result.tool_calls
+    assert len(result.tool_calls) == 1
+    assert result.tool_calls[0].function
+    assert result.tool_calls[0].function.name == 'another_calculator'
+
+
+if __name__ == '__main__':
+    asyncio.run(main())

examples/sync/function_calling/completions/tool_choice.py

@@ -0,0 +1,75 @@
+#!/usr/bin/env python3
+
+from __future__ import annotations
+
+from yandex_cloud_ml_sdk import YCloudML
+
+SCHEMA = {
+    "type": "object",
+    "properties": {
+        "expression": {
+            "type": "string",
+            "description": "The mathematical expression to evaluate (e.g., '2 + 3 * 4').",
+        }
+    },
+    "required": ["expression"],
+}
+
+
+def main() -> None:
+    sdk = YCloudML(folder_id='b1ghsjum2v37c2un8h64')
+    sdk.setup_default_logging()
+
+    calculator_tool = sdk.tools.function(
+        name="calculator_tool",
+        description="A simple calculator that performs basic arithmetic and @ operations.",
+        parameters=SCHEMA  # type: ignore[arg-type]
+    )
+    another_calculator = sdk.tools.function(
+        name="another_calculator",
+        description="A simple calculator that performs basic arithmetic and % operations.",
+        parameters=SCHEMA  # type: ignore[arg-type]
+    )
+
+    model = sdk.models.completions('yandexgpt', model_version='rc').configure(
+        tools=[calculator_tool, another_calculator],
+        temperature=0,
+        # auto is equivalent to default
+        # tool_choice='auto'
+    )
+
+    request = "How much it would be 7@8?"
+    result = model.run(request)
+
+    # Model could call the tool, but it depends on many things, for example - model version.
+    # Right now I writing this example it does not calling the tool
+    assert result.status.name == 'FINAL'
+
+    # You could configure that you don't want to call any tool
+    model = model.configure(tool_choice='none')
+    result = model.run(request)
+    assert result.status.name == 'FINAL'
+
+    # You could configure the model to always call some tool
+    model = model.configure(tool_choice='required')
+    result = model.run(request)
+    assert result.status.name =='TOOL_CALLS'
+    assert result.tool_calls
+    assert len(result.tool_calls) == 1
+    assert result.tool_calls[0].function
+    assert result.tool_calls[0].function.name == 'calculator_tool'
+
+    # Or configure to call specific tool
+    model = model.configure(tool_choice={'type': 'function', 'function': {'name': 'another_calculator'}})
+    # You could pass just a function tool object instead of this big dict
+    model = model.configure(tool_choice=another_calculator)
+    result = model.run(request)
+    assert result.status.name =='TOOL_CALLS'
+    assert result.tool_calls
+    assert len(result.tool_calls) == 1
+    assert result.tool_calls[0].function
+    assert result.tool_calls[0].function.name == 'another_calculator'
+
+
+if __name__ == '__main__':
+    main()

src/yandex_cloud_ml_sdk/_models/completions/config.py

@@ -12,6 +12,7 @@
 from yandex_cloud_ml_sdk._tools.tool import FunctionTool
 from yandex_cloud_ml_sdk._types.model_config import BaseModelConfig
 from yandex_cloud_ml_sdk._types.schemas import ResponseType
+from yandex_cloud_ml_sdk._types.tool_choice import ToolChoiceType
 from yandex_cloud_ml_sdk._utils.proto import ProtoEnumBase
 
 _m = ProtoReasoningOptions.ReasoningMode
@@ -35,3 +36,4 @@ class GPTModelConfig(BaseModelConfig):
     response_format: ResponseType | None = None
     tools: Sequence[CompletionTool] | CompletionTool | None = None
     parallel_tool_calls: bool | None = None
+    tool_choice: ToolChoiceType | None = None

src/yandex_cloud_ml_sdk/_models/completions/model.py

@@ -9,6 +9,7 @@
 from typing_extensions import Self, override
 from yandex.cloud.ai.foundation_models.v1.text_common_pb2 import CompletionOptions, ReasoningOptions
 from yandex.cloud.ai.foundation_models.v1.text_common_pb2 import Tool as ProtoCompletionsTool
+from yandex.cloud.ai.foundation_models.v1.text_common_pb2 import ToolChoice as ProtoToolChoice
 from yandex.cloud.ai.foundation_models.v1.text_generation.text_generation_service_pb2 import (
     BatchCompletionMetadata, BatchCompletionRequest, BatchCompletionResponse, CompletionRequest, CompletionResponse,
     TokenizeResponse
@@ -29,6 +30,8 @@
 )
 from yandex_cloud_ml_sdk._types.operation import AsyncOperation, Operation
 from yandex_cloud_ml_sdk._types.schemas import ResponseType, make_response_format_kwargs
+from yandex_cloud_ml_sdk._types.tool_choice import ToolChoiceType
+from yandex_cloud_ml_sdk._types.tool_choice import coerce_to_proto as coerce_to_proto_tool_choice
 from yandex_cloud_ml_sdk._types.tuning.datasets import TuningDatasetsType
 from yandex_cloud_ml_sdk._types.tuning.optimizers import BaseOptimizer
 from yandex_cloud_ml_sdk._types.tuning.schedulers import BaseScheduler
@@ -85,6 +88,7 @@ def configure(  # type: ignore[override]
         response_format: UndefinedOr[ResponseType] = UNDEFINED,
         tools: UndefinedOr[Sequence[CompletionTool] | CompletionTool] = UNDEFINED,
         parallel_tool_calls: UndefinedOr[bool] = UNDEFINED,
+        tool_choice: UndefinedOr[ToolChoiceType] = UNDEFINED,
     ) -> Self:
         return super().configure(
             temperature=temperature,
@@ -93,6 +97,7 @@ def configure(  # type: ignore[override]
             response_format=response_format,
             tools=tools,
             parallel_tool_calls=parallel_tool_calls,
+            tool_choice=tool_choice,
         )
 
     def _make_completion_options(self, *, stream: bool | None) -> CompletionOptions:
@@ -132,17 +137,22 @@ def _make_request(
         if c.parallel_tool_calls is not None:
             parallel_tool_calls = BoolValue(value=c.parallel_tool_calls)
 
+        tool_choice: None | ProtoToolChoice = None
+        if c.tool_choice is not None:
+            tool_choice = coerce_to_proto_tool_choice(c.tool_choice, expected_type=ProtoToolChoice)
+
         return CompletionRequest(
             model_uri=self._uri,
             completion_options=self._make_completion_options(stream=stream),
             messages=messages_to_proto(messages),
             tools=[tool._to_proto(ProtoCompletionsTool) for tool in tools],
             parallel_tool_calls=parallel_tool_calls,
+            tool_choice=tool_choice,
             **response_format_kwargs,
         )
 
     def _make_batch_request(self, dataset_id: str) -> BatchCompletionRequest:
-        for field in ('tools', 'response_format'):
+        for field in ('tools', 'response_format', 'tool_choice', 'parallel_tool_calls'):
             value = getattr(self.config, field)
             if value is not None:
                 warnings.warn(

src/yandex_cloud_ml_sdk/_types/tool_choice.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+from typing import Literal, TypedDict, TypeVar, Union, cast
+
+from typing_extensions import TypeAlias
+# pylint: disable=no-name-in-module
+from yandex.cloud.ai.foundation_models.v1.text_common_pb2 import ToolChoice as ProtoCompletionsToolChoice
+
+from yandex_cloud_ml_sdk._tools.tool import FunctionTool
+
+ProtoToolChoice: TypeAlias = ProtoCompletionsToolChoice
+ProtoToolChoiceTypeT = TypeVar('ProtoToolChoiceTypeT', bound=ProtoToolChoice)
+
+
+class FunctionNameType(TypedDict):
+    name: str
+
+
+class ToolChoiceDictType(TypedDict):
+    type: Literal['function']
+    function: FunctionNameType
+
+
+ToolChoiceStringType: TypeAlias = Literal[
+    'none', 'None', 'NONE',
+    'auto', 'Auto', 'AUTO',
+    'required', 'Required', 'REQUIRED'
+]
+
+ToolChoiceType: TypeAlias = Union[ToolChoiceStringType, ToolChoiceDictType, FunctionTool]
+
+STRING_TOOL_CHOICES = ('NONE', 'AUTO', 'REQUIRED')
+
+
+def coerce_to_proto(
+    tool_choice: ToolChoiceType, expected_type: type[ProtoToolChoiceTypeT]
+) -> ProtoToolChoiceTypeT:
+    if isinstance(tool_choice, str):
+        tool_choice = cast(ToolChoiceStringType, tool_choice.upper())
+        if tool_choice not in STRING_TOOL_CHOICES:
+            raise ValueError(f'wrong {tool_choice=}, use one of {STRING_TOOL_CHOICES}')
+
+        tool_choice_value = expected_type.ToolChoiceMode.Value(tool_choice)
+
+        return expected_type(mode=tool_choice_value)
+
+    if isinstance(tool_choice, dict):
+        if (
+            tool_choice.get('type') != 'function' or
+            not isinstance(tool_choice.get('function'), dict) or
+            not isinstance(tool_choice['function'].get('name'), str)
+        ):
+            raise ValueError(
+                'wrong dict structure for tool_choice, expected '
+                '`{"type": "function", "function": {"name": function_name}}`, '
+                'got {tool_choice}'
+            )
+
+        tool_choice = cast(ToolChoiceDictType, tool_choice)
+
+        return expected_type(function_name=tool_choice['function']['name'])
+
+    if isinstance(tool_choice, FunctionTool):
+        return expected_type(function_name=tool_choice.name)
+
+    raise TypeError(f'wrong {type(tool_choice)=}, expected string or dict')

src/yandex_cloud_ml_sdk/_utils/proto.py

@@ -101,7 +101,7 @@ def _coerce(cls, value: str | int | ProtoEnumBase) -> Self:
         if isinstance(value, str):
             if member := cls.__members__.get(value.upper()):
                 return member
-            raise ValueError(f'wrong value "{value}" for use as an alisas for {cls}')
+            raise ValueError(f'wrong value "{value}" for use as an alias for {cls}')
         raise TypeError(f'wrong type "{type(value)}" for use as an alias for {cls}')
 
     def _to_proto(self) -> int:

tests/conftest.py

@@ -56,7 +56,7 @@ def patch_operation(request, monkeypatch):
 
 @pytest.fixture(name='folder_id')
 def fixture_folder_id():
-    return 'yc.fomo.storage.prod.service'
+    return 'b1ghsjum2v37c2un8h64'
 
 
 @pytest.fixture(name='servicers')