Fix all old tests

Author: vhaldemar

examples/async/completions/batch_async.py

@@ -57,7 +57,7 @@ async def main() -> None:
 
     with tqdm.tqdm(total=len(input_data), desc='creatig tasks') as t:
         operations = await run_chunked_tasks(
-            function=model.run_async,
+            function=model.run_deferred,
             data=input_data,
             chunk_size=100,
             tqdm_callback=t.update,

examples/async/completions/chat.py

@@ -14,7 +14,7 @@ async def main() -> None:
     model = sdk.models.completions('yandexgpt')
     model = model.configure(temperature=0.5)
 
-    messages: list[dict[str, str] | str] = [{'role': 'system', 'text': 'Ты - Аркадий'}]
+    messages: list = [{'role': 'system', 'text': 'Ты - Аркадий'}]
     while True:
         message = input()
         messages.append(message)

examples/async/function_calling/completions/raw_tool_calls_processing.py

@@ -17,7 +17,7 @@ def weather(location: str, date: str) -> str:
     return "-10 celsius"
 
 
-def process_tool_calls(tool_calls) -> list[dict]:
+def process_tool_calls(tool_calls) -> dict[str, list[dict]]:
     """
     This function is an example how you could organize
     dispatching of function calls in general case
@@ -35,7 +35,7 @@ def process_tool_calls(tool_calls) -> list[dict]:
 
         function = function_map[tool_call.function.name]
 
-        answer = function(**tool_call.function.arguments)
+        answer = function(**tool_call.function.arguments)  # type: ignore[operator]
 
         result.append({'name': tool_call.function.name, 'content': answer})
 
@@ -86,7 +86,7 @@ class Weather(BaseModel):
 
     for question in ["How much it would be 7@8?", "What is the weather like in Paris at 12 of March?"]:
         # it is required to carefully maintain context for passing tool_results back to the model after function call
-        messages = [
+        messages: list = [
             {"role": "system", "text": "Please use English language for answer"},
             question
         ]

src/yandex_cloud_ml_sdk/_models/completions/message.py

@@ -74,7 +74,7 @@ def messages_to_proto(messages: MessageInputType) -> list[ProtoMessage]:
             kwargs = {'role': 'user', 'text': message}
         elif isinstance(message, TextMessageProtocol):
             if isinstance(message, TextMessageWithToolCallsProtocol) and message.tool_calls:
-                # pylint: disable[protected-access]
+                # pylint: disable=protected-access
                 kwargs = {'role': message.role, 'tool_call_list': message.tool_calls._proto_origin}
             else:
                 kwargs = {'role': message.role, 'text': message.text}
@@ -91,7 +91,7 @@ def messages_to_proto(messages: MessageInputType) -> list[ProtoMessage]:
                     'tool_result_list': tool_results,
                 }
             else:
-                raise ValueError(f'{message=!r} should have a "text" or "tool_results" key')
+                raise TypeError(f'{message=!r} should have a "text" or "tool_results" key')
         else:
             raise TypeError(f'{message=!r} should be str, dict with "text" or "tool_results" key or TextMessage instance')
 

src/yandex_cloud_ml_sdk/_tools/function.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TypeVar
+from typing import TypeVar, cast
 
 from yandex_cloud_ml_sdk._types.domain import BaseDomain
 from yandex_cloud_ml_sdk._types.misc import UNDEFINED, UndefinedOr, get_defined_value
@@ -21,8 +21,14 @@ def __call__(
         description: UndefinedOr[str] = UNDEFINED,
     ) -> FunctionTool:
         schema = schema_from_parameters(parameters)
-        description_ = get_defined_value(description, None) or schema.get('description')
-        name_ = get_defined_value(name, None) or schema.get('title')
+        description_ = (
+            get_defined_value(description, None) or
+            cast(str | None, schema.get('description'))
+        )
+        name_ = (
+            get_defined_value(name, None) or
+            cast(str | None, schema.get('title'))
+        )
 
         if not name_:
             raise TypeError(

src/yandex_cloud_ml_sdk/_tools/tool_call.py

@@ -1,7 +1,6 @@
 # pylint: disable=no-name-in-module
 from __future__ import annotations
 
-from collections.abc import Iterable
 from dataclasses import dataclass, field
 from typing import Generic, TypeVar, Union
 

src/yandex_cloud_ml_sdk/_tools/tool_result.py

@@ -78,7 +78,7 @@ def tool_results_to_proto(
         ProtoCompletionsToolResultList: ProtoCompletionsToolResult,
     }[proto_type])
 
-    tool_results = coerce_tuple(tool_results, (dict, ))
+    tool_results = coerce_tuple(tool_results, cast(type[FunctionResultDict], dict))
 
     proto_tool_results: list[object] = []
     for tool_result in tool_results:

src/yandex_cloud_ml_sdk/_types/operation.py

@@ -4,7 +4,7 @@
 import abc
 import asyncio
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Awaitable, Callable, Generic, Iterable, TypeVar, cast
+from typing import TYPE_CHECKING, Any, Awaitable, Callable, Generic, Iterable, TypeVar, cast, get_origin
 
 from typing_extensions import Self
 # pylint: disable-next=no-name-in-module
@@ -186,13 +186,17 @@ async def _default_result_transofrmer(self, proto: Any, timeout: float) -> Resul
         # NB: default_result_transformer should be used only with _result_type
         # which are BaseResult-compatible, but I don't know how to express it with typing,
         # maybe we need special operation class, which support transforming (probably a base one)
-        assert isinstance(self._result_type, ProtoBasedType)
+        # NB: issubclass don't like if instead of SomeClass object pass SomeClass[T];
+        # because we use _result_type also for a generic typing reasons, sometimes it requires
+        # unwrapping for issubclass check
+        result_type = get_origin(self._result_type) or self._result_type
+        assert issubclass(result_type, ProtoBasedType), f'{self._result_type} is not ProtoBasedType'
 
         # NB: mypy can't figure out that self._result_type._from_proto is
         # returning instance of self._result_type which is also is a ResultTypeT_co
         return cast(
             ResultTypeT_co,
-            self._result_type._from_proto(proto=proto, sdk=self._sdk)
+            self._result_type._from_proto(proto=proto, sdk=self._sdk)  # type: ignore[attr-defined]
         )
 
     @property

tests/models/test_image_generation.py

@@ -63,7 +63,7 @@ def check_messages(messages, expected):
     assert messages[0].weight == 2
     assert messages[1].weight == 1
 
-    messages = messages_to_proto(Alternative(role='foo', text='bar', status=None))
+    messages = messages_to_proto(Alternative(role='foo', text='bar', status=None, tool_calls=None))
     check_messages(messages, ['bar'])
     assert messages[0].weight == 0
 
@@ -73,8 +73,8 @@ def check_messages(messages, expected):
 
     gpt_model_result = GPTModelResult(
         alternatives=[
-            Alternative(role='1', text='1', status=None),
-            Alternative(role='2', text='2', status=None),
+            Alternative(role='1', text='1', status=None, tool_calls=None),
+            Alternative(role='2', text='2', status=None, tool_calls=None),
         ],
         usage=None,
         model_version=''

tests/types/test_structured_output.py

@@ -7,8 +7,8 @@
 
 import pytest
 
-import yandex_cloud_ml_sdk._types.structured_output
-from yandex_cloud_ml_sdk._types.structured_output import schema_from_response_format
+import yandex_cloud_ml_sdk._types.schemas
+from yandex_cloud_ml_sdk._types.schemas import schema_from_response_format
 
 
 def test_string_type() -> None:
@@ -25,7 +25,7 @@ def test_dict_type() -> None:
 
 @pytest.mark.require_env('pydantic')
 def test_pydantic_model() -> None:
-    assert yandex_cloud_ml_sdk._types.structured_output.PYDANTIC is True
+    assert yandex_cloud_ml_sdk._types.schemas.PYDANTIC is True
 
     import pydantic
 
@@ -116,25 +116,25 @@ class B:
 @pytest.fixture(name='no_pydantic')
 def fixture_no_pydantic(monkeypatch) -> typing.Iterator[bool]:
     # pylint: disable=reimported
-    sys.modules.pop(yandex_cloud_ml_sdk._types.structured_output.__name__, None)
+    sys.modules.pop(yandex_cloud_ml_sdk._types.schemas.__name__, None)
 
     with monkeypatch.context() as m:
         m.setitem(sys.modules, 'pydantic', None)
-        import yandex_cloud_ml_sdk._types.structured_output as _m
+        import yandex_cloud_ml_sdk._types.schemas as _m
 
         yield True
 
-    sys.modules.pop(yandex_cloud_ml_sdk._types.structured_output.__name__, None)
-    import yandex_cloud_ml_sdk._types.structured_output as _m2
+    sys.modules.pop(yandex_cloud_ml_sdk._types.schemas.__name__, None)
+    import yandex_cloud_ml_sdk._types.schemas as _m2
 
     assert _m
     assert _m2
 
 
 def test_no_pydantic(no_pydantic) -> None:
     assert no_pydantic
-    assert yandex_cloud_ml_sdk._types.structured_output.PYDANTIC is False
-    assert yandex_cloud_ml_sdk._types.structured_output.PYDANTIC_V2 is False
+    assert yandex_cloud_ml_sdk._types.schemas.PYDANTIC is False
+    assert yandex_cloud_ml_sdk._types.schemas.PYDANTIC_V2 is False
 
     class A:
         a: int