Add .add_files method for search indices (#68)

Author: vhaldemar

examples/async/assistants/search_indexes.py

@@ -28,7 +28,7 @@ async def main() -> None:
     files = await asyncio.gather(*file_coros)
 
     operation = await sdk.search_indexes.create_deferred(
-        files,
+        [files[0]],
         index_type=TextSearchIndexType(
             chunking_strategy=StaticIndexChunkingStrategy(
                 max_chunk_size_tokens=700,
@@ -45,11 +45,16 @@ async def main() -> None:
     await search_index.update(name="foo")
     print(f"now with a name {search_index=}")
 
-    # NB: it doesn't work at the moment
-    # index_files = [file async for file in search_index.list_files()]
-    # print(f"search index files: {index_files}")
-    # index_file = await search_index.get_file(index_files[0].id)
-    # print(f"search index file: {index_file}")
+    # We could also add files to index later:
+    add_operation = await search_index.add_files_deferred(files[1])
+    new_index_files = await add_operation.wait()
+    print(f"{new_index_files=}")
+
+    index_files = [file async for file in search_index.list_files()]
+    print(f"search index files: {index_files}")
+
+    index_file = await search_index.get_file(index_files[0].id)
+    print(f"search index file: {index_file}")
 
     for file in files:
         await file.delete()

examples/sync/assistants/search_indexes.py

@@ -26,7 +26,7 @@ def main() -> None:
         files.append(file)
 
     operation = sdk.search_indexes.create_deferred(
-        files,
+        [files[0]],
         index_type=TextSearchIndexType(
             chunking_strategy=StaticIndexChunkingStrategy(
                 max_chunk_size_tokens=700,
@@ -43,11 +43,16 @@ def main() -> None:
     search_index.update(name="foo")
     print(f"now with a name {search_index=}")
 
-    # NB: it doesn't work at the moment
-    # index_files = [file for file in search_index.list_files()]
-    # print(f"search index files: {index_files}")
-    # index_file = search_index.get_file(index_files[0].id)
-    # print(f"search index file: {index_file}")
+    # We could also add files to index later:
+    add_operation = search_index.add_files_deferred(files[1])
+    new_index_files = add_operation.wait()
+    print(f"{new_index_files=}")
+
+    index_files = [file for file in search_index.list_files()]
+    print(f"search index files: {index_files}")
+
+    index_file = search_index.get_file(index_files[0].id)
+    print(f"search index file: {index_file}")
 
     for file in files:
         file.delete()

src/yandex_cloud_ml_sdk/_search_indexes/search_index.py

@@ -3,24 +3,29 @@
 
 import dataclasses
 from datetime import datetime
-from typing import TYPE_CHECKING, Any, AsyncIterator, Iterator, TypeVar
+from typing import TYPE_CHECKING, Any, AsyncIterator, Iterator, TypeVar, cast
 
-from typing_extensions import Self
+from typing_extensions import Self, TypeAlias
 from yandex.cloud.ai.assistants.v1.searchindex.search_index_file_pb2 import SearchIndexFile as ProtoSearchIndexFile
 from yandex.cloud.ai.assistants.v1.searchindex.search_index_file_service_pb2 import (
-    GetSearchIndexFileRequest, ListSearchIndexFilesRequest, ListSearchIndexFilesResponse
+    BatchCreateSearchIndexFileRequest, BatchCreateSearchIndexFileResponse, GetSearchIndexFileRequest,
+    ListSearchIndexFilesRequest, ListSearchIndexFilesResponse
 )
 from yandex.cloud.ai.assistants.v1.searchindex.search_index_file_service_pb2_grpc import SearchIndexFileServiceStub
 from yandex.cloud.ai.assistants.v1.searchindex.search_index_pb2 import SearchIndex as ProtoSearchIndex
 from yandex.cloud.ai.assistants.v1.searchindex.search_index_service_pb2 import (
     DeleteSearchIndexRequest, DeleteSearchIndexResponse, UpdateSearchIndexRequest
 )
 from yandex.cloud.ai.assistants.v1.searchindex.search_index_service_pb2_grpc import SearchIndexServiceStub
+from yandex.cloud.operation.operation_pb2 import Operation as ProtoOperation
 
+from yandex_cloud_ml_sdk._files.file import BaseFile
 from yandex_cloud_ml_sdk._types.expiration import ExpirationConfig, ExpirationPolicyAlias
 from yandex_cloud_ml_sdk._types.misc import UNDEFINED, UndefinedOr, get_defined_value
+from yandex_cloud_ml_sdk._types.operation import AsyncOperation, Operation, OperationTypeT, ReturnsOperationMixin
 from yandex_cloud_ml_sdk._types.resource import ExpirableResource, safe_on_delete
 from yandex_cloud_ml_sdk._types.result import BaseResult
+from yandex_cloud_ml_sdk._utils.coerce import ResourceType, coerce_resource_ids
 from yandex_cloud_ml_sdk._utils.sync import run_sync, run_sync_generator
 
 from .file import SearchIndexFile
@@ -30,8 +35,11 @@
     from yandex_cloud_ml_sdk._sdk import BaseSDK
 
 
+SearchIndexFileTuple: TypeAlias = tuple[SearchIndexFile, ...]
+
+
 @dataclasses.dataclass(frozen=True)
-class BaseSearchIndex(ExpirableResource, BaseResult):
+class BaseSearchIndex(ExpirableResource, BaseResult, ReturnsOperationMixin[OperationTypeT]):
     @classmethod
     def _kwargs_from_message(cls, proto: ProtoSearchIndex, sdk: BaseSDK) -> dict[str, Any]:  # type: ignore[override]
         kwargs = super()._kwargs_from_message(proto, sdk=sdk)
@@ -128,6 +136,42 @@ async def _get_file(
 
         return SearchIndexFile._from_proto(proto=response, sdk=self._sdk)
 
+    # pylint: disable=unused-argument
+    async def _transform_add_files(self, proto: BatchCreateSearchIndexFileResponse, timeout: float) -> SearchIndexFileTuple:
+        return tuple(
+            SearchIndexFile._from_proto(proto=f, sdk=self._sdk)
+            for f in proto.files
+        )
+
+    @safe_on_delete
+    async def _add_files_deferred(
+        self,
+        files: ResourceType[BaseFile],
+        *,
+        timeout: float = 60,
+    ) -> OperationTypeT:
+        file_ids = coerce_resource_ids(files, BaseFile)
+        request = BatchCreateSearchIndexFileRequest(
+            file_ids=file_ids,
+            search_index_id=self.id
+        )
+
+        async with self._client.get_service_stub(SearchIndexFileServiceStub, timeout=timeout) as stub:
+            response = await self._client.call_service(
+                stub.BatchCreate,
+                request,
+                timeout=timeout,
+                expected_type=ProtoOperation
+            )
+
+        return self._operation_impl(
+            id=response.id,
+            sdk=self._sdk,
+            proto_result_type=BatchCreateSearchIndexFileResponse,
+            result_type=SearchIndexFileTuple,
+            transformer=self._transform_add_files
+        )
+
     async def _list_files(
         self,
         *,
@@ -161,7 +205,7 @@ async def _list_files(
 
 
 @dataclasses.dataclass(frozen=True)
-class RichSearchIndex(BaseSearchIndex):
+class RichSearchIndex(BaseSearchIndex[OperationTypeT]):
     folder_id: str
     name: str | None
     description: str | None
@@ -174,7 +218,9 @@ class RichSearchIndex(BaseSearchIndex):
     index_type: BaseSearchIndexType
 
 
-class AsyncSearchIndex(RichSearchIndex):
+class AsyncSearchIndex(RichSearchIndex[AsyncOperation[SearchIndexFileTuple]]):
+    _operation_impl = AsyncOperation[SearchIndexFileTuple]
+
     async def update(
         self,
         *,
@@ -224,12 +270,26 @@ async def list_files(
         ):
             yield file
 
+    async def add_files_deferred(
+        self,
+        files: ResourceType[BaseFile],
+        *,
+        timeout: float = 60,
+    ) -> AsyncOperation[SearchIndexFileTuple]:
+        return await self._add_files_deferred(
+            files=files,
+            timeout=timeout
+        )
+
 
-class SearchIndex(RichSearchIndex):
+# pylint: disable=protected-access
+class SearchIndex(RichSearchIndex[Operation[SearchIndexFileTuple]]):
+    _operation_impl = Operation[SearchIndexFileTuple]
     __update = run_sync(RichSearchIndex._update)
     __delete = run_sync(RichSearchIndex._delete)
     __get_file = run_sync(RichSearchIndex._get_file)
     __list_files = run_sync_generator(RichSearchIndex._list_files)
+    __add_files_deferred = run_sync(RichSearchIndex._add_files_deferred)
 
     def update(
         self,
@@ -279,5 +339,17 @@ def list_files(
             timeout=timeout,
         )
 
+    def add_files_deferred(
+        self,
+        files: ResourceType[BaseFile],
+        *,
+        timeout: float = 60,
+    ) -> Operation[SearchIndexFileTuple]:
+        # mypy is going crazy as always, with run_sync over generic
+        return cast(
+            Operation[SearchIndexFileTuple],
+            self.__add_files_deferred(files=files, timeout=timeout)
+        )
+
 
 SearchIndexTypeT = TypeVar('SearchIndexTypeT', bound=BaseSearchIndex)

src/yandex_cloud_ml_sdk/_types/operation.py

@@ -25,7 +25,7 @@
 logger = get_logger(__name__)
 
 AnyResultTypeT_co = TypeVar('AnyResultTypeT_co', covariant=True)
-ResultTypeT_co = TypeVar('ResultTypeT_co', bound=BaseResult, covariant=True)
+ResultTypeT_co = TypeVar('ResultTypeT_co', covariant=True)
 
 
 @dataclass(frozen=True)
@@ -183,6 +183,11 @@ def __repr__(self) -> str:
 
     # pylint: disable=unused-argument
     async def _default_result_transofrmer(self, proto: Any, timeout: float) -> ResultTypeT_co:
+        # 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, BaseResult)
+
         # 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(

src/yandex_cloud_ml_sdk/_types/result.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Protocol
+from typing import TYPE_CHECKING, Protocol, runtime_checkable
 
 from typing_extensions import Self
 
@@ -9,7 +9,7 @@
 if TYPE_CHECKING:
     from yandex_cloud_ml_sdk._sdk import BaseSDK
 
-
+@runtime_checkable
 class BaseResult(Protocol):
     @classmethod
     def _from_proto(cls, *, proto: ProtoMessage, sdk: BaseSDK) -> Self: