unstructured-ingest 0.3.13__py3-none-any.whl

unstructured_ingest/v2/processes/connectors/astradb.py

@@ -0,0 +1,402 @@
+import csv
+import hashlib
+from dataclasses import dataclass, field
+from pathlib import Path
+from time import time
+from typing import TYPE_CHECKING, Any, Generator, Optional
+
+from pydantic import BaseModel, Field, Secret
+
+from unstructured_ingest import __name__ as integration_name
+from unstructured_ingest.__version__ import __version__ as integration_version
+from unstructured_ingest.error import (
+    DestinationConnectionError,
+    SourceConnectionError,
+    SourceConnectionNetworkError,
+)
+from unstructured_ingest.utils.data_prep import batch_generator, get_data
+from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.utils.string_and_date_utils import truncate_string_bytes
+from unstructured_ingest.v2.constants import RECORD_ID_LABEL
+from unstructured_ingest.v2.interfaces import (
+    AccessConfig,
+    BatchFileData,
+    BatchItem,
+    ConnectionConfig,
+    Downloader,
+    DownloaderConfig,
+    DownloadResponse,
+    FileData,
+    FileDataSourceMetadata,
+    Indexer,
+    IndexerConfig,
+    SourceIdentifiers,
+    Uploader,
+    UploaderConfig,
+    UploadStager,
+    UploadStagerConfig,
+    download_responses,
+)
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
+    SourceRegistryEntry,
+)
+
+if TYPE_CHECKING:
+    from astrapy import AsyncCollection as AstraDBAsyncCollection
+    from astrapy import Collection as AstraDBCollection
+    from astrapy import DataAPIClient as AstraDBClient
+
+
+CONNECTOR_TYPE = "astradb"
+
+MAX_CONTENT_PARAM_BYTE_SIZE = 8000
+
+
+class AstraDBAdditionalMetadata(BaseModel):
+    collection_name: str
+    keyspace: Optional[str] = None
+
+
+class AstraDBBatchFileData(BatchFileData):
+    additional_metadata: AstraDBAdditionalMetadata
+
+
+class AstraDBAccessConfig(AccessConfig):
+    token: str = Field(description="Astra DB Token with access to the database.")
+    api_endpoint: str = Field(description="The API endpoint for the Astra DB.")
+
+
+class AstraDBConnectionConfig(ConnectionConfig):
+    connector_type: str = Field(default=CONNECTOR_TYPE, init=False)
+    access_config: Secret[AstraDBAccessConfig]
+
+    @requires_dependencies(["astrapy"], extras="astradb")
+    def get_client(self) -> "AstraDBClient":
+        from astrapy import DataAPIClient as AstraDBClient
+
+        # Create a client object to interact with the Astra DB
+        # caller_name/version for Astra DB tracking
+        return AstraDBClient(
+            caller_name=integration_name,
+            caller_version=integration_version,
+        )
+
+
+def get_astra_collection(
+    connection_config: AstraDBConnectionConfig,
+    collection_name: str,
+    keyspace: str,
+) -> "AstraDBCollection":
+    # Build the Astra DB object.
+    access_configs = connection_config.access_config.get_secret_value()
+
+    # Create a client object to interact with the Astra DB
+    # caller_name/version for Astra DB tracking
+    client = connection_config.get_client()
+
+    # Get the database object
+    astra_db = client.get_database(
+        api_endpoint=access_configs.api_endpoint,
+        token=access_configs.token,
+        keyspace=keyspace,
+    )
+
+    # Connect to the collection
+    astra_db_collection = astra_db.get_collection(name=collection_name)
+    return astra_db_collection
+
+
+async def get_async_astra_collection(
+    connection_config: AstraDBConnectionConfig,
+    collection_name: str,
+    keyspace: str,
+) -> "AstraDBAsyncCollection":
+    # Build the Astra DB object.
+    access_configs = connection_config.access_config.get_secret_value()
+
+    # Create a client object to interact with the Astra DB
+    client = connection_config.get_client()
+
+    # Get the async database object
+    async_astra_db = client.get_async_database(
+        api_endpoint=access_configs.api_endpoint,
+        token=access_configs.token,
+        keyspace=keyspace,
+    )
+
+    # Get async collection from AsyncDatabase
+    async_astra_db_collection = await async_astra_db.get_collection(name=collection_name)
+    return async_astra_db_collection
+
+
+class AstraDBUploadStagerConfig(UploadStagerConfig):
+    pass
+
+
+class AstraDBIndexerConfig(IndexerConfig):
+    collection_name: str = Field(
+        description="The name of the Astra DB collection. "
+        "Note that the collection name must only include letters, "
+        "numbers, and underscores."
+    )
+    keyspace: Optional[str] = Field(default=None, description="The Astra DB connection keyspace.")
+    batch_size: int = Field(default=20, description="Number of records per batch")
+
+
+class AstraDBDownloaderConfig(DownloaderConfig):
+    fields: list[str] = field(default_factory=list)
+
+
+class AstraDBUploaderConfig(UploaderConfig):
+    collection_name: str = Field(
+        description="The name of the Astra DB collection. "
+        "Note that the collection name must only include letters, "
+        "numbers, and underscores."
+    )
+    keyspace: Optional[str] = Field(default=None, description="The Astra DB connection keyspace.")
+    requested_indexing_policy: Optional[dict[str, Any]] = Field(
+        default=None,
+        description="The indexing policy to use for the collection.",
+        examples=['{"deny": ["metadata"]}'],
+    )
+    batch_size: int = Field(default=20, description="Number of records per batch")
+    record_id_key: str = Field(
+        default=RECORD_ID_LABEL,
+        description="searchable key to find entries for the same record on previous runs",
+    )
+
+
+@dataclass
+class AstraDBIndexer(Indexer):
+    connection_config: AstraDBConnectionConfig
+    index_config: AstraDBIndexerConfig
+
+    def get_collection(self) -> "AstraDBCollection":
+        return get_astra_collection(
+            connection_config=self.connection_config,
+            collection_name=self.index_config.collection_name,
+            keyspace=self.index_config.keyspace,
+        )
+
+    def precheck(self) -> None:
+        try:
+            self.get_collection().options()
+        except Exception as e:
+            logger.error(f"Failed to validate connection {e}", exc_info=True)
+            raise SourceConnectionError(f"failed to validate connection: {e}")
+
+    def _get_doc_ids(self) -> set[str]:
+        """Fetches all document ids in an index"""
+        # Get the collection
+        collection = self.get_collection()
+
+        # Perform the find operation to get all items
+        astra_db_docs_cursor = collection.find({}, projection={"_id": True})
+
+        # Iterate over the cursor
+        astra_db_docs = []
+        for result in astra_db_docs_cursor:
+            astra_db_docs.append(result)
+
+        # Create file data for each astra record
+        ids = sorted([astra_record["_id"] for astra_record in astra_db_docs])
+
+        return set(ids)
+
+    def run(self, **kwargs: Any) -> Generator[AstraDBBatchFileData, None, None]:
+        all_ids = self._get_doc_ids()
+        ids = list(all_ids)
+        id_batches = batch_generator(ids, self.index_config.batch_size)
+
+        for batch in id_batches:
+            fd = AstraDBBatchFileData(
+                connector_type=CONNECTOR_TYPE,
+                metadata=FileDataSourceMetadata(
+                    date_processed=str(time()),
+                ),
+                additional_metadata=AstraDBAdditionalMetadata(
+                    collection_name=self.index_config.collection_name,
+                    keyspace=self.index_config.keyspace,
+                ),
+                batch_items=[BatchItem(identifier=b) for b in batch],
+            )
+            yield fd
+
+
+@dataclass
+class AstraDBDownloader(Downloader):
+    connection_config: AstraDBConnectionConfig
+    download_config: AstraDBDownloaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    def is_async(self) -> bool:
+        return True
+
+    def get_identifier(self, record_id: str) -> str:
+        f = f"{record_id}"
+        if self.download_config.fields:
+            f = "{}-{}".format(
+                f,
+                hashlib.sha256(",".join(self.download_config.fields).encode()).hexdigest()[:8],
+            )
+        return f
+
+    def write_astra_result_to_csv(self, astra_result: dict, download_path: str) -> None:
+        with open(download_path, "w", encoding="utf8") as f:
+            writer = csv.writer(f)
+            writer.writerow(astra_result.keys())
+            writer.writerow(astra_result.values())
+
+    def generate_download_response(
+        self, result: dict, file_data: AstraDBBatchFileData
+    ) -> DownloadResponse:
+        record_id = result["_id"]
+        filename_id = self.get_identifier(record_id=record_id)
+        filename = f"{filename_id}.csv"  # csv to preserve column info
+        download_path = self.download_dir / Path(filename)
+        logger.debug(f"Downloading results from record {record_id} as csv to {download_path}")
+        download_path.parent.mkdir(parents=True, exist_ok=True)
+        try:
+            self.write_astra_result_to_csv(astra_result=result, download_path=str(download_path))
+        except Exception as e:
+            logger.error(
+                f"failed to download from record {record_id} to {download_path}: {e}",
+                exc_info=True,
+            )
+            raise SourceConnectionNetworkError(f"failed to download file {file_data.identifier}")
+
+        # modify input file_data for download_response
+        file_data.source_identifiers = SourceIdentifiers(filename=filename, fullpath=filename)
+        cast_file_data = FileData.cast(file_data=file_data)
+        cast_file_data.identifier = filename
+        cast_file_data.metadata.date_processed = str(time())
+        cast_file_data.metadata.record_locator = {"document_id": record_id}
+        return super().generate_download_response(
+            file_data=cast_file_data, download_path=download_path
+        )
+
+    def run(self, file_data: FileData, **kwargs: Any) -> download_responses:
+        raise NotImplementedError("Use astradb run_async instead")
+
+    async def run_async(self, file_data: FileData, **kwargs: Any) -> download_responses:
+        # Get metadata from file_data
+        astra_file_data = AstraDBBatchFileData.cast(file_data=file_data)
+        ids: list[str] = [item.identifier for item in astra_file_data.batch_items]
+        collection_name: str = astra_file_data.additional_metadata.collection_name
+        keyspace: str = astra_file_data.additional_metadata.keyspace
+
+        # Retrieve results from async collection
+        download_responses = []
+        async_astra_collection = await get_async_astra_collection(
+            connection_config=self.connection_config,
+            collection_name=collection_name,
+            keyspace=keyspace,
+        )
+        async for result in async_astra_collection.find({"_id": {"$in": ids}}):
+            download_responses.append(
+                self.generate_download_response(result=result, file_data=astra_file_data)
+            )
+        return download_responses
+
+
+@dataclass
+class AstraDBUploadStager(UploadStager):
+    upload_stager_config: AstraDBUploadStagerConfig = field(
+        default_factory=lambda: AstraDBUploadStagerConfig()
+    )
+
+    def truncate_dict_elements(self, element_dict: dict) -> None:
+        text = element_dict.pop("text", None)
+        if text is not None:
+            element_dict["text"] = truncate_string_bytes(text, MAX_CONTENT_PARAM_BYTE_SIZE)
+        metadata = element_dict.get("metadata")
+        if metadata is not None and isinstance(metadata, dict):
+            text_as_html = element_dict["metadata"].pop("text_as_html", None)
+            if text_as_html is not None:
+                element_dict["metadata"]["text_as_html"] = truncate_string_bytes(
+                    text_as_html, MAX_CONTENT_PARAM_BYTE_SIZE
+                )
+
+    def conform_dict(self, element_dict: dict, file_data: FileData) -> dict:
+        self.truncate_dict_elements(element_dict)
+        return {
+            "$vector": element_dict.pop("embeddings", None),
+            "content": element_dict.pop("text", None),
+            RECORD_ID_LABEL: file_data.identifier,
+            "metadata": element_dict,
+        }
+
+
+@dataclass
+class AstraDBUploader(Uploader):
+    connection_config: AstraDBConnectionConfig
+    upload_config: AstraDBUploaderConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    def precheck(self) -> None:
+        try:
+            get_astra_collection(
+                connection_config=self.connection_config,
+                collection_name=self.upload_config.collection_name,
+                keyspace=self.upload_config.keyspace,
+            ).options()
+        except Exception as e:
+            logger.error(f"Failed to validate connection {e}", exc_info=True)
+            raise DestinationConnectionError(f"failed to validate connection: {e}")
+
+    @requires_dependencies(["astrapy"], extras="astradb")
+    def get_collection(self) -> "AstraDBCollection":
+        return get_astra_collection(
+            connection_config=self.connection_config,
+            collection_name=self.upload_config.collection_name,
+            keyspace=self.upload_config.keyspace,
+        )
+
+    def delete_by_record_id(self, collection: "AstraDBCollection", file_data: FileData):
+        logger.debug(
+            f"deleting records from collection {collection.name} "
+            f"with {self.upload_config.record_id_key} "
+            f"set to {file_data.identifier}"
+        )
+        delete_filter = {self.upload_config.record_id_key: {"$eq": file_data.identifier}}
+        delete_resp = collection.delete_many(filter=delete_filter)
+        logger.debug(
+            f"deleted {delete_resp.deleted_count} records from collection {collection.name}"
+        )
+
+    def run_data(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None:
+        logger.info(
+            f"writing {len(data)} objects to destination "
+            f"collection {self.upload_config.collection_name}"
+        )
+
+        astra_db_batch_size = self.upload_config.batch_size
+        collection = self.get_collection()
+
+        self.delete_by_record_id(collection=collection, file_data=file_data)
+
+        for chunk in batch_generator(data, astra_db_batch_size):
+            collection.insert_many(chunk)
+
+    def run(self, path: Path, file_data: FileData, **kwargs: Any) -> None:
+        data = get_data(path=path)
+        self.run_data(data=data, file_data=file_data, **kwargs)
+
+
+astra_db_source_entry = SourceRegistryEntry(
+    indexer=AstraDBIndexer,
+    indexer_config=AstraDBIndexerConfig,
+    downloader=AstraDBDownloader,
+    downloader_config=AstraDBDownloaderConfig,
+    connection_config=AstraDBConnectionConfig,
+)
+
+astra_db_destination_entry = DestinationRegistryEntry(
+    connection_config=AstraDBConnectionConfig,
+    upload_stager_config=AstraDBUploadStagerConfig,
+    upload_stager=AstraDBUploadStager,
+    uploader_config=AstraDBUploaderConfig,
+    uploader=AstraDBUploader,
+)

unstructured_ingest/v2/processes/connectors/azure_ai_search.py

@@ -0,0 +1,276 @@
+import json
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Generator
+
+from pydantic import Field, Secret
+
+from unstructured_ingest.error import DestinationConnectionError, WriteError
+from unstructured_ingest.utils.data_prep import batch_generator
+from unstructured_ingest.utils.dep_check import requires_dependencies
+from unstructured_ingest.v2.constants import RECORD_ID_LABEL
+from unstructured_ingest.v2.interfaces import (
+    AccessConfig,
+    ConnectionConfig,
+    FileData,
+    Uploader,
+    UploaderConfig,
+    UploadStager,
+    UploadStagerConfig,
+)
+from unstructured_ingest.v2.logger import logger
+from unstructured_ingest.v2.processes.connector_registry import (
+    DestinationRegistryEntry,
+)
+from unstructured_ingest.v2.processes.connectors.utils import parse_datetime
+from unstructured_ingest.v2.utils import get_enhanced_element_id
+
+if TYPE_CHECKING:
+    from azure.search.documents import SearchClient
+    from azure.search.documents.indexes import SearchIndexClient
+
+CONNECTOR_TYPE = "azure_ai_search"
+
+
+class AzureAISearchAccessConfig(AccessConfig):
+    azure_ai_search_key: str = Field(
+        alias="key", description="Credential that is used for authenticating to an Azure service"
+    )
+
+
+class AzureAISearchConnectionConfig(ConnectionConfig):
+    endpoint: str = Field(
+        description="The URL endpoint of an Azure AI (Cognitive) search service. "
+        "In the form of https://{{service_name}}.search.windows.net"
+    )
+    index: str = Field(
+        description="The name of the Azure AI (Cognitive) Search index to connect to."
+    )
+    access_config: Secret[AzureAISearchAccessConfig]
+
+    @requires_dependencies(["azure.search", "azure.core"], extras="azure-ai-search")
+    @contextmanager
+    def get_search_client(self) -> Generator["SearchClient", None, None]:
+        from azure.core.credentials import AzureKeyCredential
+        from azure.search.documents import SearchClient
+
+        with SearchClient(
+            endpoint=self.endpoint,
+            index_name=self.index,
+            credential=AzureKeyCredential(
+                self.access_config.get_secret_value().azure_ai_search_key
+            ),
+        ) as client:
+            yield client
+
+    @requires_dependencies(["azure.search", "azure.core"], extras="azure-ai-search")
+    @contextmanager
+    def get_search_index_client(self) -> Generator["SearchIndexClient", None, None]:
+        from azure.core.credentials import AzureKeyCredential
+        from azure.search.documents.indexes import SearchIndexClient
+
+        with SearchIndexClient(
+            endpoint=self.endpoint,
+            credential=AzureKeyCredential(
+                self.access_config.get_secret_value().azure_ai_search_key
+            ),
+        ) as search_index_client:
+            yield search_index_client
+
+
+class AzureAISearchUploadStagerConfig(UploadStagerConfig):
+    pass
+
+
+class AzureAISearchUploaderConfig(UploaderConfig):
+    batch_size: int = Field(default=100, description="Number of records per batch")
+    record_id_key: str = Field(
+        default=RECORD_ID_LABEL,
+        description="searchable key to find entries for the same record on previous runs",
+    )
+
+
+@dataclass
+class AzureAISearchUploadStager(UploadStager):
+    upload_stager_config: AzureAISearchUploadStagerConfig = field(
+        default_factory=lambda: AzureAISearchUploadStagerConfig()
+    )
+
+    def conform_dict(self, element_dict: dict, file_data: FileData) -> dict:
+        """
+        updates the dictionary that is from each Element being converted into a dict/json
+        into a dictionary that conforms to the schema expected by the
+        Azure Cognitive Search index
+        """
+        data = element_dict.copy()
+        data["id"] = get_enhanced_element_id(element_dict=data, file_data=file_data)
+        data[RECORD_ID_LABEL] = file_data.identifier
+
+        if points := data.get("metadata", {}).get("coordinates", {}).get("points"):
+            data["metadata"]["coordinates"]["points"] = json.dumps(points)
+        if version := data.get("metadata", {}).get("data_source", {}).get("version"):
+            data["metadata"]["data_source"]["version"] = str(version)
+        if record_locator := data.get("metadata", {}).get("data_source", {}).get("record_locator"):
+            data["metadata"]["data_source"]["record_locator"] = json.dumps(record_locator)
+        if permissions_data := (
+            data.get("metadata", {}).get("data_source", {}).get("permissions_data")
+        ):
+            data["metadata"]["data_source"]["permissions_data"] = json.dumps(permissions_data)
+        if links := data.get("metadata", {}).get("links"):
+            data["metadata"]["links"] = [json.dumps(link) for link in links]
+        if last_modified := data.get("metadata", {}).get("last_modified"):
+            data["metadata"]["last_modified"] = parse_datetime(last_modified).strftime(
+                "%Y-%m-%dT%H:%M:%S.%fZ"
+            )
+        if date_created := data.get("metadata", {}).get("data_source", {}).get("date_created"):
+            data["metadata"]["data_source"]["date_created"] = parse_datetime(date_created).strftime(
+                "%Y-%m-%dT%H:%M:%S.%fZ"
+            )
+
+        if date_modified := data.get("metadata", {}).get("data_source", {}).get("date_modified"):
+            data["metadata"]["data_source"]["date_modified"] = parse_datetime(
+                date_modified
+            ).strftime("%Y-%m-%dT%H:%M:%S.%fZ")
+
+        if date_processed := data.get("metadata", {}).get("data_source", {}).get("date_processed"):
+            data["metadata"]["data_source"]["date_processed"] = parse_datetime(
+                date_processed
+            ).strftime("%Y-%m-%dT%H:%M:%S.%fZ")
+
+        if regex_metadata := data.get("metadata", {}).get("regex_metadata"):
+            data["metadata"]["regex_metadata"] = json.dumps(regex_metadata)
+        if page_number := data.get("metadata", {}).get("page_number"):
+            data["metadata"]["page_number"] = str(page_number)
+        return data
+
+
+@dataclass
+class AzureAISearchUploader(Uploader):
+    upload_config: AzureAISearchUploaderConfig
+    connection_config: AzureAISearchConnectionConfig
+    connector_type: str = CONNECTOR_TYPE
+
+    def query_docs(self, record_id: str, index_key: str) -> list[str]:
+        with self.connection_config.get_search_client() as search_client:
+            results = list(
+                search_client.search(filter=f"record_id eq '{record_id}'", select=[index_key])
+            )
+        return [result[index_key] for result in results]
+
+    def delete_by_record_id(self, file_data: FileData, index_key: str) -> None:
+        logger.debug(
+            f"deleting any content with metadata "
+            f"{self.upload_config.record_id_key}={file_data.identifier} "
+            f"from azure cognitive search index: {self.connection_config.index}"
+        )
+        doc_ids_to_delete = self.query_docs(record_id=file_data.identifier, index_key=index_key)
+        if not doc_ids_to_delete:
+            return
+        with self.connection_config.get_search_client() as search_client:
+            results = search_client.delete_documents(
+                documents=[{index_key: doc_id} for doc_id in doc_ids_to_delete]
+            )
+        errors = []
+        success = []
+        for result in results:
+            if result.succeeded:
+                success.append(result)
+            else:
+                errors.append(result)
+        logger.debug(f"results: {len(success)} successes, {len(errors)} failures")
+        if errors:
+            raise WriteError(
+                ", ".join(
+                    [f"[{error.status_code}] {error.error_message}" for error in errors],
+                ),
+            )
+
+    @DestinationConnectionError.wrap
+    @requires_dependencies(["azure"], extras="azure-ai-search")
+    def write_dict(
+        self, elements_dict: list[dict[str, Any]], search_client: "SearchClient"
+    ) -> None:
+        import azure.core.exceptions
+
+        logger.info(
+            f"writing {len(elements_dict)} documents to destination "
+            f"index at {self.connection_config.index}",
+        )
+        try:
+            results = search_client.upload_documents(documents=elements_dict)
+        except azure.core.exceptions.HttpResponseError as http_error:
+            raise WriteError(f"http error: {http_error}") from http_error
+
+        errors = []
+        success = []
+        for result in results:
+            if result.succeeded:
+                success.append(result)
+            else:
+                errors.append(result)
+        logger.debug(f"results: {len(success)} successes, {len(errors)} failures")
+        if errors:
+            raise WriteError(
+                ", ".join(
+                    [
+                        f"{error.key}: " f"[{error.status_code}] {error.error_message}"
+                        for error in errors
+                    ],
+                ),
+            )
+
+    def can_delete(self) -> bool:
+        with self.connection_config.get_search_index_client() as search_index_client:
+            index = search_index_client.get_index(name=self.connection_config.index)
+        index_fields = index.fields
+        record_id_fields = [
+            field for field in index_fields if field.name == self.upload_config.record_id_key
+        ]
+        if not record_id_fields:
+            return False
+        record_id_field = record_id_fields[0]
+        return record_id_field.filterable
+
+    def get_index_key(self) -> str:
+        with self.connection_config.get_search_index_client() as search_index_client:
+            index = search_index_client.get_index(name=self.connection_config.index)
+        index_fields = index.fields
+        key_fields = [field for field in index_fields if field.key]
+        if not key_fields:
+            raise ValueError("no key field found in index fields")
+        return key_fields[0].name
+
+    def precheck(self) -> None:
+        try:
+            with self.connection_config.get_search_client() as search_client:
+                search_client.get_document_count()
+        except Exception as e:
+            logger.error(f"failed to validate connection: {e}", exc_info=True)
+            raise DestinationConnectionError(f"failed to validate connection: {e}")
+
+    def run_data(self, data: list[dict], file_data: FileData, **kwargs: Any) -> None:
+        logger.info(
+            f"writing document batches to destination"
+            f" endpoint at {str(self.connection_config.endpoint)}"
+            f" index at {str(self.connection_config.index)}"
+            f" with batch size {str(self.upload_config.batch_size)}"
+        )
+        if self.can_delete():
+            index_key = self.get_index_key()
+            self.delete_by_record_id(file_data=file_data, index_key=index_key)
+        else:
+            logger.warning("criteria for deleting previous content not met, skipping")
+
+        batch_size = self.upload_config.batch_size
+        with self.connection_config.get_search_client() as search_client:
+            for chunk in batch_generator(data, batch_size):
+                self.write_dict(elements_dict=chunk, search_client=search_client)  # noqa: E203
+
+
+azure_ai_search_destination_entry = DestinationRegistryEntry(
+    connection_config=AzureAISearchConnectionConfig,
+    uploader=AzureAISearchUploader,
+    uploader_config=AzureAISearchUploaderConfig,
+    upload_stager=AzureAISearchUploadStager,
+    upload_stager_config=AzureAISearchUploadStagerConfig,
+)