Azureaisearch

CognitiveSearchVectorStore module-attribute

CognitiveSearchVectorStore = AzureAISearchVectorStore

AzureAISearchVectorStore

Bases: BasePydanticVectorStore

# Azure AI Search向量存储。

# 示例：
# `pip install llama-index-vector-stores-azureaisearch`

from azure.core.credentials import AzureKeyCredential
from azure.search.documents import SearchClient
from azure.search.documents.indexes import SearchIndexClient
from llama_index.vector_stores.azureaisearch import AzureAISearchVectorStore
from llama_index.vector_stores.azureaisearch import IndexManagement, MetadataIndexFieldType

# Azure AI Search设置
search_service_api_key = "YOUR-AZURE-SEARCH-SERVICE-ADMIN-KEY"
search_service_endpoint = "YOUR-AZURE-SEARCH-SERVICE-ENDPOINT"
search_service_api_version = "2023-11-01"
credential = AzureKeyCredential(search_service_api_key)

# 要使用的索引名称
index_name = "llamaindex-vector-demo"

# 使用索引客户端来演示创建索引
index_client = SearchIndexClient(
    endpoint=search_service_endpoint,
    credential=credential,
)

metadata_fields = {
    "author": "author",
    "theme": ("topic", MetadataIndexFieldType.STRING),
    "director": "director",
}

# 创建Azure AI Search向量存储
vector_store = AzureAISearchVectorStore(
    search_or_index_client=index_client,
    filterable_metadata_field_keys=metadata_fields,
    index_name=index_name,
    index_management=IndexManagement.CREATE_IF_NOT_EXISTS,
    id_field_key="id",
    chunk_field_key="chunk",
    embedding_field_key="embedding",
    embedding_dimensionality=1536,
    metadata_string_field_key="metadata",
    doc_id_field_key="doc_id",
    language_analyzer="en.lucene",
    vector_algorithm_type="exhaustiveKnn",
)

Source code in llama_index/vector_stores/azureaisearch/base.py

