# Copyright (c) Microsoft. All rights reserved. import ast import asyncio import logging import sys from collections.abc import Sequence from typing import Any, ClassVar, Final, Generic, TypeVar from azure.core.credentials import AzureKeyCredential from azure.core.credentials_async import AsyncTokenCredential from azure.search.documents.aio import SearchClient from azure.search.documents.indexes.aio import SearchIndexClient from azure.search.documents.indexes.models import ( ExhaustiveKnnAlgorithmConfiguration, ExhaustiveKnnParameters, HnswAlgorithmConfiguration, HnswParameters, SearchField, SearchFieldDataType, SearchIndex, SearchResourceEncryptionKey, SimpleField, VectorSearchAlgorithmMetric, VectorSearchProfile, ) from azure.search.documents.indexes.models import VectorSearch as AZSVectorSearch from azure.search.documents.models import VectorizableTextQuery, VectorizedQuery from pydantic import SecretStr, ValidationError from semantic_kernel.connectors.ai.embedding_generator_base import EmbeddingGeneratorBase from semantic_kernel.data.vector import ( DistanceFunction, FieldTypes, GetFilteredRecordOptions, IndexKind, KernelSearchResults, SearchType, TModel, VectorSearch, VectorSearchOptions, VectorSearchResult, VectorStore, VectorStoreCollection, VectorStoreCollectionDefinition, _get_collection_name_from_model, ) from semantic_kernel.exceptions import ( ServiceInitializationError, VectorSearchExecutionException, VectorStoreInitializationException, VectorStoreOperationException, ) from semantic_kernel.kernel_pydantic import HttpsUrl, KernelBaseSettings from semantic_kernel.utils.feature_stage_decorator import release_candidate from semantic_kernel.utils.telemetry.user_agent import APP_INFO, prepend_semantic_kernel_to_user_agent if sys.version_info >= (3, 12): from typing import override # pragma: no cover else: from typing_extensions import override # pragma: no cover logger: logging.Logger = logging.getLogger(__name__) TKey = TypeVar("TKey", bound=str) INDEX_ALGORITHM_MAP: Final[dict[IndexKind, tuple[type, type]]] = { IndexKind.HNSW: (HnswAlgorithmConfiguration, HnswParameters), IndexKind.FLAT: (ExhaustiveKnnAlgorithmConfiguration, ExhaustiveKnnParameters), IndexKind.DEFAULT: (HnswAlgorithmConfiguration, HnswParameters), } DISTANCE_FUNCTION_MAP: Final[dict[DistanceFunction, VectorSearchAlgorithmMetric]] = { DistanceFunction.COSINE_DISTANCE: VectorSearchAlgorithmMetric.COSINE, DistanceFunction.DOT_PROD: VectorSearchAlgorithmMetric.DOT_PRODUCT, DistanceFunction.EUCLIDEAN_DISTANCE: VectorSearchAlgorithmMetric.EUCLIDEAN, DistanceFunction.HAMMING: VectorSearchAlgorithmMetric.HAMMING, DistanceFunction.DEFAULT: VectorSearchAlgorithmMetric.COSINE, } TYPE_MAP_DATA: Final[dict[str, str]] = { "default": SearchFieldDataType.String, "str": SearchFieldDataType.String, "int": SearchFieldDataType.Int64, "float": SearchFieldDataType.Double, "bool": SearchFieldDataType.Boolean, "list[str]": SearchFieldDataType.Collection(SearchFieldDataType.String), "list[int]": SearchFieldDataType.Collection(SearchFieldDataType.Int64), "list[float]": SearchFieldDataType.Collection(SearchFieldDataType.Double), "list[bool]": SearchFieldDataType.Collection(SearchFieldDataType.Boolean), "list[dict]": SearchFieldDataType.Collection(SearchFieldDataType.ComplexType), "dict": SearchFieldDataType.ComplexType, SearchFieldDataType.ComplexType: SearchFieldDataType.ComplexType, SearchFieldDataType.String: SearchFieldDataType.String, SearchFieldDataType.Int64: SearchFieldDataType.Int64, SearchFieldDataType.Double: SearchFieldDataType.Double, SearchFieldDataType.Boolean: SearchFieldDataType.Boolean, SearchFieldDataType.Collection(SearchFieldDataType.String): SearchFieldDataType.Collection( SearchFieldDataType.String ), SearchFieldDataType.Collection(SearchFieldDataType.Int64): SearchFieldDataType.Collection( SearchFieldDataType.Int64 ), SearchFieldDataType.Collection(SearchFieldDataType.Double): SearchFieldDataType.Collection( SearchFieldDataType.Double ), SearchFieldDataType.Collection(SearchFieldDataType.Boolean): SearchFieldDataType.Collection( SearchFieldDataType.Boolean ), SearchFieldDataType.Collection(SearchFieldDataType.ComplexType): SearchFieldDataType.Collection( SearchFieldDataType.ComplexType ), SearchFieldDataType.Collection(SearchFieldDataType.Single): SearchFieldDataType.Collection( SearchFieldDataType.Single ), SearchFieldDataType.DateTimeOffset: SearchFieldDataType.DateTimeOffset, SearchFieldDataType.GeographyPoint: SearchFieldDataType.GeographyPoint, SearchFieldDataType.Single: SearchFieldDataType.Single, } TYPE_MAP_VECTOR: Final[dict[str, str]] = { "float": SearchFieldDataType.Collection(SearchFieldDataType.Single), "int": "Collection(Edm.Int16)", "binary": "Collection(Edm.Byte)", "default": SearchFieldDataType.Collection(SearchFieldDataType.Single), } __all__ = [ "AzureAISearchCollection", "AzureAISearchSettings", "AzureAISearchStore", ] @release_candidate class AzureAISearchSettings(KernelBaseSettings): """Azure AI Search model settings currently used by the AzureCognitiveSearchMemoryStore connector. Args: - api_key: SecretStr - Azure AI Search API key (Env var AZURE_AI_SEARCH_API_KEY) - endpoint: HttpsUrl - Azure AI Search endpoint (Env var AZURE_AI_SEARCH_ENDPOINT) - index_name: str - Azure AI Search index name (Env var AZURE_AI_SEARCH_INDEX_NAME) """ env_prefix: ClassVar[str] = "AZURE_AI_SEARCH_" api_key: SecretStr | None = None endpoint: HttpsUrl index_name: str | None = None def _get_search_client( endpoint: str, collection_name: str | None, credential: "AzureKeyCredential | AsyncTokenCredential", **kwargs: Any, ) -> SearchClient: """Create a search client for a collection.""" if not collection_name: raise VectorStoreInitializationException("Collection name is required to create a search client.") try: return SearchClient(endpoint, collection_name, credential, **kwargs) except ValueError as exc: raise VectorStoreInitializationException( f"Failed to create Azure Cognitive Search client for collection {collection_name}." ) from exc def _resolve_credential( azure_ai_search_settings: AzureAISearchSettings, azure_credential: AzureKeyCredential | None = None, token_credential: "AsyncTokenCredential | None" = None, ) -> "AzureKeyCredential | AsyncTokenCredential": """Resolve the credential to use for Azure AI Search. Args: azure_ai_search_settings: Azure AI Search settings. azure_credential: Optional Azure credentials (default: {None}). token_credential: Optional Token credential (default: {None}). """ if azure_credential: return azure_credential if token_credential: return token_credential if azure_ai_search_settings.api_key: return AzureKeyCredential(azure_ai_search_settings.api_key.get_secret_value()) raise ServiceInitializationError("Error: missing Azure AI Search client credentials.") def _get_search_index_client( azure_ai_search_settings: AzureAISearchSettings, azure_credential: AzureKeyCredential | None = None, token_credential: "AsyncTokenCredential | None" = None, ) -> SearchIndexClient: """Return a client for Azure AI Search. Args: azure_ai_search_settings: Azure AI Search settings. azure_credential: Optional Azure credentials (default: {None}). token_credential: Optional Token credential (default: {None}). """ credential = _resolve_credential(azure_ai_search_settings, azure_credential, token_credential) return SearchIndexClient( endpoint=str(azure_ai_search_settings.endpoint), credential=credential, headers=prepend_semantic_kernel_to_user_agent({}) if APP_INFO else None, ) def _definition_to_azure_ai_search_index( collection_name: str, definition: VectorStoreCollectionDefinition, encryption_key: SearchResourceEncryptionKey | None = None, ) -> SearchIndex: """Convert a VectorStoreRecordDefinition to an Azure AI Search index.""" fields = [] search_profiles = [] search_algos = [] for field in definition.fields: if field.field_type == FieldTypes.DATA: if not field.type_: logger.debug(f"Field {field.name} has not specified type, defaulting to Edm.String.") if field.type_ and field.type_ not in TYPE_MAP_DATA: if field.type_.startswith("dict"): type_ = TYPE_MAP_DATA["dict"] elif field.type_.startswith("list") and "dict" in field.type_: type_ = TYPE_MAP_DATA["list[dict]"] else: raise VectorStoreOperationException(f"{field.type_} not supported in Azure AI Search.") else: type_ = TYPE_MAP_DATA[field.type_ or "default"] fields.append( SearchField( name=field.storage_name or field.name, type=type_, filterable=field.is_indexed or field.is_full_text_indexed, # searchable is set first on the value of is_full_text_searchable, # if it is None it checks the field type, if text then it is searchable searchable=type_ in ("Edm.String", "Collection(Edm.String)") if field.is_full_text_indexed is None else field.is_full_text_indexed, sortable=not type_.startswith("Collection") or type_ == "Edm.ComplexType", hidden=False, ) ) elif field.field_type == FieldTypes.KEY: fields.append( SimpleField( name=field.storage_name or field.name, type="Edm.String", # hardcoded, only allowed type for key key=True, filterable=True, searchable=True, ) ) elif field.field_type == FieldTypes.VECTOR: if not field.type_: logger.debug(f"Field {field.name} has not specified type, defaulting to Collection(Edm.Single).") if field.index_kind not in INDEX_ALGORITHM_MAP: raise VectorStoreOperationException(f"{field.index_kind} not supported in Azure AI Search.") if field.distance_function not in DISTANCE_FUNCTION_MAP: raise VectorStoreOperationException(f"{field.distance_function} not supported in Azure AI Search.") profile_name = f"{field.storage_name or field.name}_profile" algo_name = f"{field.storage_name or field.name}_algorithm" fields.append( SearchField( name=field.storage_name or field.name, type=TYPE_MAP_VECTOR[field.type_ or "default"], searchable=True, vector_search_dimensions=field.dimensions, vector_search_profile_name=profile_name, hidden=False, ) ) search_profiles.append( VectorSearchProfile( name=profile_name, algorithm_configuration_name=algo_name, ) ) algo_class, algo_params = INDEX_ALGORITHM_MAP[field.index_kind] distance_metric = DISTANCE_FUNCTION_MAP[field.distance_function] search_algos.append(algo_class(name=algo_name, parameters=algo_params(metric=distance_metric))) return SearchIndex( name=collection_name, fields=fields, vector_search=AZSVectorSearch(profiles=search_profiles, algorithms=search_algos), encryption_key=encryption_key, ) @release_candidate class AzureAISearchCollection( VectorStoreCollection[TKey, TModel], VectorSearch[TKey, TModel], Generic[TKey, TModel], ): """Azure AI Search collection implementation.""" search_client: SearchClient search_index_client: SearchIndexClient search_endpoint: str | None = None search_credential: Any = None supported_key_types: ClassVar[set[str] | None] = {"str"} supported_vector_types: ClassVar[set[str] | None] = {"float", "int"} supported_search_types: ClassVar[set[SearchType]] = {SearchType.VECTOR, SearchType.KEYWORD_HYBRID} managed_search_index_client: bool = True def __init__( self, record_type: type[TModel], definition: VectorStoreCollectionDefinition | None = None, collection_name: str | None = None, search_index_client: SearchIndexClient | None = None, search_client: SearchClient | None = None, embedding_generator: "EmbeddingGeneratorBase | None" = None, search_credential: "AzureKeyCredential | AsyncTokenCredential | None" = None, **kwargs: Any, ) -> None: """Initializes a new instance of the AzureAISearchCollection class. The collection name can be set in four ways: 1. By passing it in the constructor. 2. By passing it in the data model definition or record_type. 3. By passing it in the search client. 4. By setting the AZURE_AI_SEARCH_INDEX_NAME environment variable. They are checked in that order, so if the collection name is passed in the constructor it is used. Args: record_type: The type of the data model. definition: The model definition, optional. collection_name: The name of the collection, optional. search_index_client: The search index client for interacting with Azure AI Search, used for creating and deleting indexes. search_client: The search client for interacting with Azure AI Search, used for record operations. search_credential: The credential used to authenticate with Azure AI Search. If not provided, it will be resolved from azure_credentials, token_credentials, or api_key in kwargs/environment. embedding_generator: The embedding generator, optional. **kwargs: Additional keyword arguments, including: The same keyword arguments used for AzureAISearchVectorStore: search_endpoint: The endpoint of the Azure AI Search service, optional. api_key: str | None = None, azure_credentials: AzureKeyCredential | None = None, token_credentials: AsyncTokenCredential | None = None, env_file_path: str | None = None, env_file_encoding: str | None = None """ if not collection_name: collection_name = _get_collection_name_from_model(record_type, definition) if not collection_name and search_client: collection_name = search_client._index_name if search_client and search_index_client: if collection_name and search_client._index_name != collection_name: search_client._index_name = collection_name super().__init__( record_type=record_type, definition=definition, collection_name=collection_name, search_client=search_client, search_index_client=search_index_client, search_endpoint=kwargs.get("search_endpoint"), search_credential=search_credential, managed_search_index_client=False, managed_client=False, embedding_generator=embedding_generator, ) return if search_index_client: try: azure_ai_search_settings = AzureAISearchSettings( env_file_path=kwargs.get("env_file_path"), endpoint=kwargs.get("search_endpoint"), api_key=kwargs.get("api_key"), env_file_encoding=kwargs.get("env_file_encoding"), index_name=collection_name, ) except ValidationError as exc: raise VectorStoreInitializationException("Failed to create Azure Cognitive Search settings.") from exc endpoint = str(azure_ai_search_settings.endpoint) credential = search_credential or _resolve_credential( azure_ai_search_settings, azure_credential=kwargs.get("azure_credentials"), token_credential=kwargs.get("token_credentials"), ) super().__init__( record_type=record_type, definition=definition, collection_name=azure_ai_search_settings.index_name, search_client=_get_search_client( endpoint=endpoint, collection_name=azure_ai_search_settings.index_name, credential=credential, ), search_index_client=search_index_client, search_endpoint=endpoint, search_credential=credential, managed_search_index_client=False, embedding_generator=embedding_generator, ) return try: azure_ai_search_settings = AzureAISearchSettings( env_file_path=kwargs.get("env_file_path"), endpoint=kwargs.get("search_endpoint"), api_key=kwargs.get("api_key"), env_file_encoding=kwargs.get("env_file_encoding"), index_name=collection_name, ) except ValidationError as exc: raise VectorStoreInitializationException("Failed to create Azure Cognitive Search settings.") from exc endpoint = str(azure_ai_search_settings.endpoint) credential = search_credential or _resolve_credential( azure_ai_search_settings, azure_credential=kwargs.get("azure_credentials"), token_credential=kwargs.get("token_credentials"), ) search_index_client = _get_search_index_client( azure_ai_search_settings=azure_ai_search_settings, azure_credential=kwargs.get("azure_credentials"), token_credential=kwargs.get("token_credentials"), ) super().__init__( record_type=record_type, definition=definition, collection_name=azure_ai_search_settings.index_name, search_client=_get_search_client( endpoint=endpoint, collection_name=azure_ai_search_settings.index_name, credential=credential, ), search_index_client=search_index_client, search_endpoint=endpoint, search_credential=credential, embedding_generator=embedding_generator, ) @override async def _inner_upsert( self, records: Sequence[Any], **kwargs: Any, ) -> Sequence[TKey]: if not isinstance(records, list): records = list(records) results = await self.search_client.merge_or_upload_documents(documents=records, **kwargs) return [result.key for result in results] # type: ignore @override async def _inner_get( self, keys: Sequence[TKey] | None = None, options: GetFilteredRecordOptions | None = None, **kwargs: Any, ) -> Sequence[dict[str, Any]]: client = self.search_client if "selected_fields" in kwargs: selected_fields = kwargs["selected_fields"] elif kwargs.get("include_vectors"): selected_fields = ["*"] else: selected_fields = self.definition.get_storage_names(include_vector_fields=False) if keys is not None: gather_result = await asyncio.gather( *[client.get_document(key=key, selected_fields=selected_fields) for key in keys], # type: ignore return_exceptions=True, ) return [res for res in gather_result if not isinstance(res, BaseException)] if options is not None: ordering = [] if options.order_by: for field, asc_flag in options.order_by.items(): if field not in self.definition.storage_names: logger.warning(f"Field {field} not in data model, skipping.") continue ordering.append(field if asc_flag else f"{field} desc") result = await client.search( search_text="*", top=options.top, skip=options.skip, select=selected_fields, order_by=ordering, ) return [res async for res in result] raise VectorStoreOperationException("No keys or options provided for get operation.") @override async def _inner_delete(self, keys: Sequence[TKey], **kwargs: Any) -> None: await self.search_client.delete_documents(documents=[{self._key_field_name: key} for key in keys]) @override def _serialize_dicts_to_store_models(self, records: Sequence[dict[str, Any]], **kwargs: Any) -> Sequence[Any]: return records @override def _deserialize_store_models_to_dicts(self, records: Sequence[Any], **kwargs: Any) -> Sequence[dict[str, Any]]: return records @override async def ensure_collection_exists(self, **kwargs) -> None: """Create a new collection in Azure AI Search. Args: **kwargs: Additional keyword arguments. index (SearchIndex): The search index to create, if this is supplied this is used instead of a index created based on the definition. encryption_key (SearchResourceEncryptionKey): The encryption key to use, not used when index is supplied. other kwargs are passed to the create_index method. """ if index := kwargs.pop("index", None): if isinstance(index, SearchIndex): await self.search_index_client.create_index(index=index, **kwargs) return raise VectorStoreOperationException("Invalid index type supplied, should be a SearchIndex object.") await self.search_index_client.create_index( index=_definition_to_azure_ai_search_index( collection_name=self.collection_name, definition=self.definition, encryption_key=kwargs.pop("encryption_key", None), ), **kwargs, ) @override async def collection_exists(self, **kwargs) -> bool: if "params" not in kwargs: kwargs["params"] = {"select": ["name"]} return self.collection_name in [ index_name async for index_name in self.search_index_client.list_index_names(**kwargs) ] @override async def ensure_collection_deleted(self, **kwargs) -> None: await self.search_index_client.delete_index(self.collection_name, **kwargs) @override async def _inner_search( self, search_type: SearchType, options: VectorSearchOptions, values: Any | None = None, vector: Sequence[float | int] | None = None, **kwargs: Any, ) -> KernelSearchResults[VectorSearchResult[TModel]]: search_args: dict[str, Any] = { "top": options.top, "skip": options.skip, "include_total_count": options.include_total_count, } if options.include_vectors: search_args["select"] = ["*"] else: search_args["select"] = self.definition.get_storage_names(include_vector_fields=False) if filter := self._build_filter(options.filter): search_args["filter"] = filter if isinstance(filter, str) else " and ".join(filter) match search_type: case SearchType.VECTOR: if vector is not None: vector_field = self.definition.try_get_vector_field(options.vector_property_name) search_args["vector_queries"] = [ VectorizedQuery( vector=vector, # type: ignore fields=vector_field.storage_name or vector_field.name if vector_field else None, ) ] elif values is not None: generated_vector = await self._generate_vector_from_values(values or "*", options) vector_field = self.definition.try_get_vector_field(options.vector_property_name) if generated_vector is not None: search_args["vector_queries"] = [ VectorizedQuery( vector=generated_vector, # type: ignore fields=vector_field.storage_name or vector_field.name if vector_field else None, ) ] else: search_args["vector_queries"] = [ VectorizableTextQuery( text=values, fields=vector_field.storage_name or vector_field.name if vector_field else None, ) ] else: raise VectorStoreOperationException("No vector or keywords provided for vector search.") case SearchType.KEYWORD_HYBRID: if values is None: raise VectorStoreOperationException("No vector and/or keywords provided for search.") vector_field = self.definition.try_get_vector_field(options.vector_property_name) search_args["search_fields"] = ( [options.additional_property_name] if options.additional_property_name is not None else [ field.name for field in self.definition.fields if field.field_type == FieldTypes.DATA and field.is_full_text_indexed ] ) if not search_args["search_fields"]: raise VectorStoreOperationException("No searchable fields found for hybrid search.") search_args["search_text"] = values vector = await self._generate_vector_from_values(values, options) if vector is None else vector if vector is not None: search_args["vector_queries"] = [ VectorizedQuery( vector=vector, # type: ignore fields=vector_field.name if vector_field else None, ) ] else: search_args["vector_queries"] = [ VectorizableTextQuery( text=values, fields=vector_field.name if vector_field else None, ) ] try: raw_results = await self.search_client.search(**search_args, **kwargs) except Exception as exc: raise VectorSearchExecutionException("Failed to search the collection.") from exc return KernelSearchResults( results=self._get_vector_search_results_from_results(raw_results, options), total_count=await raw_results.get_count() if options.include_total_count else None, ) @override def _lambda_parser(self, node: ast.AST) -> Any: def _parse_attribute_chain(attr_node: ast.Attribute) -> str: parts = [] current = attr_node while isinstance(current, ast.Attribute): parts.append(current.attr) current = current.value # type: ignore if isinstance(current, ast.Name): # skip the root variable name (e.g., 'x') pass else: raise NotImplementedError(f"Unsupported attribute chain root: {type(current)}") # reverse to get the correct order prop_path = "/".join(reversed(parts)) # Check if the top-level property is in the data model top_level = parts[-1] if parts else None if top_level and top_level not in self.definition.storage_names: raise VectorStoreOperationException( f"Field '{top_level}' not in data model (storage property names are used)." ) return prop_path match node: case ast.Compare(): if len(node.ops) > 1: values: list[ast.expr] = [] for idx in range(len(node.ops)): if idx == 0: values.append( ast.Compare( left=node.left, ops=[node.ops[idx]], comparators=[node.comparators[idx]], ) ) else: values.append( ast.Compare( left=node.comparators[idx - 1], ops=[node.ops[idx]], comparators=[node.comparators[idx]], ) ) return self._lambda_parser(ast.BoolOp(op=ast.And(), values=values)) left = self._lambda_parser(node.left) right = self._lambda_parser(node.comparators[0]) op = node.ops[0] match op: case ast.In(): return f"search.ismatch({left}, '{right}')" case ast.NotIn(): return f"not search.ismatch({left}, '{right}')" case ast.Eq(): return f"{left} eq {right}" case ast.NotEq(): return f"{left} ne {right}" case ast.Gt(): return f"{left} gt {right}" case ast.GtE(): return f"{left} ge {right}" case ast.Lt(): return f"{left} lt {right}" case ast.LtE(): return f"{left} le {right}" raise NotImplementedError(f"Unsupported operator: {type(op)}") case ast.BoolOp(): op_str = "and" if isinstance(node.op, ast.And) else "or" return "(" + f" {op_str} ".join([self._lambda_parser(v) for v in node.values]) + ")" case ast.UnaryOp(): match node.op: case ast.UAdd(): return f"+{self._lambda_parser(node.operand)}" case ast.USub(): return f"-{self._lambda_parser(node.operand)}" case ast.Invert(): raise NotImplementedError("Invert operation is not supported.") case ast.Not(): return f"not {self._lambda_parser(node.operand)}" case ast.Attribute(): # Support nested property chains return _parse_attribute_chain(node) case ast.Name(): raise NotImplementedError("Constants are not supported, make sure to use a value or a attribute.") case ast.Constant(): value = node.value if isinstance(value, str): return "'" + value.replace("'", "''") + "'" if isinstance(value, bytes): return "'" + value.decode("utf-8").replace("'", "''") + "'" if isinstance(value, bool): return str(value).lower() if value is None: return "null" if isinstance(value, (int, float)): return str(value) raise VectorStoreOperationException(f"Unsupported constant type: {type(value)}") raise NotImplementedError(f"Unsupported AST node: {type(node)}") @override def _get_record_from_result(self, result: dict[str, Any]) -> dict[str, Any]: return result @override def _get_score_from_result(self, result: dict[str, Any]) -> float | None: return result.get("@search.score") @override async def __aexit__(self, exc_type, exc_value, traceback) -> None: """Exit the context manager.""" if self.managed_client: await self.search_client.close() if self.managed_search_index_client: await self.search_index_client.close() @release_candidate class AzureAISearchStore(VectorStore): """Azure AI Search store implementation.""" search_index_client: SearchIndexClient search_endpoint: str | None = None search_credential: Any = None def __init__( self, search_endpoint: str | None = None, api_key: str | None = None, azure_credentials: "AzureKeyCredential | None" = None, token_credentials: "AsyncTokenCredential | None" = None, search_index_client: SearchIndexClient | None = None, embedding_generator: "EmbeddingGeneratorBase | None" = None, env_file_path: str | None = None, env_file_encoding: str | None = None, ) -> None: """Initializes a new instance of the AzureAISearchStore class.""" managed_client: bool = False endpoint: str | None = None credential: AzureKeyCredential | AsyncTokenCredential | None = None if not search_index_client: try: azure_ai_search_settings = AzureAISearchSettings( env_file_path=env_file_path, endpoint=search_endpoint, api_key=api_key, env_file_encoding=env_file_encoding, ) except ValidationError as exc: raise VectorStoreInitializationException("Failed to create Azure AI Search settings.") from exc endpoint = str(azure_ai_search_settings.endpoint) credential = _resolve_credential( azure_ai_search_settings, azure_credential=azure_credentials, token_credential=token_credentials, ) search_index_client = _get_search_index_client( azure_ai_search_settings=azure_ai_search_settings, azure_credential=azure_credentials, token_credential=token_credentials, ) managed_client = True else: endpoint = search_endpoint credential = azure_credentials or token_credentials or (AzureKeyCredential(api_key) if api_key else None) super().__init__( search_index_client=search_index_client, search_endpoint=endpoint, search_credential=credential, managed_client=managed_client, embedding_generator=embedding_generator, ) @override def get_collection( self, record_type: type[TModel], *, definition: VectorStoreCollectionDefinition | None = None, collection_name: str | None = None, embedding_generator: EmbeddingGeneratorBase | None = None, search_client: SearchClient | None = None, **kwargs: Any, ) -> AzureAISearchCollection: """Get a AzureAISearchCollection tied to a collection. Args: collection_name: The name of the collection. record_type: The type of the data model. definition: The model fields, optional. search_client: The search client for interacting with Azure AI Search, will be created if not supplied. embedding_generator: The embedding generator, optional. **kwargs: Additional keyword arguments, passed to the collection constructor. """ return AzureAISearchCollection( record_type=record_type, definition=definition, collection_name=collection_name, search_index_client=self.search_index_client, search_client=search_client, embedding_generator=embedding_generator or self.embedding_generator, search_credential=self.search_credential, search_endpoint=self.search_endpoint, **kwargs, ) @override async def list_collection_names(self, **kwargs: Any) -> list[str]: if "params" not in kwargs: kwargs["params"] = {"select": ["name"]} return [index async for index in self.search_index_client.list_index_names(**kwargs)] @override async def __aexit__(self, exc_type, exc_value, traceback) -> None: """Exit the context manager.""" if self.managed_client: await self.search_index_client.close()