diff --git a/examples/async/assistants/hybrid_search_index.py b/examples/async/assistants/hybrid_search_index.py new file mode 100755 index 0000000..bd85448 --- /dev/null +++ b/examples/async/assistants/hybrid_search_index.py @@ -0,0 +1,78 @@ +#!/usr/bin/env python3 + +from __future__ import annotations + +import asyncio +import pathlib +import pprint + +from yandex_cloud_ml_sdk import AsyncYCloudML +from yandex_cloud_ml_sdk.search_indexes import ( + HybridSearchIndexType, ReciprocalRankFusionIndexCombinationStrategy, StaticIndexChunkingStrategy, + TextSearchIndexType, VectorSearchIndexType +) + + +def local_path(path: str) -> pathlib.Path: + return pathlib.Path(__file__).parent / path + + +async def main() -> None: + sdk = AsyncYCloudML(folder_id='b1ghsjum2v37c2un8h64') + + file_coros = ( + sdk.files.upload( + local_path(path), + ttl_days=5, + expiration_policy="static", + ) + for path in ['turkey_example.txt', 'maldives_example.txt'] + ) + files = await asyncio.gather(*file_coros) + + # How to create search index with all default settings: + operation = await sdk.search_indexes.create_deferred( + files, + index_type=HybridSearchIndexType() + ) + default_search_index = await operation.wait() + print("new hybrid search index with default settings:") + pprint.pprint(default_search_index) + + # But you could override any default: + operation = await sdk.search_indexes.create_deferred( + files, + index_type=HybridSearchIndexType( + chunking_strategy=StaticIndexChunkingStrategy( + max_chunk_size_tokens=700, + chunk_overlap_tokens=300 + ), + # you could also override some text/vector indexes settings + text_search_index=TextSearchIndexType(), + vector_search_index=VectorSearchIndexType(), + normalization_strategy='L2', + # you don't really want to change `k` parameter if you don't + # really know what you are doing + combination_strategy=ReciprocalRankFusionIndexCombinationStrategy( + k=60 + ) + ) + ) + search_index = await operation.wait() + print("new hybrid search index with overridden settings:") + pprint.pprint(search_index) + + # And how to use your index you could learn in example file "assistant_with_search_index.py". + # Working with hybrid index does not differ from working with any other index besides creation. + + # Created resources cleanup: + for file in files: + await file.delete() + + for search_index in [default_search_index, search_index]: + print(f"delete {search_index.id=}") + await search_index.delete() + + +if __name__ == '__main__': + asyncio.run(main()) diff --git a/examples/sync/assistants/hybrid_search_index.py b/examples/sync/assistants/hybrid_search_index.py new file mode 100755 index 0000000..a7871ed --- /dev/null +++ b/examples/sync/assistants/hybrid_search_index.py @@ -0,0 +1,76 @@ +#!/usr/bin/env python3 + +from __future__ import annotations + +import pathlib +import pprint + +from yandex_cloud_ml_sdk import YCloudML +from yandex_cloud_ml_sdk.search_indexes import ( + HybridSearchIndexType, ReciprocalRankFusionIndexCombinationStrategy, StaticIndexChunkingStrategy, + TextSearchIndexType, VectorSearchIndexType +) + + +def local_path(path: str) -> pathlib.Path: + return pathlib.Path(__file__).parent / path + + +def main() -> None: + sdk = YCloudML(folder_id='b1ghsjum2v37c2un8h64') + + files = [] + for path in ['turkey_example.txt', 'maldives_example.txt']: + file = sdk.files.upload( + local_path(path), + ttl_days=5, + expiration_policy="static", + ) + files.append(file) + + # How to create search index with all default settings: + operation = sdk.search_indexes.create_deferred( + files, + index_type=HybridSearchIndexType() + ) + default_search_index = operation.wait() + print("new hybrid search index with default settings:") + pprint.pprint(default_search_index) + + # But you could override any default: + operation = sdk.search_indexes.create_deferred( + files, + index_type=HybridSearchIndexType( + chunking_strategy=StaticIndexChunkingStrategy( + max_chunk_size_tokens=700, + chunk_overlap_tokens=300 + ), + # you could also override some text/vector indexes settings + text_search_index=TextSearchIndexType(), + vector_search_index=VectorSearchIndexType(), + normalization_strategy='L2', + # you don't really want to change `k` parameter if you don't + # really know what you are doing + combination_strategy=ReciprocalRankFusionIndexCombinationStrategy( + k=60 + ) + ) + ) + search_index = operation.wait() + print("new hybrid search index with overridden settings:") + pprint.pprint(search_index) + + # And how to use your index you could learn in example file "assistant_with_search_index.py". + # Working with hybrid index does not differ from working with any other index besides creation. + + # Created resources cleanup: + for file in files: + file.delete() + + for search_index in [default_search_index, search_index]: + print(f"delete {search_index.id=}") + search_index.delete() + + +if __name__ == '__main__': + main() diff --git a/pyproject.toml b/pyproject.toml index d6ad724..1422ac1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -238,6 +238,8 @@ exclude-protected= [ "_to_proto", "_result_type", "_proto_result_type", + "_proto_field_name", + "_coerce", ] valid-classmethod-first-arg="cls" valid-metaclass-classmethod-first-arg="cls" diff --git a/src/yandex_cloud_ml_sdk/_search_indexes/combination_strategy.py b/src/yandex_cloud_ml_sdk/_search_indexes/combination_strategy.py new file mode 100644 index 0000000..92b4132 --- /dev/null +++ b/src/yandex_cloud_ml_sdk/_search_indexes/combination_strategy.py @@ -0,0 +1,109 @@ +# pylint: disable=no-name-in-module,protected-access +from __future__ import annotations + +import abc +import enum +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Collection + +from google.protobuf.wrappers_pb2 import Int64Value +from yandex.cloud.ai.assistants.v1.searchindex.common_pb2 import CombinationStrategy as ProtoCombinationStrategy +from yandex.cloud.ai.assistants.v1.searchindex.common_pb2 import MeanCombinationStrategy as ProtoMeanCombinationStrategy +from yandex.cloud.ai.assistants.v1.searchindex.common_pb2 import ( + ReciprocalRankFusionCombinationStrategy as ProtoReciprocalRankFusionCombinationStrategy +) + +if TYPE_CHECKING: + from yandex_cloud_ml_sdk._sdk import BaseSDK + + +class BaseIndexCombinationStrategy(abc.ABC): + @classmethod + @abc.abstractmethod + def _from_proto(cls, proto: Any, sdk: BaseSDK) -> BaseIndexCombinationStrategy: + pass + + @abc.abstractmethod + def _to_proto(self) -> ProtoCombinationStrategy: + pass + + @classmethod + def _from_upper_proto(cls, proto: ProtoCombinationStrategy, sdk: BaseSDK) -> BaseIndexCombinationStrategy: + if proto.HasField('mean_combination'): + return MeanIndexCombinationStrategy._from_proto( + proto=proto.mean_combination, + sdk=sdk + ) + if proto.HasField('rrf_combination'): + return ReciprocalRankFusionIndexCombinationStrategy._from_proto( + proto=proto.rrf_combination, + sdk=sdk + ) + raise NotImplementedError( + 'combination strategies other then Mean and RRF are not supported in this SDK version' + ) + + +_orig = ProtoMeanCombinationStrategy.MeanEvaluationTechnique + +class MeanIndexEvaluationTechnique(enum.IntEnum): + MEAN_EVALUATION_TECHNIQUE_UNSPECIFIED = _orig.MEAN_EVALUATION_TECHNIQUE_UNSPECIFIED + ARITHMETIC = _orig.ARITHMETIC + GEOMETRIC = _orig.GEOMETRIC + HARMONIC = _orig.HARMONIC + + @classmethod + def _coerce(cls, technique: str | int ) -> MeanIndexEvaluationTechnique: + if isinstance(technique, str): + technique = _orig.Value(technique.upper()) + return cls(technique) + + +@dataclass(frozen=True) +class MeanIndexCombinationStrategy(BaseIndexCombinationStrategy): + mean_evaluation_technique: MeanIndexEvaluationTechnique | None + weights: Collection[float] | None + + @classmethod + # pylint: disable=unused-argument + def _from_proto(cls, proto: ProtoMeanCombinationStrategy, sdk: BaseSDK) -> MeanIndexCombinationStrategy: + return cls( + mean_evaluation_technique=MeanIndexEvaluationTechnique._coerce(proto.mean_evaluation_technique), + weights=tuple(proto.weights) + ) + + def _to_proto(self) -> ProtoCombinationStrategy: + kwargs: dict[str, Any] = {} + if self.mean_evaluation_technique: + kwargs['mean_evaluation_technique'] = int(self.mean_evaluation_technique) + if self.weights is not None: + kwargs['weghts'] = tuple(self.weights) + + return ProtoCombinationStrategy( + mean_combination=ProtoMeanCombinationStrategy(**kwargs) + ) + + +@dataclass(frozen=True) +class ReciprocalRankFusionIndexCombinationStrategy(BaseIndexCombinationStrategy): + k: int | None = None + + @classmethod + # pylint: disable=unused-argument + def _from_proto( + cls, proto: ProtoReciprocalRankFusionCombinationStrategy, sdk: BaseSDK + ) -> ReciprocalRankFusionIndexCombinationStrategy: + kwargs = {} + if proto.HasField('k'): + kwargs['k'] = proto.k.value + return ReciprocalRankFusionIndexCombinationStrategy( + **kwargs + ) + + def _to_proto(self) -> ProtoCombinationStrategy: + kwargs = {} + if self.k is not None: + kwargs['k'] = Int64Value(value=self.k) + return ProtoCombinationStrategy( + rrf_combination=ProtoReciprocalRankFusionCombinationStrategy(**kwargs) + ) diff --git a/src/yandex_cloud_ml_sdk/_search_indexes/domain.py b/src/yandex_cloud_ml_sdk/_search_indexes/domain.py index c39a008..e9cedfa 100644 --- a/src/yandex_cloud_ml_sdk/_search_indexes/domain.py +++ b/src/yandex_cloud_ml_sdk/_search_indexes/domain.py @@ -4,7 +4,6 @@ from typing import AsyncIterator, Generic, Iterator from yandex.cloud.ai.assistants.v1.searchindex.search_index_pb2 import SearchIndex as ProtoSearchIndex -from yandex.cloud.ai.assistants.v1.searchindex.search_index_pb2 import TextSearchIndex, VectorSearchIndex from yandex.cloud.ai.assistants.v1.searchindex.search_index_service_pb2 import ( CreateSearchIndexRequest, GetSearchIndexRequest, ListSearchIndicesRequest, ListSearchIndicesResponse ) @@ -19,7 +18,7 @@ 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 .index_type import BaseSearchIndexType, TextSearchIndexType, VectorSearchIndexType +from .index_type import BaseSearchIndexType from .search_index import AsyncSearchIndex, SearchIndex, SearchIndexTypeT @@ -47,14 +46,11 @@ async def _create_deferred( expiration_config = ExpirationConfig.coerce(ttl_days=ttl_days, expiration_policy=expiration_policy) - vector_search_index: VectorSearchIndex | None = None - text_search_index: TextSearchIndex | None = None - if isinstance(index_type, VectorSearchIndexType): - vector_search_index = index_type._to_proto() - elif isinstance(index_type, TextSearchIndexType): - text_search_index = index_type._to_proto() - elif is_defined(index_type): - raise TypeError('index type must be instance of SearchIndexType') + kwargs = {} + if is_defined(index_type): + if not isinstance(index_type, BaseSearchIndexType): + raise TypeError('index type must be instance of BaseSearchIndexType') + kwargs[index_type._proto_field_name] = index_type._to_proto() request = CreateSearchIndexRequest( folder_id=self._folder_id, @@ -63,8 +59,7 @@ async def _create_deferred( description=get_defined_value(description, ''), labels=get_defined_value(labels, {}), expiration_config=expiration_config.to_proto(), - vector_search_index=vector_search_index, - text_search_index=text_search_index, + **kwargs, # type: ignore[arg-type] ) async with self._client.get_service_stub(SearchIndexServiceStub, timeout=timeout) as stub: diff --git a/src/yandex_cloud_ml_sdk/_search_indexes/index_type.py b/src/yandex_cloud_ml_sdk/_search_indexes/index_type.py index c6f666a..1c9cd11 100644 --- a/src/yandex_cloud_ml_sdk/_search_indexes/index_type.py +++ b/src/yandex_cloud_ml_sdk/_search_indexes/index_type.py @@ -3,27 +3,32 @@ import abc from dataclasses import dataclass -from typing import TYPE_CHECKING, Union +from typing import TYPE_CHECKING, ClassVar, Union +from typing_extensions import Self +from yandex.cloud.ai.assistants.v1.searchindex.search_index_pb2 import HybridSearchIndex from yandex.cloud.ai.assistants.v1.searchindex.search_index_pb2 import SearchIndex as ProtoSearchIndex from yandex.cloud.ai.assistants.v1.searchindex.search_index_pb2 import TextSearchIndex, VectorSearchIndex from .chunking_strategy import BaseIndexChunkingStrategy +from .combination_strategy import BaseIndexCombinationStrategy +from .normalization_strategy import IndexNormalizationStrategy if TYPE_CHECKING: from yandex_cloud_ml_sdk._sdk import BaseSDK -ProtoSearchIndexType = Union[TextSearchIndex, VectorSearchIndex] +ProtoSearchIndexType = Union[TextSearchIndex, VectorSearchIndex, HybridSearchIndex] @dataclass(frozen=True) class BaseSearchIndexType(abc.ABC): + _proto_field_name: ClassVar[str] chunking_strategy: BaseIndexChunkingStrategy | None = None @classmethod @abc.abstractmethod - def _from_proto(cls, proto, sdk: BaseSDK) -> BaseSearchIndexType: + def _from_proto(cls, proto, sdk: BaseSDK) -> Self: pass @abc.abstractmethod @@ -31,30 +36,40 @@ def _to_proto(self) -> ProtoSearchIndexType: pass @classmethod - def _from_upper_proto(cls, proto: ProtoSearchIndex, sdk: BaseSDK) -> BaseSearchIndexType: - if proto.HasField('text_search_index'): - return TextSearchIndexType._from_proto( - proto=proto.text_search_index, - sdk=sdk - ) - if proto.HasField('vector_search_index'): - return VectorSearchIndexType._from_proto( - proto=proto.vector_search_index, - sdk=sdk - ) + def _parse_chunking_strategy(cls, proto: ProtoSearchIndexType, sdk: BaseSDK) -> BaseIndexChunkingStrategy | None: + if proto.HasField('chunking_strategy'): + return BaseIndexChunkingStrategy._from_upper_proto(proto=proto.chunking_strategy, sdk=sdk) + return None - raise NotImplementedError('search index types other then text&vector are not supported in this SDK version') + @classmethod + def _from_upper_proto(cls, proto: ProtoSearchIndex, sdk: BaseSDK) -> BaseSearchIndexType: + klasses = ( + TextSearchIndexType, + VectorSearchIndexType, + HybridSearchIndexType, + ) + field_klasses = {klass._proto_field_name: klass for klass in klasses} + # TODO: registering metaclass? + for field, klass in field_klasses.items(): + if proto.HasField(field): # type: ignore[arg-type] + value = getattr(proto, field) + return klass._from_proto( + proto=value, + sdk=sdk + ) + raise NotImplementedError( + f'search index types other than {list(field_klasses)} are not supported in this SDK version' + ) @dataclass(frozen=True) class TextSearchIndexType(BaseSearchIndexType): + _proto_field_name: ClassVar[str] = 'text_search_index' + @classmethod def _from_proto(cls, proto: TextSearchIndex, sdk: BaseSDK) -> TextSearchIndexType: return cls( - chunking_strategy=BaseIndexChunkingStrategy._from_upper_proto( - proto=proto.chunking_strategy, - sdk=sdk, - ) + chunking_strategy=cls._parse_chunking_strategy(proto, sdk), ) def _to_proto(self) -> TextSearchIndex: @@ -66,6 +81,8 @@ def _to_proto(self) -> TextSearchIndex: @dataclass(frozen=True) class VectorSearchIndexType(BaseSearchIndexType): + _proto_field_name: ClassVar[str] = 'vector_search_index' + doc_embedder_uri: str | None = None query_embedder_uri: str | None = None @@ -75,10 +92,7 @@ def _from_proto(cls, proto: VectorSearchIndex, sdk: BaseSDK) -> VectorSearchInde return cls( doc_embedder_uri=proto.doc_embedder_uri, query_embedder_uri=proto.query_embedder_uri, - chunking_strategy=BaseIndexChunkingStrategy._from_upper_proto( - proto=proto.chunking_strategy, - sdk=sdk, - ) + chunking_strategy=cls._parse_chunking_strategy(proto, sdk), ) def _to_proto(self) -> VectorSearchIndex: @@ -88,3 +102,48 @@ def _to_proto(self) -> VectorSearchIndex: doc_embedder_uri=self.doc_embedder_uri or '', query_embedder_uri=self.query_embedder_uri or '', ) + + +@dataclass(frozen=True) +class HybridSearchIndexType(BaseSearchIndexType): + _proto_field_name: ClassVar[str] = 'hybrid_search_index' + + text_search_index: TextSearchIndexType | None = None + vector_search_index: VectorSearchIndexType | None = None + normalization_strategy: IndexNormalizationStrategy | str | int | None = None + combination_strategy: BaseIndexCombinationStrategy | None = None + + @classmethod + def _from_proto(cls, proto: HybridSearchIndex, sdk: BaseSDK) -> HybridSearchIndexType: + return cls( + chunking_strategy=cls._parse_chunking_strategy(proto, sdk), + text_search_index=TextSearchIndexType._from_proto( + proto=proto.text_search_index, + sdk=sdk, + ), + vector_search_index=VectorSearchIndexType._from_proto( + proto=proto.vector_search_index, + sdk=sdk, + ), + normalization_strategy=IndexNormalizationStrategy._coerce(proto.normalization_strategy), + combination_strategy=BaseIndexCombinationStrategy._from_upper_proto( + proto.combination_strategy, sdk=sdk + ) + ) + + def _to_proto(self) -> HybridSearchIndex: + chunking_strategy = self.chunking_strategy._to_proto() if self.chunking_strategy else None + text_search_index = self.text_search_index._to_proto() if self.text_search_index else None + vector_search_index = self.vector_search_index._to_proto() if self.vector_search_index else None + normalization_strategy = IndexNormalizationStrategy._coerce( + self.normalization_strategy + ) if self.normalization_strategy else 0 + combination_strategy = self.combination_strategy._to_proto() if self.combination_strategy else None + + return HybridSearchIndex( + chunking_strategy=chunking_strategy, + text_search_index=text_search_index, + vector_search_index=vector_search_index, + normalization_strategy=normalization_strategy, # type: ignore[arg-type] + combination_strategy=combination_strategy, + ) diff --git a/src/yandex_cloud_ml_sdk/_search_indexes/normalization_strategy.py b/src/yandex_cloud_ml_sdk/_search_indexes/normalization_strategy.py new file mode 100644 index 0000000..04362a2 --- /dev/null +++ b/src/yandex_cloud_ml_sdk/_search_indexes/normalization_strategy.py @@ -0,0 +1,18 @@ +# pylint: disable=no-name-in-module,protected-access +from __future__ import annotations + +import enum + +from yandex.cloud.ai.assistants.v1.searchindex.common_pb2 import NormalizationStrategy + + +class IndexNormalizationStrategy(enum.IntEnum): + NORMALIZATION_STRATEGY_UNSPECIFIED = NormalizationStrategy.NORMALIZATION_STRATEGY_UNSPECIFIED + MIN_MAX = NormalizationStrategy.MIN_MAX + L2 = NormalizationStrategy.L2 + + @classmethod + def _coerce(cls, strategy: str | int | IndexNormalizationStrategy) -> IndexNormalizationStrategy: + if isinstance(strategy, str): + strategy = NormalizationStrategy.Value(strategy.upper()) + return cls(strategy) diff --git a/src/yandex_cloud_ml_sdk/search_indexes.py b/src/yandex_cloud_ml_sdk/search_indexes.py index 197d284..4b577c5 100644 --- a/src/yandex_cloud_ml_sdk/search_indexes.py +++ b/src/yandex_cloud_ml_sdk/search_indexes.py @@ -1,10 +1,19 @@ from __future__ import annotations from ._search_indexes.chunking_strategy import StaticIndexChunkingStrategy -from ._search_indexes.index_type import TextSearchIndexType, VectorSearchIndexType +from ._search_indexes.combination_strategy import ( + MeanIndexCombinationStrategy, MeanIndexEvaluationTechnique, ReciprocalRankFusionIndexCombinationStrategy +) +from ._search_indexes.index_type import HybridSearchIndexType, TextSearchIndexType, VectorSearchIndexType +from ._search_indexes.normalization_strategy import IndexNormalizationStrategy __all__ = [ + 'IndexNormalizationStrategy', + 'HybridSearchIndexType', + 'MeanIndexCombinationStrategy', + 'MeanIndexEvaluationTechnique', + 'ReciprocalRankFusionIndexCombinationStrategy', 'StaticIndexChunkingStrategy', + 'TextSearchIndexType', 'VectorSearchIndexType', - 'TextSearchIndexType' ] diff --git a/tests/assistants/cassettes/test_search_indexes/test_hybrid_search_index.gprc.json b/tests/assistants/cassettes/test_search_indexes/test_hybrid_search_index.gprc.json new file mode 100644 index 0000000..f2d1ffc --- /dev/null +++ b/tests/assistants/cassettes/test_search_indexes/test_hybrid_search_index.gprc.json @@ -0,0 +1,550 @@ +{ + "interactions": [ + { + "request": { + "cls": "ListApiEndpointsRequest", + "module": "yandex.cloud.endpoint.api_endpoint_service_pb2", + "message": {} + }, + "response": { + "cls": "ListApiEndpointsResponse", + "module": "yandex.cloud.endpoint.api_endpoint_service_pb2", + "message": { + "endpoints": [ + { + "id": "ai-assistants", + "address": "assistant.api.cloud.yandex.net:443" + }, + { + "id": "ai-files", + "address": "assistant.api.cloud.yandex.net:443" + }, + { + "id": "ai-foundation-models", + "address": "llm.api.cloud.yandex.net:443" + }, + { + "id": "ai-llm", + "address": "llm.api.cloud.yandex.net:443" + }, + { + "id": "ai-speechkit", + "address": "transcribe.api.cloud.yandex.net:443" + }, + { + "id": "ai-stt", + "address": "transcribe.api.cloud.yandex.net:443" + }, + { + "id": "ai-stt-v3", + "address": "stt.api.cloud.yandex.net:443" + }, + { + "id": "ai-translate", + "address": "translate.api.cloud.yandex.net:443" + }, + { + "id": "ai-vision", + "address": "vision.api.cloud.yandex.net:443" + }, + { + "id": "ai-vision-ocr", + "address": "ocr.api.cloud.yandex.net:443" + }, + { + "id": "alb", + "address": "alb.api.cloud.yandex.net:443" + }, + { + "id": "apigateway-connections", + "address": "apigateway-connections.api.cloud.yandex.net:443" + }, + { + "id": "application-load-balancer", + "address": "alb.api.cloud.yandex.net:443" + }, + { + "id": "apploadbalancer", + "address": "alb.api.cloud.yandex.net:443" + }, + { + "id": "audittrails", + "address": "audittrails.api.cloud.yandex.net:443" + }, + { + "id": "baas", + "address": "backup.api.cloud.yandex.net:443" + }, + { + "id": "backup", + "address": "backup.api.cloud.yandex.net:443" + }, + { + "id": "billing", + "address": "billing.api.cloud.yandex.net:443" + }, + { + "id": "broker-data", + "address": "iot-data.api.cloud.yandex.net:443" + }, + { + "id": "cdn", + "address": "cdn.api.cloud.yandex.net:443" + }, + { + "id": "certificate-manager", + "address": "certificate-manager.api.cloud.yandex.net:443" + }, + { + "id": "certificate-manager-data", + "address": "data.certificate-manager.api.cloud.yandex.net:443" + }, + { + "id": "cic", + "address": "cic-api.api.cloud.yandex.net:443" + }, + { + "id": "cloud-registry", + "address": "registry.api.cloud.yandex.net:443" + }, + { + "id": "cloudapps", + "address": "cloudapps.api.cloud.yandex.net:443" + }, + { + "id": "cloudbackup", + "address": "backup.api.cloud.yandex.net:443" + }, + { + "id": "clouddesktops", + "address": "clouddesktops.api.cloud.yandex.net:443" + }, + { + "id": "cloudrouter", + "address": "cic-api.api.cloud.yandex.net:443" + }, + { + "id": "cloudvideo", + "address": "video.api.cloud.yandex.net:443" + }, + { + "id": "compute", + "address": "compute.api.cloud.yandex.net:443" + }, + { + "id": "container-registry", + "address": "container-registry.api.cloud.yandex.net:443" + }, + { + "id": "dataproc", + "address": "dataproc.api.cloud.yandex.net:443" + }, + { + "id": "dataproc-manager", + "address": "dataproc-manager.api.cloud.yandex.net:443" + }, + { + "id": "datasphere", + "address": "datasphere.api.cloud.yandex.net:443" + }, + { + "id": "datatransfer", + "address": "datatransfer.api.cloud.yandex.net:443" + }, + { + "id": "dns", + "address": "dns.api.cloud.yandex.net:443" + }, + { + "id": "endpoint", + "address": "api.cloud.yandex.net:443" + }, + { + "id": "fomo-dataset", + "address": "fomo-dataset.api.cloud.yandex.net:443" + }, + { + "id": "fomo-tuning", + "address": "fomo-tuning.api.cloud.yandex.net:443" + }, + { + "id": "iam", + "address": "iam.api.cloud.yandex.net:443" + }, + { + "id": "iot-broker", + "address": "iot-broker.api.cloud.yandex.net:443" + }, + { + "id": "iot-data", + "address": "iot-data.api.cloud.yandex.net:443" + }, + { + "id": "iot-devices", + "address": "iot-devices.api.cloud.yandex.net:443" + }, + { + "id": "k8s", + "address": "mks.api.cloud.yandex.net:443" + }, + { + "id": "kms", + "address": "kms.api.cloud.yandex.net:443" + }, + { + "id": "kms-crypto", + "address": "kms.yandex:443" + }, + { + "id": "load-balancer", + "address": "load-balancer.api.cloud.yandex.net:443" + }, + { + "id": "loadtesting", + "address": "loadtesting.api.cloud.yandex.net:443" + }, + { + "id": "locator", + "address": "locator.api.cloud.yandex.net:443" + }, + { + "id": "lockbox", + "address": "lockbox.api.cloud.yandex.net:443" + }, + { + "id": "lockbox-payload", + "address": "payload.lockbox.api.cloud.yandex.net:443" + }, + { + "id": "log-ingestion", + "address": "ingester.logging.yandexcloud.net:443" + }, + { + "id": "log-reading", + "address": "reader.logging.yandexcloud.net:443" + }, + { + "id": "logging", + "address": "logging.api.cloud.yandex.net:443" + }, + { + "id": "managed-airflow", + "address": "airflow.api.cloud.yandex.net:443" + }, + { + "id": "managed-clickhouse", + "address": "mdb.api.cloud.yandex.net:443" + }, + { + "id": "managed-elasticsearch", + "address": "mdb.api.cloud.yandex.net:443" + }, + { + "id": "managed-greenplum", + "address": "mdb.api.cloud.yandex.net:443" + }, + { + "id": "managed-kafka", + "address": "mdb.api.cloud.yandex.net:443" + }, + { + "id": "managed-kubernetes", + "address": "mks.api.cloud.yandex.net:443" + }, + { + "id": "managed-mongodb", + "address": "mdb.api.cloud.yandex.net:443" + }, + { + "id": "managed-mysql", + "address": "mdb.api.cloud.yandex.net:443" + }, + { + "id": "managed-opensearch", + "address": "mdb.api.cloud.yandex.net:443" + }, + { + "id": "managed-postgresql", + "address": "mdb.api.cloud.yandex.net:443" + }, + { + "id": "managed-redis", + "address": "mdb.api.cloud.yandex.net:443" + }, + { + "id": "managed-sqlserver", + "address": "mdb.api.cloud.yandex.net:443" + }, + { + "id": "marketplace", + "address": "marketplace.api.cloud.yandex.net:443" + }, + { + "id": "marketplace-pim", + "address": "mkt.private-api.cloud.yandex.net:4446" + }, + { + "id": "mdb-clickhouse", + "address": "mdb.api.cloud.yandex.net:443" + }, + { + "id": "mdb-mongodb", + "address": "mdb.api.cloud.yandex.net:443" + }, + { + "id": "mdb-mysql", + "address": "mdb.api.cloud.yandex.net:443" + }, + { + "id": "mdb-opensearch", + "address": "mdb.api.cloud.yandex.net:443" + }, + { + "id": "mdb-postgresql", + "address": "mdb.api.cloud.yandex.net:443" + }, + { + "id": "mdb-redis", + "address": "mdb.api.cloud.yandex.net:443" + }, + { + "id": "mdbproxy", + "address": "mdbproxy.api.cloud.yandex.net:443" + }, + { + "id": "monitoring", + "address": "monitoring.api.cloud.yandex.net:443" + }, + { + "id": "operation", + "address": "operation.api.cloud.yandex.net:443" + }, + { + "id": "organization-manager", + "address": "organization-manager.api.cloud.yandex.net:443" + }, + { + "id": "organizationmanager", + "address": "organization-manager.api.cloud.yandex.net:443" + }, + { + "id": "resource-manager", + "address": "resource-manager.api.cloud.yandex.net:443" + }, + { + "id": "resourcemanager", + "address": "resource-manager.api.cloud.yandex.net:443" + }, + { + "id": "searchapi", + "address": "searchapi.api.cloud.yandex.net:443" + }, + { + "id": "serialssh", + "address": "serialssh.cloud.yandex.net:9600" + }, + { + "id": "serverless-apigateway", + "address": "serverless-apigateway.api.cloud.yandex.net:443" + }, + { + "id": "serverless-containers", + "address": "serverless-containers.api.cloud.yandex.net:443" + }, + { + "id": "serverless-eventrouter", + "address": "serverless-eventrouter.api.cloud.yandex.net:443" + }, + { + "id": "serverless-functions", + "address": "serverless-functions.api.cloud.yandex.net:443" + }, + { + "id": "serverless-gateway-connections", + "address": "apigateway-connections.api.cloud.yandex.net:443" + }, + { + "id": "serverless-triggers", + "address": "serverless-triggers.api.cloud.yandex.net:443" + }, + { + "id": "serverless-workflows", + "address": "serverless-workflows.api.cloud.yandex.net:443" + }, + { + "id": "serverlesseventrouter-events", + "address": "events.eventrouter.serverless.yandexcloud.net:443" + }, + { + "id": "smart-captcha", + "address": "smartcaptcha.api.cloud.yandex.net:443" + }, + { + "id": "smart-web-security", + "address": "smartwebsecurity.api.cloud.yandex.net:443" + }, + { + "id": "storage", + "address": "storage.yandexcloud.net:443" + }, + { + "id": "storage-api", + "address": "storage.api.cloud.yandex.net:443" + }, + { + "id": "video", + "address": "video.api.cloud.yandex.net:443" + }, + { + "id": "vpc", + "address": "vpc.api.cloud.yandex.net:443" + }, + { + "id": "ydb", + "address": "ydb.api.cloud.yandex.net:443" + } + ] + } + } + }, + { + "request": { + "cls": "CreateFileRequest", + "module": "yandex.cloud.ai.files.v1.file_service_pb2", + "message": { + "folderId": "b1ghsjum2v37c2un8h64", + "content": "dGVzdCBmaWxl" + } + }, + "response": { + "cls": "File", + "module": "yandex.cloud.ai.files.v1.file_pb2", + "message": { + "id": "fvtppj63ddll9e6u7ncj", + "folderId": "b1ghsjum2v37c2un8h64", + "mimeType": "text/plain", + "createdBy": "aje6euqn63oa635coh28", + "createdAt": "2024-12-18T15:14:47.538885Z", + "updatedBy": "aje6euqn63oa635coh28", + "updatedAt": "2024-12-18T15:14:47.538885Z", + "expirationConfig": { + "expirationPolicy": "SINCE_LAST_ACTIVE", + "ttlDays": "7" + }, + "expiresAt": "2024-12-25T15:14:47.538885Z" + } + } + }, + { + "request": { + "cls": "CreateSearchIndexRequest", + "module": "yandex.cloud.ai.assistants.v1.searchindex.search_index_service_pb2", + "message": { + "folderId": "b1ghsjum2v37c2un8h64", + "fileIds": [ + "fvtppj63ddll9e6u7ncj" + ], + "hybridSearchIndex": { + "chunkingStrategy": { + "staticStrategy": { + "maxChunkSizeTokens": "700", + "chunkOverlapTokens": "300" + } + }, + "normalizationStrategy": "L2", + "combinationStrategy": { + "rrfCombination": { + "k": "51" + } + } + } + } + }, + "response": { + "cls": "Operation", + "module": "yandex.cloud.operation.operation_pb2", + "message": { + "id": "fvtb3tf696qo91f6qavq", + "description": "search index creation", + "createdAt": "2024-12-18T15:14:47.869730Z", + "createdBy": "aje6euqn63oa635coh28", + "modifiedAt": "2024-12-18T15:14:47.869730Z" + } + } + }, + { + "request": { + "cls": "GetOperationRequest", + "module": "yandex.cloud.operation.operation_service_pb2", + "message": { + "operationId": "fvtb3tf696qo91f6qavq" + } + }, + "response": { + "cls": "Operation", + "module": "yandex.cloud.operation.operation_pb2", + "message": { + "id": "fvtb3tf696qo91f6qavq", + "description": "search index creation", + "createdAt": "2024-12-18T15:14:47.869730Z", + "createdBy": "aje6euqn63oa635coh28", + "modifiedAt": "2024-12-18T15:14:47.869730Z" + } + } + }, + { + "request": { + "cls": "GetOperationRequest", + "module": "yandex.cloud.operation.operation_service_pb2", + "message": { + "operationId": "fvtb3tf696qo91f6qavq" + } + }, + "response": { + "cls": "Operation", + "module": "yandex.cloud.operation.operation_pb2", + "message": { + "id": "fvtb3tf696qo91f6qavq", + "description": "search index creation", + "createdAt": "2024-12-18T15:14:47.869730Z", + "createdBy": "aje6euqn63oa635coh28", + "modifiedAt": "2024-12-18T15:14:49.925279Z", + "done": true, + "response": { + "@type": "type.googleapis.com/yandex.cloud.ai.assistants.v1.searchindex.SearchIndex", + "id": "fvt2dcunqph3m3sq0lsq", + "folderId": "b1ghsjum2v37c2un8h64", + "createdBy": "aje6euqn63oa635coh28", + "createdAt": "2024-12-18T15:14:47.802010Z", + "updatedBy": "aje6euqn63oa635coh28", + "updatedAt": "2024-12-18T15:14:47.802010Z", + "expirationConfig": { + "expirationPolicy": "SINCE_LAST_ACTIVE", + "ttlDays": "7" + }, + "expiresAt": "2024-12-25T15:14:47.802010Z", + "hybridSearchIndex": { + "textSearchIndex": {}, + "vectorSearchIndex": { + "docEmbedderUri": "emb://yc.ml.rag-prod.common/text-search-doc/latest", + "queryEmbedderUri": "emb://yc.ml.rag-prod.common/text-search-query/latest", + "chunkingStrategy": { + "staticStrategy": { + "maxChunkSizeTokens": "700", + "chunkOverlapTokens": "300" + } + } + }, + "normalizationStrategy": "L2", + "combinationStrategy": { + "rrfCombination": { + "k": "51" + } + } + } + } + } + } + } + ] +} diff --git a/tests/assistants/test_search_indexes.py b/tests/assistants/test_search_indexes.py index bb786be..35ff4c5 100644 --- a/tests/assistants/test_search_indexes.py +++ b/tests/assistants/test_search_indexes.py @@ -2,7 +2,11 @@ import pytest -from yandex_cloud_ml_sdk.search_indexes import StaticIndexChunkingStrategy, TextSearchIndexType, VectorSearchIndexType +from yandex_cloud_ml_sdk import AsyncYCloudML +from yandex_cloud_ml_sdk.search_indexes import ( + HybridSearchIndexType, IndexNormalizationStrategy, ReciprocalRankFusionIndexCombinationStrategy, + StaticIndexChunkingStrategy, TextSearchIndexType, VectorSearchIndexType +) pytestmark = pytest.mark.asyncio @@ -118,3 +122,29 @@ async def test_assistant_with_search_index(async_sdk, tmp_path): await thread.delete() await assistant.delete() await file.delete() + + +@pytest.mark.allow_grpc +async def test_hybrid_search_index(async_sdk: AsyncYCloudML, test_file_path): + file = await async_sdk.files.upload(test_file_path) + operation = await async_sdk.search_indexes.create_deferred( + file, + index_type=HybridSearchIndexType( + chunking_strategy=StaticIndexChunkingStrategy( + max_chunk_size_tokens=700, + chunk_overlap_tokens=300 + ), + normalization_strategy='L2', + combination_strategy=ReciprocalRankFusionIndexCombinationStrategy(k=51) + ) + ) + search_index = await operation.wait() + + assert isinstance(search_index.index_type, HybridSearchIndexType) + assert isinstance(search_index.index_type.text_search_index, TextSearchIndexType) + assert isinstance(search_index.index_type.vector_search_index, VectorSearchIndexType) + assert isinstance(search_index.index_type.vector_search_index.chunking_strategy, StaticIndexChunkingStrategy) + assert search_index.index_type.vector_search_index.chunking_strategy.max_chunk_size_tokens == 700 + assert search_index.index_type.normalization_strategy == IndexNormalizationStrategy.L2 + assert isinstance(search_index.index_type.combination_strategy, ReciprocalRankFusionIndexCombinationStrategy) + assert search_index.index_type.combination_strategy.k == 51