class AzureAISearchVectorStore(BasePydanticVectorStore):
    """```python
# Azure AI Search向量存储。

# 示例：
# `pip install llama-index-vector-stores-azureaisearch`

from azure.core.credentials import AzureKeyCredential
from azure.search.documents import SearchClient
from azure.search.documents.indexes import SearchIndexClient
from llama_index.vector_stores.azureaisearch import AzureAISearchVectorStore
from llama_index.vector_stores.azureaisearch import IndexManagement, MetadataIndexFieldType

# Azure AI Search设置
search_service_api_key = "YOUR-AZURE-SEARCH-SERVICE-ADMIN-KEY"
search_service_endpoint = "YOUR-AZURE-SEARCH-SERVICE-ENDPOINT"
search_service_api_version = "2023-11-01"
credential = AzureKeyCredential(search_service_api_key)

# 要使用的索引名称
index_name = "llamaindex-vector-demo"

# 使用索引客户端来演示创建索引
index_client = SearchIndexClient(
    endpoint=search_service_endpoint,
    credential=credential,
)

metadata_fields = {
    "author": "author",
    "theme": ("topic", MetadataIndexFieldType.STRING),
    "director": "director",
}

# 创建Azure AI Search向量存储
vector_store = AzureAISearchVectorStore(
    search_or_index_client=index_client,
    filterable_metadata_field_keys=metadata_fields,
    index_name=index_name,
    index_management=IndexManagement.CREATE_IF_NOT_EXISTS,
    id_field_key="id",
    chunk_field_key="chunk",
    embedding_field_key="embedding",
    embedding_dimensionality=1536,
    metadata_string_field_key="metadata",
    doc_id_field_key="doc_id",
    language_analyzer="en.lucene",
    vector_algorithm_type="exhaustiveKnn",
)
```"""

    stores_text: bool = True
    flat_metadata: bool = True

    _index_client: SearchIndexClient = PrivateAttr()
    _search_client: SearchClient = PrivateAttr()
    _embedding_dimensionality: int = PrivateAttr()
    _language_analyzer: str = PrivateAttr()
    _field_mapping: Dict[str, str] = PrivateAttr()
    _index_management: IndexManagement = PrivateAttr()
    _index_mapping: Callable[
        [Dict[str, str], Dict[str, Any]], Dict[str, str]
    ] = PrivateAttr()
    _metadata_to_index_field_map: Dict[
        str, Tuple[str, MetadataIndexFieldType]
    ] = PrivateAttr()
    _vector_profile_name: str = PrivateAttr()

    def _normalise_metadata_to_index_fields(
        self,
        filterable_metadata_field_keys: Union[
            List[str],
            Dict[str, str],
            Dict[str, Tuple[str, MetadataIndexFieldType]],
            None,
        ] = [],
    ) -> Dict[str, Tuple[str, MetadataIndexFieldType]]:
        index_field_spec: Dict[str, Tuple[str, MetadataIndexFieldType]] = {}

        if isinstance(filterable_metadata_field_keys, List):
            for field in filterable_metadata_field_keys:
                # Index field name and the metadata field name are the same
                # Use String as the default index field type
                index_field_spec[field] = (field, MetadataIndexFieldType.STRING)

        elif isinstance(filterable_metadata_field_keys, Dict):
            for k, v in filterable_metadata_field_keys.items():
                if isinstance(v, tuple):
                    # Index field name and metadata field name may differ
                    # The index field type used is as supplied
                    index_field_spec[k] = v
                else:
                    # Index field name and metadata field name may differ
                    # Use String as the default index field type
                    index_field_spec[k] = (v, MetadataIndexFieldType.STRING)

        return index_field_spec

    def _create_index_if_not_exists(self, index_name: str) -> None:
        if index_name not in self._index_client.list_index_names():
            logger.info(
                f"Index {index_name} does not exist in Azure AI Search, creating index"
            )
            self._create_index(index_name)

    def _create_metadata_index_fields(self) -> List[Any]:
        """创建一个用于存储元数据值的索引字段列表。"""
        from azure.search.documents.indexes.models import SimpleField

        index_fields = []

        # create search fields
        for v in self._metadata_to_index_field_map.values():
            field_name, field_type = v

            if field_type == MetadataIndexFieldType.STRING:
                index_field_type = "Edm.String"
            elif field_type == MetadataIndexFieldType.INT32:
                index_field_type = "Edm.Int32"
            elif field_type == MetadataIndexFieldType.INT64:
                index_field_type = "Edm.Int64"
            elif field_type == MetadataIndexFieldType.DOUBLE:
                index_field_type = "Edm.Double"
            elif field_type == MetadataIndexFieldType.BOOLEAN:
                index_field_type = "Edm.Boolean"

            field = SimpleField(name=field_name, type=index_field_type, filterable=True)
            index_fields.append(field)

        return index_fields

    def _create_index(self, index_name: Optional[str]) -> None:
        """根据提供的索引名称、键字段名称和元数据过滤键创建默认索引。
"""
        from azure.search.documents.indexes.models import (
            ExhaustiveKnnAlgorithmConfiguration,
            ExhaustiveKnnParameters,
            HnswAlgorithmConfiguration,
            HnswParameters,
            SearchableField,
            SearchField,
            SearchFieldDataType,
            SearchIndex,
            SemanticConfiguration,
            SemanticField,
            SemanticPrioritizedFields,
            SemanticSearch,
            SimpleField,
            VectorSearch,
            VectorSearchAlgorithmKind,
            VectorSearchAlgorithmMetric,
            VectorSearchProfile,
        )

        logger.info(f"Configuring {index_name} fields for Azure AI Search")
        fields = [
            SimpleField(name=self._field_mapping["id"], type="Edm.String", key=True),
            SearchableField(
                name=self._field_mapping["chunk"],
                type="Edm.String",
                analyzer_name=self._language_analyzer,
            ),
            SearchField(
                name=self._field_mapping["embedding"],
                type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
                searchable=True,
                vector_search_dimensions=self._embedding_dimensionality,
                vector_search_profile_name=self._vector_profile_name,
            ),
            SimpleField(name=self._field_mapping["metadata"], type="Edm.String"),
            SimpleField(
                name=self._field_mapping["doc_id"], type="Edm.String", filterable=True
            ),
        ]
        logger.info(f"Configuring {index_name} metadata fields")
        metadata_index_fields = self._create_metadata_index_fields()
        fields.extend(metadata_index_fields)
        logger.info(f"Configuring {index_name} vector search")
        # Configure the vector search algorithms and profiles
        vector_search = VectorSearch(
            algorithms=[
                HnswAlgorithmConfiguration(
                    name="myHnsw",
                    kind=VectorSearchAlgorithmKind.HNSW,
                    # For more information on HNSw parameters, visit https://learn.microsoft.com//azure/search/vector-search-ranking#creating-the-hnsw-graph
                    parameters=HnswParameters(
                        m=4,
                        ef_construction=400,
                        ef_search=500,
                        metric=VectorSearchAlgorithmMetric.COSINE,
                    ),
                ),
                ExhaustiveKnnAlgorithmConfiguration(
                    name="myExhaustiveKnn",
                    kind=VectorSearchAlgorithmKind.EXHAUSTIVE_KNN,
                    parameters=ExhaustiveKnnParameters(
                        metric=VectorSearchAlgorithmMetric.COSINE,
                    ),
                ),
            ],
            profiles=[
                VectorSearchProfile(
                    name="myHnswProfile",
                    algorithm_configuration_name="myHnsw",
                ),
                # Add more profiles if needed
                VectorSearchProfile(
                    name="myExhaustiveKnnProfile",
                    algorithm_configuration_name="myExhaustiveKnn",
                ),
                # Add more profiles if needed
            ],
        )
        logger.info(f"Configuring {index_name} semantic search")
        semantic_config = SemanticConfiguration(
            name="mySemanticConfig",
            prioritized_fields=SemanticPrioritizedFields(
                content_fields=[SemanticField(field_name=self._field_mapping["chunk"])],
            ),
        )

        semantic_search = SemanticSearch(configurations=[semantic_config])

        index = SearchIndex(
            name=index_name,
            fields=fields,
            vector_search=vector_search,
            semantic_search=semantic_search,
        )
        logger.debug(f"Creating {index_name} search index")
        self._index_client.create_index(index)

    def _validate_index(self, index_name: Optional[str]) -> None:
        if self._index_client and index_name:
            if index_name not in self._index_client.list_index_names():
                raise ValueError(
                    f"Validation failed, index {index_name} does not exist."
                )

    def __init__(
        self,
        search_or_index_client: Any,
        id_field_key: str,
        chunk_field_key: str,
        embedding_field_key: str,
        metadata_string_field_key: str,
        doc_id_field_key: str,
        filterable_metadata_field_keys: Optional[
            Union[
                List[str],
                Dict[str, str],
                Dict[str, Tuple[str, MetadataIndexFieldType]],
            ]
        ] = None,
        index_name: Optional[str] = None,
        index_mapping: Optional[
            Callable[[Dict[str, str], Dict[str, Any]], Dict[str, str]]
        ] = None,
        index_management: IndexManagement = IndexManagement.NO_VALIDATION,
        embedding_dimensionality: int = 1536,
        vector_algorithm_type: str = "exhaustiveKnn",
        # If we have content in other languages, it is better to enable the language analyzer to be adjusted in searchable fields.
        # https://learn.microsoft.com/en-us/azure/search/index-add-language-analyzers
        language_analyzer: str = "en.lucene",
        **kwargs: Any,
    ) -> None:
        # ruff: noqa: E501
        """嵌入和文档存储在 Azure AI Search 索引中，添加嵌入时使用合并或上传方法。当添加多个嵌入时，索引会以每批 10 个文档的方式进行更新，如果批处理字节大小超出限制，可能会导致失败。

Args:
    search_client (azure.search.documents.SearchClient)：
        用于填充/查询的索引客户端。
    id_field_key (str)：存储 id 的索引字段
    chunk_field_key (str)：存储节点文本的索引字段
    embedding_field_key (str)：存储嵌入向量的索引字段
    metadata_string_field_key (str)：
        将节点元数据存储为 JSON 字符串的索引字段。
        架构是任意的，要对元数据值进行过滤，它们必须存储为索引中的单独字段，使用 filterable_metadata_field_keys 指定应存储在这些可过滤字段中的元数据值
    doc_id_field_key (str)：存储 doc_id 的索引字段
    index_mapping：
        可选函数，具有定义
        (enriched_doc: Dict[str, str], metadata: Dict[str, Any]): Dict[str,str]
        用于将文档字段映射到 AI 搜索索引字段（函数的返回值）。
        如果未指定，则提供默认映射，使用字段键。enriched_doc 中的键为 ["id", "chunk", "embedding", "metadata"]。
        默认映射为：
            - "id" 到 id_field_key
            - "chunk" 到 chunk_field_key
            - "embedding" 到 embedding_field_key
            - "metadata" 到 metadata_field_key
    *kwargs (Any)：其他关键字参数。

抛出：
    ImportError：无法导入 `azure-search-documents`
    ValueError：如果未提供 `search_or_index_client`
    ValueError：如果未提供 `index_name`，且 `search_or_index_client` 为 azure.search.documents.SearchIndexClient 类型
    ValueError：如果提供了 `index_name`，且 `search_or_index_client` 为 azure.search.documents.SearchClient 类型
    ValueError：如果 `create_index_if_not_exists` 为 true，且 `search_or_index_client` 为 azure.search.documents.SearchClient 类型
"""
        import_err_msg = (
            "`azure-search-documents` package not found, please run "
            "`pip install azure-search-documents==11.4.0`"
        )

        try:
            import azure.search.documents  # noqa
            from azure.search.documents import SearchClient
            from azure.search.documents.indexes import SearchIndexClient
        except ImportError:
            raise ImportError(import_err_msg)

        self._index_client: SearchIndexClient = cast(SearchIndexClient, None)
        self._search_client: SearchClient = cast(SearchClient, None)
        self._embedding_dimensionality = embedding_dimensionality

        if vector_algorithm_type == "exhaustiveKnn":
            self._vector_profile_name = "myExhaustiveKnnProfile"
        elif vector_algorithm_type == "hnsw":
            self._vector_profile_name = "myHnswProfile"
        else:
            raise ValueError(
                "Only 'exhaustiveKnn' and 'hnsw' are supported for vector_algorithm_type"
            )

        self._language_analyzer = language_analyzer

        # Validate search_or_index_client
        if search_or_index_client is not None:
            if isinstance(search_or_index_client, SearchIndexClient):
                # If SearchIndexClient is supplied so must index_name
                self._index_client = cast(SearchIndexClient, search_or_index_client)

                if not index_name:
                    raise ValueError(
                        "index_name must be supplied if search_or_index_client is of "
                        "type azure.search.documents.SearchIndexClient"
                    )

                self._search_client = self._index_client.get_search_client(
                    index_name=index_name
                )

            elif isinstance(search_or_index_client, SearchClient):
                self._search_client = cast(SearchClient, search_or_index_client)

                # Validate index_name
                if index_name:
                    raise ValueError(
                        "index_name cannot be supplied if search_or_index_client "
                        "is of type azure.search.documents.SearchClient"
                    )

            if not self._index_client and not self._search_client:
                raise ValueError(
                    "search_or_index_client must be of type "
                    "azure.search.documents.SearchClient or "
                    "azure.search.documents.SearchIndexClient"
                )
        else:
            raise ValueError("search_or_index_client not specified")

        if (
            index_management == IndexManagement.CREATE_IF_NOT_EXISTS
            and not self._index_client
        ):
            raise ValueError(
                "index_management has value of IndexManagement.CREATE_IF_NOT_EXISTS "
                "but search_or_index_client is not of type "
                "azure.search.documents.SearchIndexClient"
            )

        self._index_management = index_management

        # Default field mapping
        field_mapping = {
            "id": id_field_key,
            "chunk": chunk_field_key,
            "embedding": embedding_field_key,
            "metadata": metadata_string_field_key,
            "doc_id": doc_id_field_key,
        }

        self._field_mapping = field_mapping

        self._index_mapping = (
            self._default_index_mapping if index_mapping is None else index_mapping
        )

        # self._filterable_metadata_field_keys = filterable_metadata_field_keys
        self._metadata_to_index_field_map = self._normalise_metadata_to_index_fields(
            filterable_metadata_field_keys
        )

        if self._index_management == IndexManagement.CREATE_IF_NOT_EXISTS:
            if index_name:
                self._create_index_if_not_exists(index_name)

        if self._index_management == IndexManagement.VALIDATE_INDEX:
            self._validate_index(index_name)

        super().__init__()

    @property
    def client(self) -> Any:
        """获取客户端。"""
        return self._search_client

    def _default_index_mapping(
        self, enriched_doc: Dict[str, str], metadata: Dict[str, Any]
    ) -> Dict[str, str]:
        index_doc: Dict[str, str] = {}

        for field in self._field_mapping:
            index_doc[self._field_mapping[field]] = enriched_doc[field]

        for metadata_field_name, (
            index_field_name,
            _,
        ) in self._metadata_to_index_field_map.items():
            metadata_value = metadata.get(metadata_field_name)
            if metadata_value:
                index_doc[index_field_name] = metadata_value

        return index_doc

    def add(
        self,
        nodes: List[BaseNode],
        **add_kwargs: Any,
    ) -> List[str]:
        """将节点添加到与配置的搜索客户端相关联的索引中。

Args:
    节点：List[BaseNode]：具有嵌入的节点
"""
        if not self._search_client:
            raise ValueError("Search client not initialized")

        documents = []
        ids = []

        for node in nodes:
            logger.debug(f"Processing embedding: {node.node_id}")
            ids.append(node.node_id)

            index_document = self._create_index_document(node)

            documents.append(index_document)

            if len(documents) >= 10:
                logger.info(
                    f"Uploading batch of size {len(documents)}, "
                    f"current progress {len(ids)} of {len(nodes)}"
                )
                self._search_client.merge_or_upload_documents(documents)
                documents = []

        # Upload remaining batch of less than 10 documents
        if len(documents) > 0:
            logger.info(
                f"Uploading remaining batch of size {len(documents)}, "
                f"current progress {len(ids)} of {len(nodes)}"
            )
            self._search_client.merge_or_upload_documents(documents)
            documents = []

        return ids

    def _create_index_document(self, node: BaseNode) -> Dict[str, Any]:
        """从嵌入结果中创建AI搜索索引文档。"""
        doc: Dict[str, Any] = {}
        doc["id"] = node.node_id
        doc["chunk"] = node.get_content(metadata_mode=MetadataMode.NONE) or ""
        doc["embedding"] = node.get_embedding()
        doc["doc_id"] = node.ref_doc_id

        node_metadata = node_to_metadata_dict(
            node,
            remove_text=True,
            flat_metadata=self.flat_metadata,
        )

        doc["metadata"] = json.dumps(node_metadata)

        return self._index_mapping(doc, node_metadata)

    def delete(self, ref_doc_id: str, **delete_kwargs: Any) -> None:
        """从AI搜索索引中删除文档，其中doc_id_field_key字段等于ref_doc_id。
"""
        # Locate documents to delete
        filter = f'{self._field_mapping["doc_id"]} eq \'{ref_doc_id}\''
        results = self._search_client.search(search_text="*", filter=filter)

        logger.debug(f"Searching with filter {filter}")

        docs_to_delete = []
        for result in results:
            doc = {}
            doc["id"] = result[self._field_mapping["id"]]
            logger.debug(f"Found document to delete: {doc}")
            docs_to_delete.append(doc)

        if len(docs_to_delete) > 0:
            logger.debug(f"Deleting {len(docs_to_delete)} documents")
            self._search_client.delete_documents(docs_to_delete)

    def _create_odata_filter(self, metadata_filters: MetadataFilters) -> str:
        """使用提供的元数据过滤器生成一个OData过滤字符串。"""
        odata_filter: List[str] = []
        for f in metadata_filters.legacy_filters():
            if not isinstance(f, ExactMatchFilter):
                raise NotImplementedError(
                    "Only `ExactMatchFilter` filters are supported"
                )

            # Raise error if filtering on a metadata field that lacks a mapping to
            # an index field
            metadata_mapping = self._metadata_to_index_field_map.get(f.key)

            if not metadata_mapping:
                raise ValueError(
                    f"Metadata field '{f.key}' is missing a mapping to an index field, "
                    "provide entry in 'filterable_metadata_field_keys' for this "
                    "vector store"
                )

            index_field = metadata_mapping[0]

            if len(odata_filter) > 0:
                odata_filter.append(f" {metadata_filters.condition.value} ")
            if isinstance(f.value, str):
                escaped_value = "".join([("''" if s == "'" else s) for s in f.value])
                odata_filter.append(f"{index_field} eq '{escaped_value}'")
            else:
                odata_filter.append(f"{index_field} eq {f.value}")

        odata_expr = "".join(odata_filter)

        logger.info(f"Odata filter: {odata_expr}")

        return odata_expr

    def query(self, query: VectorStoreQuery, **kwargs: Any) -> VectorStoreQueryResult:
        odata_filter = None
        if query.filters is not None:
            odata_filter = self._create_odata_filter(query.filters)
        azure_query_result_search: AzureQueryResultSearchBase = (
            AzureQueryResultSearchDefault(
                query, self._field_mapping, odata_filter, self._search_client
            )
        )
        if query.mode == VectorStoreQueryMode.SPARSE:
            azure_query_result_search = AzureQueryResultSearchSparse(
                query, self._field_mapping, odata_filter, self._search_client
            )
        elif query.mode == VectorStoreQueryMode.HYBRID:
            azure_query_result_search = AzureQueryResultSearchHybrid(
                query, self._field_mapping, odata_filter, self._search_client
            )
        elif query.mode == VectorStoreQueryMode.SEMANTIC_HYBRID:
            azure_query_result_search = AzureQueryResultSearchSemanticHybrid(
                query, self._field_mapping, odata_filter, self._search_client
            )
        return azure_query_result_search.search()

