chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,852 @@
|
||||
# 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()
|
||||
Reference in New Issue
Block a user