client property

client: Any

获取客户端。

add

add(nodes: List[BaseNode], **add_kwargs: Any) -> List[str]

将节点添加到与配置的搜索客户端相关联的索引中。

Source code in llama_index/vector_stores/azureaisearch/base.py

    def add(
        self,
        nodes: List[BaseNode],
        **add_kwargs: Any,
    ) -> List[str]:
        """将节点添加到与配置的搜索客户端相关联的索引中。

Args:
    节点：List[BaseNode]：具有嵌入的节点
"""
        if not self._search_client:
            raise ValueError("Search client not initialized")

        documents = []
        ids = []

        for node in nodes:
            logger.debug(f"Processing embedding: {node.node_id}")
            ids.append(node.node_id)

            index_document = self._create_index_document(node)

            documents.append(index_document)

            if len(documents) >= 10:
                logger.info(
                    f"Uploading batch of size {len(documents)}, "
                    f"current progress {len(ids)} of {len(nodes)}"
                )
                self._search_client.merge_or_upload_documents(documents)
                documents = []

        # Upload remaining batch of less than 10 documents
        if len(documents) > 0:
            logger.info(
                f"Uploading remaining batch of size {len(documents)}, "
                f"current progress {len(ids)} of {len(nodes)}"
            )
            self._search_client.merge_or_upload_documents(documents)
            documents = []

        return ids

delete

delete(ref_doc_id: str, **delete_kwargs: Any) -> None

从AI搜索索引中删除文档，其中doc_id_field_key字段等于ref_doc_id。

Source code in llama_index/vector_stores/azureaisearch/base.py

    def delete(self, ref_doc_id: str, **delete_kwargs: Any) -> None:
        """从AI搜索索引中删除文档，其中doc_id_field_key字段等于ref_doc_id。
"""
        # Locate documents to delete
        filter = f'{self._field_mapping["doc_id"]} eq \'{ref_doc_id}\''
        results = self._search_client.search(search_text="*", filter=filter)

        logger.debug(f"Searching with filter {filter}")

        docs_to_delete = []
        for result in results:
            doc = {}
            doc["id"] = result[self._field_mapping["id"]]
            logger.debug(f"Found document to delete: {doc}")
            docs_to_delete.append(doc)

        if len(docs_to_delete) > 0:
            logger.debug(f"Deleting {len(docs_to_delete)} documents")
            self._search_client.delete_documents(docs_to_delete)