setup.py

import argparse
import json
import logging
import time

import requests
from azure.core.credentials import AccessToken, TokenCredential
from azure.core.exceptions import ClientAuthenticationError
from azure.identity import (
    AzureCliCredential,
    ChainedTokenCredential,
    ManagedIdentityCredential,
)
from azure.mgmt.web import WebSiteManagementClient

# Set up logging configuration globally
logging.getLogger("azure").setLevel(logging.WARNING)


def format_timespan(seconds: int) -> str:
    hours = seconds // 3600
    minutes = (seconds - hours * 3600) // 60
    remaining_seconds = seconds - hours * 3600 - minutes * 60
    timespan = f"{hours} hours {minutes} minutes {remaining_seconds:.4f} seconds"
    return timespan


def call_search_api(
    search_service: str,
    search_api_version: str,
    resource_type: str,
    resource_name: str,
    method: str,
    credential: TokenCredential,
    body: dict | None = None,
) -> None:
    """
    Calls the Azure AI Search API with the specified parameters.

    Args:
        search_service (str): The name of the Azure AI Search service.
        search_api_version (str): The version of the Azure AI Search API.
        resource_type (str): The type of resource to access (e.g. "indexes", "docs").
        resource_name (str): The name of the resource to access.
        method (str): The HTTP method to use (either GET, PUT or DELETE).
        credential (TokenCredential): An instance of a TokenCredential class that can provide an access token.
        body (dict | None, optional):  The JSON payload to include in the request body (for PUT requests).
            Defaults to None.

    Returns:
        None

    Raises:
        ValueError: If the specified HTTP method is not GET, PUT or DELETE.
        HTTPError: If the response status code is 400 or greater.
    """
    # Get the token
    access_token: AccessToken = credential.get_token(
        "https://search.azure.com/.default"
    )
    token = access_token.token
    headers = {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json",
        # "api-key": SEARCH_API_KEY
    }
    search_endpoint = f"https://{search_service}.search.windows.net/{resource_type}/{resource_name}?api-version={search_api_version}"
    response = None

    try:
        if method not in ["get", "put", "delete"]:
            logging.warning(f"[call_search_api] Invalid method: {method}.")

        # GET
        if method == "get":
            response = requests.get(search_endpoint, headers=headers)
        # PUT
        elif method == "put":
            response = requests.put(search_endpoint, headers=headers, json=body)
        # DELETE
        elif method == "delete":
            response = requests.delete(search_endpoint, headers=headers)
            status_code = response.status_code
            logging.info(
                f"[call_search_api] Successfully called DELETE on {resource_type} "
                f"for {resource_name} with status code: {status_code}."
            )

        if response is not None:
            status_code = response.status_code
            if status_code >= 400:
                logging.warning(
                    f"[call_search_api] Unsucessfully called {method} on {resource_type} "
                    f"for {resource_name} with status code: {status_code}. Reason: {response.reason}."
                )
                try:
                    response_text_dict = json.loads(response.text)
                    error_message = response_text_dict["error"]["message"]
                    logging.warning(f"[call_search_api] Message: '{error_message}'.")
                except json.JSONDecodeError:
                    logging.warning(
                        f"[call_search_api] Response is not valid JSON. Raw response: '{response.text}'."
                    )
            else:
                logging.info(
                    f"[call_search_api] Successfully called {method} on {resource_type} "
                    f"for {resource_name} with status code: {status_code}."
                )

    except Exception as e:
        logging.error(
            f"[call_search_api] Error when calling {method} on {resource_type} "
            f"for {resource_name}. Error: '{e}'."
        )


def get_function_key(
    subscription_id: str,
    resource_group: str,
    function_app_name: str,
    credential: TokenCredential,
) -> str | None:
    """
    Returns an API key for the given function.

    Args:
        subscription_id (str): The subscription ID of the Azure subscription.
        resource_group (str): The name of the resource group containing the solution resources.
        function_app_name (str): The name of the function app.
        credential (TokenCredential): An instance of a TokenCredential class that can provide an access token.

    Returns:
        str | None: A unique key for the function. Returns None if the key could not be obtained.
    """
    logging.info(
        "[get_function_key] Obtaining function key after creating or updating its value."
    )
    # Get the token
    access_token: AccessToken = credential.get_token(
        "https://management.azure.com/.default"
    )
    token = access_token.token
    # Get key
    request_url = f"https://management.azure.com/subscriptions/{subscription_id}/resourceGroups/{resource_group}/providers/Microsoft.Web/sites/{function_app_name}/functions/document_chunking/keys/mykey?api-version=2022-03-01"
    headers = {
        "Authorization": f"Bearer {token}",
        "Content-Type": "application/json",
    }
    data = {
        "properties": {
            "name": "mykey"  # omit the "value" field to let Azure generate the key
        }
    }

    response = requests.put(request_url, headers=headers, json=data)
    response_json = json.loads(response.content.decode("utf-8"))
    # print(response_json)
    try:
        function_key = response_json["properties"]["value"]
    except Exception as e:
        function_key = None
        logging.error(f"[get_function_key] Error when getting function key. Error: {e}")
    return function_key


def approve_private_link_connections(
    bearer_token: str,
    subscription_id: str,
    resource_group: str,
    service_name: str,
    service_type: str,
    api_version: str,
) -> None:
    """
    Approves private link service connections for a given service using
    the GET-then-PUT pattern to ensure all required fields are present.

    Args:
        bearer_token (str): The bearer token used for authorization.
        subscription_id (str): The subscription ID of the Azure subscription.
        resource_group (str): The name of the resource group containing the solution resources.
        service_name (str): The name of the service.
        service_type (str):  The type of the service (e.g., 'Microsoft.Storage/storageAccounts').
        api_version (str): The API version.

    Returns:
        None

    Note:
        Instead of raising an exception on errors, we log a warning.
        This updated version performs a GET-then-PUT for each connection
        to avoid 'InvalidValuesForRequestParameters' errors.
    """
    logging.info(
        f"[approve_private_link_connections] Bearer token: {bearer_token[:5]}...{bearer_token[-5:]}."
    )
    logging.info(
        f"[approve_private_link_connections] Subscription ID: {subscription_id}."
    )
    logging.info(
        f"[approve_private_link_connections] Resource group: {resource_group}."
    )
    logging.info(f"[approve_private_link_connections] Service name: {service_name}.")
    logging.info(f"[approve_private_link_connections] Service type: {service_type}.")
    logging.info(f"[approve_private_link_connections] API version: {api_version}.")

    # List all private endpoint connections for the given resource
    list_url = (
        f"https://management.azure.com/subscriptions/{subscription_id}"
        f"/resourceGroups/{resource_group}/providers/{service_type}/{service_name}"
        f"/privateEndpointConnections?api-version={api_version}"
    )
    logging.debug(f"[approve_private_link_connections] Request URL: {list_url}")

    headers = {"Authorization": bearer_token, "Content-Type": "application/json"}

    try:
        response = requests.get(list_url, headers=headers)
        response.raise_for_status()
        response_json = response.json()

        if "value" not in response_json:
            logging.error(
                f"Unexpected response structure when fetching private link connections. "
                f"Response content: {response.content}."
            )
            return  # no connections to approve or error in structure

        # Iterate over all connections in the array
        for connection in response_json["value"]:
            connection_id = connection["id"]  # full ARM ID
            connection_name = connection["name"]
            status = connection["properties"]["privateLinkServiceConnectionState"][
                "status"
            ]
            logging.info(
                f"[approve_private_link_connections] Checking connection '{connection_name}'. Status: {status}."
            )

            # Approve only if status is "Pending"
            if status.lower() == "pending":
                # 1. GET the entire connection resource so we can PUT it back intact
                single_connection_url = f"https://management.azure.com{connection_id}?api-version={api_version}"
                logging.debug(
                    f"[approve_private_link_connections] GET single connection URL: {single_connection_url}"
                )
                try:
                    single_conn_response = requests.get(
                        single_connection_url, headers=headers
                    )
                    single_conn_response.raise_for_status()
                    full_conn_resource = single_conn_response.json()
                except requests.HTTPError as http_err:
                    logging.warning(
                        f"Failed to GET full connection resource for '{connection_name}': {http_err}. "
                        f"Response: {single_conn_response.text if 'single_conn_response' in locals() else ''}"
                    )
                    continue

                # 2. Update the status to "Approved" within the retrieved resource
                full_conn_resource["properties"]["privateLinkServiceConnectionState"][
                    "status"
                ] = "Approved"
                full_conn_resource["properties"]["privateLinkServiceConnectionState"][
                    "description"
                ] = "Approved by setup script"

                # 3. PUT the entire resource (with updated status)
                logging.debug(
                    f"[approve_private_link_connections] PUT single connection URL: {single_connection_url}"
                )
                approve_response = requests.put(
                    single_connection_url,
                    headers=headers,
                    json=full_conn_resource,
                )

                if approve_response.status_code in [200, 202]:
                    logging.info(
                        f"Approved private endpoint connection '{connection_name}' for service '{service_name}'."
                    )
                else:
                    logging.warning(
                        f"Warning: Failed to approve private endpoint connection '{connection_name}' "
                        f"for service '{service_name}'. Status Code: {approve_response.status_code}, "
                        f"Response: {approve_response.text}."
                    )

            elif status.lower() == "approved":
                logging.info(
                    f"[approve_private_link_connections] Connection '{connection_name}' is already Approved. Skipping re-approval."
                )
                continue

    except requests.HTTPError as http_err:
        logging.warning(
            f"HTTP error occurred when listing/approving private link connections: {http_err}. "
            f"Response: {response.text}."
        )
    except Exception as e:
        logging.warning(
            f"Error occurred when approving private link connections: Error: {e}"
        )


def approve_search_shared_private_access(
    subscription_id: str,
    resource_group: str,
    storage_resource_group: str,
    aoai_resource_group: str,
    function_app_name: str,
    storage_account_name: str,
    openai_service_name: str,
    credential: TokenCredential,
) -> None:
    """
    Approves Shared Private Access requests for private endpoints for Azure AI Search,
    Storage Account, Function App, and Azure OpenAI Service.

    Args:
        subscription_id (str): The subscription ID of the Azure subscription.
        resource_group (str): The name of the resource group containing the solution resources.
        storage_resource_group (str): The resource group name where the  Storage Account is located.
        aoai_resource_group (str): The resource group name where the Azure OpenAI service is located.
        function_app_name (str): The name of the Function App.
        storage_account_name (str):  The name of the Storage Account.
        openai_service_name (str): The name of the Azure OpenAI service.
        credential (TokenCredential): An instance of a TokenCredential class that can provide an access token.

    Returns:
        None

    Raises:
        Exception: If approval fails.
    """
    try:
        logging.info(
            "[approve_search_shared_private_access] Approving Shared Private Access requests for "
            "Storage Account, Function App, and Azure OpenAI Service if needed."
        )

        # Obtain the access token
        try:
            access_token: AccessToken = credential.get_token(
                "https://management.azure.com/.default"
            )
            access_token = f"Bearer {access_token.token}"
            logging.info("Obtained access token successfully.")
        except ClientAuthenticationError as e:
            logging.error(f"Authentication failed when obtaining access token: {e}")
            raise
        except Exception as e:
            logging.error(f"Unexpected error when obtaining access token: {e}")
            raise

        # Approve private link connection for Storage Account
        try:
            approve_private_link_connections(
                access_token,
                subscription_id,
                storage_resource_group,
                storage_account_name,
                "Microsoft.Storage/storageAccounts",
                "2023-01-01",
            )
            logging.info(
                f"[approve_private_link_connections] Approved private link connections for Storage Account: {storage_account_name}."
            )
        except Exception as e:
            logging.error(
                f"Failed to approve private link connections for Storage Account '{storage_account_name}': {e}"
            )
            raise

        # Approve private link connection for Function App
        try:
            approve_private_link_connections(
                access_token,
                subscription_id,
                resource_group,
                function_app_name,
                "Microsoft.Web/sites",
                "2022-09-01",
            )
            logging.info(
                f"[approve_private_link_connections] Approved private link connections for Function App: {function_app_name}."
            )
        except Exception as e:
            logging.error(
                f"Failed to approve private link connections for Function App '{function_app_name}': {e}"
            )
            raise

        # Approve private link connection for Azure OpenAI Service
        try:
            approve_private_link_connections(
                access_token,
                subscription_id,
                aoai_resource_group,
                openai_service_name,
                "Microsoft.CognitiveServices/accounts",
                "2022-10-01",
            )
            logging.info(
                f"Approved private link connections for Azure OpenAI Service: {openai_service_name}."
            )
        except Exception as e:
            logging.error(
                f"Failed to approve private link connections for Azure OpenAI Service '{openai_service_name}': {e}"
            )
            raise

    except Exception as e:
        logging.error(
            f"Error when approving private link service connection. Please do it manually. Error: {e}"
        )
        raise


def execute_setup(
    subscription_id: str,
    resource_group: str,
    function_app_name: str,
    search_principal_id: str,
    azure_search_use_mis: bool,
    enable_managed_identities: bool,
    enable_env_credentials: bool,
):
    """
    This function performs the necessary steps to set up the ingestion sub components, such as creating the required datastores and indexers.

    Args:
        subscription_id (str): The subscription ID of the Azure subscription.
        resource_group (str): The name of the resource group containing the solution resources.
        function_app_name (str): The name of the Function App.
        search_principal_id (str): The principal ID of the search managed identity.
        azure_search_use_mis (bool): Whether to use Search Service Managed Identity to conect to the Data Ingestion function.
        enable_managed_identities (bool):
        enable_env_credentials (bool):

    Returns:
        None
    """
    logging.info(
        f"[execute_setup] Getting function app {function_app_name} properties."
    )
    credential = ChainedTokenCredential(
        ManagedIdentityCredential(), AzureCliCredential()
    )
    web_mgmt_client = WebSiteManagementClient(credential, subscription_id)
    function_app_settings = web_mgmt_client.web_apps.list_application_settings(
        resource_group, function_app_name
    )

    function_endpoint = f"https://{function_app_name}.azurewebsites.net"
    azure_openai_service_name = function_app_settings.properties[
        "AZURE_OPENAI_SERVICE_NAME"
    ]
    search_service = function_app_settings.properties["AZURE_SEARCH_SERVICE"]
    search_analyzer_name = function_app_settings.properties["SEARCH_ANALYZER_NAME"]
    search_api_version = function_app_settings.properties.get(
        "SEARCH_API_VERSION", "2024-07-01"
    )
    search_index_interval = function_app_settings.properties["SEARCH_INDEX_INTERVAL"]
    search_index_name = function_app_settings.properties["SEARCH_INDEX_NAME"]
    storage_container = function_app_settings.properties["STORAGE_CONTAINER"]
    storage_account_name = function_app_settings.properties["STORAGE_ACCOUNT_NAME"]
    network_isolation = (
        True
        if function_app_settings.properties["NETWORK_ISOLATION"].lower() == "true"
        else False
    )
    storage_container = function_app_settings.properties["STORAGE_CONTAINER"]
    storage_account_name = function_app_settings.properties["STORAGE_ACCOUNT_NAME"]
    azure_openai_embedding_deployment = function_app_settings.properties.get(
        "AZURE_OPENAI_EMBEDDING_DEPLOYMENT", "text-embedding"
    )
    azure_openai_embedding_model = function_app_settings.properties.get(
        "AZURE_OPENAI_EMBEDDING_MODEL", "text-embedding-3-large"
    )
    azure_embeddings_vector_size = function_app_settings.properties.get(
        "AZURE_EMBEDDINGS_VECTOR_SIZE", "3072"
    )
    azure_storage_resource_group = function_app_settings.properties[
        "AZURE_STORAGE_ACCOUNT_RG"
    ]
    azure_aoai_resource_group = function_app_settings.properties["AZURE_AOAI_RG"]

    logging.info(f"[execute_setup] Function endpoint: {function_endpoint}")
    logging.info(f"[execute_setup] Search service: {search_service}")
    logging.info(f"[execute_setup] Search analyzer name: {search_analyzer_name}")
    logging.info(f"[execute_setup] Search API version: {search_api_version}")
    logging.info(f"[execute_setup] Search index interval: {search_index_interval}")
    logging.info(f"[execute_setup] Search index name: {search_index_name}")
    logging.info(f"[execute_setup] Storage container: {storage_container}")
    logging.info(f"[execute_setup] Storage account name: {storage_account_name}")
    logging.info(
        f"[execute_setup] Embedding deployment name: {azure_openai_embedding_deployment}"
    )
    logging.info(f"[execute_setup] Embedding model: {azure_openai_embedding_model}")
    logging.info(
        f"[execute_setup] Embedding vector size: {azure_embeddings_vector_size}"
    )
    logging.info(f"[execute_setup] Resource group: {resource_group}")
    logging.info(
        f"[execute_setup] Storage resource group: {azure_storage_resource_group}"
    )
    logging.info(
        f"[execute_setup] Azure OpenAI resource group: {azure_aoai_resource_group}"
    )

    # NL2SQL Elements
    storage_container_nl2sql = "nl2sql"
    search_index_name_nl2sql_queries = "nl2sql-queries"
    search_index_name_nl2sql_tables = "nl2sql-tables"

    logging.info(
        f"[execute_setup] NL2SQL Storage container: `{storage_container_nl2sql}`."
    )
    logging.info(
        f"[execute_setup] NL2SQL Search index name (queries): `{search_index_name_nl2sql_queries}`."
    )
    logging.info(
        f"[execute_setup] NL2SQL Search index name (tables): `{search_index_name_nl2sql_tables}`."
    )

    ###################################################################
    # 1. Get function key to be used later when creating the skillset #
    ###################################################################

    logging.info("Getting function key.")
    function_key = get_function_key(
        subscription_id, resource_group, function_app_name, credential
    )
    if function_key is None:
        logging.error(
            f"Could not get function key. Please make sure the function {function_app_name}/document_chunking is deployed before running this script."
        )
        exit(1)

    ######################################################
    # 2. Approve Search Shared Private Links (if needed) #
    ######################################################

    logging.info("Approving search shared private links.")
    approve_search_shared_private_access(
        subscription_id,
        resource_group,
        azure_storage_resource_group,
        azure_aoai_resource_group,
        function_app_name,
        storage_account_name,
        azure_openai_service_name,
        credential,
    )

    ###############################
    # 3. Creating blob containers #
    ###############################

    # NOTE: This step was removed since the storage account and container are already created by `azd provision`.

    ####################################
    # 4. Creating AI Search datasource #
    ####################################

    logging.info("Creating datasources.")
    start_time = time.time()

    def create_datasource(
        search_service: str,
        search_api_version: str,
        datasource_name: str,
        storage_connection_string: str,
        container_name: str,
        credential: TokenCredential,
        subfolder: str | None = None,
    ) -> None:
        """
        Creates an Azure AI Search datasource.

        Args:
            search_service (str): The name of the Azure AI Search service.
            search_api_version (str): The version of the Azure AI Search API.
            datasource_name (str): The name to give the datasource.
            storage_connection_string (str): The connection string to connect to the Storage Account.
            container_name (str): The name of the container.
            credential (TokenCredential): An instance of a TokenCredential class that can provide an access token.
            subfolder (str | None, optional): The subfolder to use. Defaults to None.

        Returns:
            None
        """
        body = {
            "description": f"Datastore for {datasource_name}",
            "type": "azureblob",
            "dataDeletionDetectionPolicy": {
                "@odata.type": "#Microsoft.Azure.Search.NativeBlobSoftDeleteDeletionDetectionPolicy"
            },
            "credentials": {"connectionString": storage_connection_string},
            "container": {
                "name": container_name,
                "query": (
                    f"{subfolder}/" if subfolder else ""
                ),  # adding subfolder path if provided
            },
        }
        call_search_api(
            search_service,
            search_api_version,
            "datasources",
            f"{datasource_name}-datasource",
            "put",
            credential,
            body,
        )

    # Define storage connection string without account key
    # TODO: Use storage account resource group
    storage_connection_string = f"ResourceId=/subscriptions/{subscription_id}/resourceGroups/{azure_storage_resource_group}/providers/Microsoft.Storage/storageAccounts/{storage_account_name}/;"

    # Creating main datasource
    create_datasource(
        search_service,
        search_api_version,
        search_index_name,
        storage_connection_string,
        storage_container,
        credential,
    )

    # Creating NL2SQL datasource in different subfolders
    nl2sql_subfolders = {
        "queries": search_index_name_nl2sql_queries,
        "tables": search_index_name_nl2sql_tables,
    }

    for subfolder, index_name in nl2sql_subfolders.items():
        create_datasource(
            search_service,
            search_api_version,
            index_name,
            storage_connection_string,
            "nl2sql",
            credential,
            subfolder=subfolder,
        )

    response_time = time.time() - start_time
    logging.info(f"Datastores created in {format_timespan(response_time)}.")

    #######################
    # 5. Creating indexes #
    #######################

    logging.info("Creating indexes.")
    start_time = time.time()

    def create_index_body(
        index_name: str,
        fields: list,
        content_fields_name: list[str],
        keyword_field_name: str,
        vector_profile_name: str = "myHnswProfile",
        vector_algorithm_name: str = "myHnswConfig",
    ) -> dict:
        """
        Create the request body to create a search index.

        Args:
            index_name (str): The name of the search index to be created.
            fields (list): The fields in the search index
            content_fields_name (list[str]): The fields to use for semantic ranking.
                Usually contains rich descriptive content or long-form text
            keyword_field_name (str): The field to perform direct keyword matching on.
                Usually contains key terms or short values likely to match important user query terms.
            vector_profile_name (str, optional): The name for the vector profile. Defaults to "myHnswProfile".
            vector_algorithm_name (str, optional): The name for the vector algorithm. Defaults to "myHnswConfig".

        Returns:
            dict: The request body.
        """
        body = {
            "name": index_name,
            "fields": fields,
            "corsOptions": {"allowedOrigins": ["*"], "maxAgeInSeconds": 60},
            "vectorSearch": {
                "profiles": [
                    {"name": vector_profile_name, "algorithm": vector_algorithm_name}
                ],
                "algorithms": [
                    {
                        "name": vector_algorithm_name,
                        "kind": "hnsw",
                        "hnswParameters": {
                            "m": 4,
                            "efConstruction": 400,
                            "efSearch": 500,
                            "metric": "cosine",
                        },
                    }
                ],
            },
            "semantic": {
                "configurations": [
                    {
                        "name": "my-semantic-config",
                        "prioritizedFields": {
                            "prioritizedContentFields": [
                                {"fieldName": field_name}
                                for field_name in content_fields_name
                            ],
                            "prioritizedKeywordsFields": [
                                {"fieldName": keyword_field_name}
                            ],
                        },
                    }
                ]
            },
        }
        return body

    # Common vector search configurations
    vector_profile_name = "myHnswProfile"
    vector_algorithm_name = "myHnswConfig"

    # Define index configurations
    indices = [
        {
            "index_name": search_index_name,  # RAG index
            "fields": [
                {
                    "name": "id",
                    "type": "Edm.String",
                    "key": True,
                    "analyzer": "keyword",
                    "searchable": True,
                    "retrievable": True,
                },
                {
                    "name": "parent_id",
                    "type": "Edm.String",
                    "searchable": False,
                    "retrievable": True,
                },
                {
                    "name": "metadata_storage_path",
                    "type": "Edm.String",
                    "searchable": False,
                    "sortable": False,
                    "filterable": False,
                    "facetable": False,
                },
                {
                    "name": "metadata_storage_name",
                    "type": "Edm.String",
                    "searchable": False,
                    "sortable": False,
                    "filterable": False,
                    "facetable": False,
                },
                {
                    "name": "metadata_storage_last_modified",
                    "type": "Edm.DateTimeOffset",
                    "searchable": False,
                    "sortable": True,
                    "retrievable": True,
                    "filterable": True,
                },
                {
                    "name": "metadata_security_id",
                    "type": "Collection(Edm.String)",
                    "searchable": False,
                    "retrievable": True,
                    "filterable": True,
                },
                {
                    "name": "chunk_id",
                    "type": "Edm.Int32",
                    "searchable": False,
                    "retrievable": True,
                },
                {
                    "name": "content",
                    "type": "Edm.String",
                    "searchable": True,
                    "retrievable": True,
                    "analyzer": search_analyzer_name,
                },
                {
                    "name": "imageCaptions",
                    "type": "Edm.String",
                    "searchable": True,
                    "retrievable": True,
                    "analyzer": search_analyzer_name,
                },
                {
                    "name": "page",
                    "type": "Edm.Int32",
                    "searchable": False,
                    "retrievable": True,
                },
                {
                    "name": "offset",
                    "type": "Edm.Int64",
                    "filterable": False,
                    "searchable": False,
                    "retrievable": True,
                },
                {
                    "name": "length",
                    "type": "Edm.Int32",
                    "filterable": False,
                    "searchable": False,
                    "retrievable": True,
                },
                {
                    "name": "title",
                    "type": "Edm.String",
                    "filterable": True,
                    "searchable": True,
                    "retrievable": True,
                    "analyzer": search_analyzer_name,
                },
                {
                    "name": "category",
                    "type": "Edm.String",
                    "filterable": True,
                    "searchable": True,
                    "retrievable": True,
                    "analyzer": search_analyzer_name,
                },
                {
                    "name": "filepath",
                    "type": "Edm.String",
                    "filterable": False,
                    "searchable": False,
                    "retrievable": True,
                },
                {
                    "name": "url",
                    "type": "Edm.String",
                    "filterable": False,
                    "searchable": False,
                    "retrievable": True,
                },
                {
                    "name": "summary",
                    "type": "Edm.String",
                    "filterable": False,
                    "searchable": True,
                    "retrievable": True,
                },
                {
                    "name": "relatedImages",
                    "type": "Collection(Edm.String)",
                    "filterable": False,
                    "searchable": False,
                    "retrievable": True,
                },
                {
                    "name": "relatedFiles",
                    "type": "Collection(Edm.String)",
                    "filterable": False,
                    "searchable": False,
                    "retrievable": True,
                },
                {
                    "name": "source",
                    "type": "Edm.String",
                    "searchable": False,
                    "retrievable": True,
                    "filterable": True,
                },
                {
                    "name": "contentVector",
                    "type": "Collection(Edm.Single)",
                    "searchable": True,
                    "retrievable": True,
                    "dimensions": azure_embeddings_vector_size,
                    "vectorSearchProfile": vector_profile_name,
                },
                {
                    "name": "captionVector",
                    "type": "Collection(Edm.Single)",
                    "searchable": True,
                    "retrievable": True,
                    "dimensions": azure_embeddings_vector_size,
                    "vectorSearchProfile": vector_profile_name,
                },
            ],
            "content_fields_name": ["content", "imageCaptions"],
            "keyword_field_name": "category",
        },
        {
            "index_name": search_index_name_nl2sql_queries,
            "fields": [
                {
                    "name": "id",
                    "type": "Edm.String",
                    "key": True,
                    "searchable": False,
                    "filterable": False,
                    "sortable": False,
                    "facetable": False,
                },
                {
                    "name": "datasource",
                    "type": "Edm.String",
                    "searchable": True,
                    "filterable": True,
                    "retrievable": True,
                    "sortable": False,
                    "facetable": False,
                    "analyzer": "standard.lucene",
                },
                {
                    "name": "question",
                    "type": "Edm.String",
                    "searchable": True,
                    "filterable": False,
                    "retrievable": True,
                    "sortable": False,
                    "facetable": False,
                    "analyzer": "standard.lucene",
                },
                {
                    "name": "query",
                    "type": "Edm.String",
                    "searchable": False,
                    "filterable": False,
                    "retrievable": True,
                    "sortable": False,
                    "facetable": False,
                },
                {
                    "name": "selected_tables",
                    "type": "Collection(Edm.String)",
                    "searchable": False,
                    "filterable": False,
                    "retrievable": True,
                },
                {
                    "name": "selected_columns",
                    "type": "Collection(Edm.String)",
                    "searchable": False,
                    "filterable": False,
                    "retrievable": True,
                },
                {
                    "name": "reasoning",
                    "type": "Edm.String",
                    "searchable": True,
                    "filterable": False,
                    "retrievable": True,
                    "sortable": False,
                    "facetable": False,
                },
                {
                    "name": "contentVector",
                    "type": "Collection(Edm.Single)",
                    "searchable": True,
                    "retrievable": True,
                    "dimensions": azure_embeddings_vector_size,
                    "vectorSearchProfile": vector_profile_name,
                },
            ],
            "content_fields_name": ["question"],
            "keyword_field_name": "question",
        },
        {
            "index_name": search_index_name_nl2sql_tables,
            "fields": [
                {
                    "name": "id",
                    "type": "Edm.String",
                    "key": True,
                    "searchable": False,
                    "filterable": False,
                    "sortable": False,
                    "facetable": False,
                },
                {
                    "name": "table_name",
                    "type": "Edm.String",
                    "searchable": True,
                    "retrievable": True,
                    "analyzer": "standard.lucene",
                },
                {
                    "name": "description",
                    "type": "Edm.String",
                    "searchable": True,
                    "retrievable": True,
                    "analyzer": "standard.lucene",
                },
                {
                    "name": "datasource",
                    "type": "Edm.String",
                    "searchable": True,
                    "retrievable": True,
                    "analyzer": "standard.lucene",
                },
                {
                    "name": "columns",
                    "type": "Collection(Edm.ComplexType)",
                    "fields": [
                        {
                            "name": "name",
                            "type": "Edm.String",
                            "searchable": True,
                            "retrievable": True,
                            "analyzer": "standard.lucene",
                        },
                        {
                            "name": "description",
                            "type": "Edm.String",
                            "searchable": True,
                            "retrievable": True,
                            "analyzer": "standard.lucene",
                        },
                    ],
                },
                {
                    "name": "contentVector",
                    "type": "Collection(Edm.Single)",
                    "searchable": True,
                    "retrievable": True,
                    "dimensions": azure_embeddings_vector_size,
                    "vectorSearchProfile": vector_profile_name,
                },
            ],
            # Use the table-level description for vectorization
            "content_fields_name": ["description"],
            "keyword_field_name": "description",
        },
    ]

    # Iterate over each index configuration and create the index
    for index in indices:
        body = create_index_body(
            index_name=index["index_name"],
            fields=index["fields"],
            content_fields_name=index["content_fields_name"],
            keyword_field_name=index["keyword_field_name"],
            vector_profile_name=vector_profile_name,
            vector_algorithm_name=vector_algorithm_name,
        )

        # Delete existing index if it exists
        call_search_api(
            search_service,
            search_api_version,
            "indexes",
            index["index_name"],
            "delete",
            credential,
        )
        # Create the index
        call_search_api(
            search_service,
            search_api_version,
            "indexes",
            index["index_name"],
            "put",
            credential,
            body,
        )

    response_time = time.time() - start_time
    logging.info(f"Indexes created in {format_timespan(response_time)}.")

    ###################################
    # 6. Creating AI Search skillsets #
    ###################################

    logging.info("Creating skillsets.")
    start_time = time.time()

    body = {
        "name": f"{search_index_name}-skillset-chunking",
        "description": "Skillset to perform document chunking",
        "skills": [
            {
                "@odata.type": "#Microsoft.Skills.Custom.WebApiSkill",
                "name": "document-chunking",
                "description": "Extract chunks from documents.",
                "httpMethod": "POST",
                "timeout": "PT230S",
                "context": "/document",
                "batchSize": 1,  # if not set, the skill will process just the last chunk (see `function_app.py`)
                "inputs": [
                    {
                        "name": "documentUrl",
                        "source": "/document/metadata_storage_path",
                    },
                    {
                        "name": "documentSasToken",
                        "source": "/document/metadata_storage_sas_token",
                    },
                    {
                        "name": "documentContentType",
                        "source": "/document/metadata_content_type",
                    },
                ],
                "outputs": [{"name": "chunks", "targetName": "chunks"}],
            }
        ],
        "indexProjections": {
            "selectors": [
                {
                    "targetIndexName": search_index_name,
                    "parentKeyFieldName": "parent_id",
                    "sourceContext": "/document/chunks/*",
                    "mappings": [
                        {
                            "name": "chunk_id",
                            "source": "/document/chunks/*/chunk_id",
                            "inputs": [],
                        },
                        {
                            "name": "offset",
                            "source": "/document/chunks/*/offset",
                            "inputs": [],
                        },
                        {
                            "name": "length",
                            "source": "/document/chunks/*/length",
                            "inputs": [],
                        },
                        {
                            "name": "page",
                            "source": "/document/chunks/*/page",
                            "inputs": [],
                        },
                        {
                            "name": "title",
                            "source": "/document/chunks/*/title",
                            "inputs": [],
                        },
                        {
                            "name": "category",
                            "source": "/document/chunks/*/category",
                            "inputs": [],
                        },
                        {
                            "name": "url",
                            "source": "/document/chunks/*/url",
                            "inputs": [],
                        },
                        {
                            "name": "relatedImages",
                            "source": "/document/chunks/*/relatedImages",
                            "inputs": [],
                        },
                        {
                            "name": "relatedFiles",
                            "source": "/document/chunks/*/relatedFiles",
                            "inputs": [],
                        },
                        {
                            "name": "filepath",
                            "source": "/document/chunks/*/filepath",
                            "inputs": [],
                        },
                        {
                            "name": "content",
                            "source": "/document/chunks/*/content",
                            "inputs": [],
                        },
                        {
                            "name": "imageCaptions",
                            "source": "/document/chunks/*/imageCaptions",
                            "inputs": [],
                        },
                        {
                            "name": "summary",
                            "source": "/document/chunks/*/summary",
                            "inputs": [],
                        },
                        {
                            "name": "source",
                            "source": "/document/chunks/*/source",
                            "inputs": [],
                        },
                        {
                            "name": "captionVector",
                            "source": "/document/chunks/*/captionVector",
                            "inputs": [],
                        },
                        {
                            "name": "contentVector",
                            "source": "/document/chunks/*/contentVector",
                            "inputs": [],
                        },
                        {
                            "name": "metadata_storage_last_modified",
                            "source": "/document/metadata_storage_last_modified",
                            "inputs": [],
                        },
                        {
                            "name": "metadata_storage_name",
                            "source": "/document/metadata_storage_name",
                            "inputs": [],
                        },
                        {
                            "name": "metadata_storage_path",
                            "source": "/document/metadata_storage_path",
                            "inputs": [],
                        },
                        {
                            "name": "metadata_security_id",
                            "source": "/document/metadata_security_id",
                            "inputs": [],
                        },
                    ],
                }
            ],
            "parameters": {"projectionMode": "skipIndexingParentDocuments"},
        },
    }
    if azure_search_use_mis:
        body["skills"][0]["uri"] = f"{function_endpoint}/api/document-chunking"
        body["skills"][0]["authResourceId"] = f"api://{search_principal_id}"
    else:
        body["skills"][0][
            "uri"
        ] = f"{function_endpoint}/api/document-chunking?code={function_key}"

    # Delete existing skillset if it exists
    call_search_api(
        search_service,
        search_api_version,
        "skillsets",
        f"{search_index_name}-skillset-chunking",
        "delete",
        credential,
    )
    # Create skillset
    call_search_api(
        search_service,
        search_api_version,
        "skillsets",
        f"{search_index_name}-skillset-chunking",
        "put",
        credential,
        body,
    )

    def create_embedding_skillset(
        skillset_name: str,
        resource_uri: str,
        deployment_id: str,
        model_name: str,
        input_field: str,
        output_field: str,
        dimensions: int,
    ) -> dict:
        """
        Creates a skillset configuration for embedding generation Azure AI Search.

        Args:
            skillset_name (str): The name of the skillset to be created.
            resource_uri (str): The URI of the Azure OpenAI resource to be used.
            deployment_id (str): The deployment ID for the model in Azure OpenAI.
            model_name (str): The name of the model to be used for embedding generation.
            input_field (str): The field in the document from which text will be taken for embeddings.
            output_field (str): The field where the generated embeddings will be stored.
            dimensions (int): The dimensionality of the generated embeddings.

        Returns:
            dict: A dictionary representing the skillset configuration for Azure AI Search.
        """
        skill = {
            "@odata.type": "#Microsoft.Skills.Text.AzureOpenAIEmbeddingSkill",
            "name": f"{skillset_name}-embedding-skill",
            "description": f"Generates embeddings for {input_field}.",
            "resourceUri": resource_uri,
            "deploymentId": deployment_id,
            "modelName": model_name,
            "dimensions": dimensions,
            "context": "/document",
            "inputs": [{"name": "text", "source": f"/document/{input_field}"}],
            "outputs": [{"name": "embedding", "targetName": output_field}],
        }

        skillset_body = {
            "name": skillset_name,
            "description": f"Skillset for generating embeddings for {skillset_name} index.",
            "skills": [skill],
        }

        return skillset_body

    # Configuration parameters
    resource_uri = f"https://{azure_openai_service_name}.openai.azure.com/"
    deployment_id = azure_openai_embedding_deployment
    model_name = azure_openai_embedding_model

    # Creating skill sets for the NL2SQL indexes

    # Define skillsets configurations
    skillsets = [
        {
            "skillset_name": "queries-skillset",
            "input_field": "question",
            "output_field": "contentVector",
        },
        {
            "skillset_name": "tables-skillset",
            "input_field": "description",
            "output_field": "contentVector",
        },
    ]

    # Iterate and create skillsets
    for skillset in skillsets:
        body = create_embedding_skillset(
            resource_uri=resource_uri,
            deployment_id=deployment_id,
            model_name=model_name,
            dimensions=azure_embeddings_vector_size,
            **skillset,
        )

        # Delete existing skillset if it exists
        call_search_api(
            search_service,
            search_api_version,
            "skillsets",
            skillset["skillset_name"],
            "delete",
            credential,
        )
        # Create the new skillset
        call_search_api(
            search_service,
            search_api_version,
            "skillsets",
            skillset["skillset_name"],
            "put",
            credential,
            body,
        )

        logging.info(f"Skillset '{skillset['skillset_name']}' created successfully.")

    response_time = time.time() - start_time
    logging.info(f"Skillsets created in {format_timespan(response_time)}.")

    ##############################
    ##### 7. Create indexers #####
    ##############################

    logging.info("Creating indexer.")
    start_time = time.time()

    body = {
        "dataSourceName": f"{search_index_name}-datasource",
        "targetIndexName": search_index_name,
        "skillsetName": f"{search_index_name}-skillset-chunking",
        "schedule": {"interval": search_index_interval},
        "fieldMappings": [
            {
                "sourceFieldName": "metadata_storage_path",
                "targetFieldName": "id",
                "mappingFunction": {"name": "fixedLengthEncode"},
            }
        ],
        "outputFieldMappings": [],
        "parameters": {
            "batchSize": 1,
            "maxFailedItems": -1,
            "maxFailedItemsPerBatch": -1,
            "base64EncodeKeys": True,
            "configuration": {"dataToExtract": "allMetadata"},
        },
    }

    if network_isolation:
        body["parameters"]["configuration"]["executionEnvironment"] = "private"

    # Create the new indexer
    call_search_api(
        search_service,
        search_api_version,
        "indexers",
        f"{search_index_name}-indexer-chunk-documents",
        "put",
        credential,
        body,
    )

    # Creating indexers for the NL2SQL indexes
    def create_indexer_body(
        indexer_name: str,
        index_name: str,
        data_source_name: str,
        skillset_name: str,
        field_mappings: list | None = None,
        indexing_parameters: dict | None = None,
    ) -> dict:
        body = {
            "name": indexer_name,
            "dataSourceName": data_source_name,
            "targetIndexName": index_name,
            "skillsetName": skillset_name,
            "schedule": {"interval": "PT2H"},  # adjust as needed
            "fieldMappings": field_mappings if field_mappings else [],
            "outputFieldMappings": [
                {
                    "sourceFieldName": "/document/contentVector",
                    "targetFieldName": "contentVector",
                }
            ],
            "parameters": {"configuration": {"parsingMode": "json"}},
        }
        if indexing_parameters:
            body["parameters"] = indexing_parameters
        return body

    # Define field mappings for the `queries-indexer`
    field_mappings_queries = [
        {
            "sourceFieldName": "metadata_storage_path",
            "targetFieldName": "id",
            "mappingFunction": {"name": "fixedLengthEncode"},
        },
        {"sourceFieldName": "datasource", "targetFieldName": "datasource"},
        {"sourceFieldName": "question", "targetFieldName": "question"},
        {"sourceFieldName": "query", "targetFieldName": "query"},
        {"sourceFieldName": "selected_tables", "targetFieldName": "selected_tables"},
        {"sourceFieldName": "selected_columns", "targetFieldName": "selected_columns"},
        {"sourceFieldName": "reasoning", "targetFieldName": "reasoning"},
    ]

    # Define field mappings for the `tables-indexer`
    field_mappings_tables = [
        {
            "sourceFieldName": "table",
            "targetFieldName": "id",
            "mappingFunction": {"name": "fixedLengthEncode"},
        },
        {"sourceFieldName": "table", "targetFieldName": "table_name"},
        {"sourceFieldName": "description", "targetFieldName": "description"},
        {"sourceFieldName": "datasource", "targetFieldName": "datasource"},
        {"sourceFieldName": "columns", "targetFieldName": "columns"},
    ]

    # Define indexing parameters for the `queries-indexer`
    indexing_parameters = {"configuration": {"parsingMode": "json"}}

    # Define indexers configurations
    indexers = [
        {
            "indexer_name": "queries-indexer",
            "index_name": search_index_name_nl2sql_queries,
            "data_source_name": f"{search_index_name_nl2sql_queries}-datasource",
            "skillset_name": "queries-skillset",
            "field_mappings": field_mappings_queries,
            "indexing_parameters": indexing_parameters,
        },
        {
            "indexer_name": "tables-indexer",
            "index_name": search_index_name_nl2sql_tables,
            "data_source_name": f"{search_index_name_nl2sql_tables}-datasource",
            "skillset_name": "tables-skillset",
            "field_mappings": field_mappings_tables,
            "indexing_parameters": indexing_parameters,
        },
    ]

    # Iterate and create indexers
    for indexer in indexers:
        body = create_indexer_body(**indexer)

        # Delete existing indexer if it exists
        call_search_api(
            search_service,
            search_api_version,
            "indexers",
            indexer["indexer_name"],
            "delete",
            credential,
        )
        # Create the new indexer
        call_search_api(
            search_service,
            search_api_version,
            "indexers",
            indexer["indexer_name"],
            "put",
            credential,
            body,
        )

        logging.info(f"Indexer '{indexer['indexer_name']}' created successfully.")

    response_time = time.time() - start_time
    logging.info(f"Indexers created in {format_timespan(response_time)}.")


def main(
    subscription_id: str | None = None,
    resource_group: str | None = None,
    function_app_name: str | None = None,
    search_principal_id: str = "",
    azure_search_use_mis: bool = False,
    enable_env_credentials: bool = False,
    enable_managed_identities: bool = False,
):
    """
    Sets up a chunking Function App in Azure.

    Args:
        subscription_id (str | None, optional): The subscription ID of the Azure subscription.
            If not provided, the user will be prompted to enter it. Defaults to None.
        resource_group (str | None, optional): The name of the resource group containing the solution resources.
            If not provided, the user will be prompted to enter it. Defaults to None.
        function_app_name (str | None, optional): The name of the Function App. Defaults to None.
        search_principal_id (str, optional): The principal ID of the search managed identity. Defaults to "".
        azure_search_use_mis (bool, optional): Whether to use Search Service Managed Identity to conect to the Data Ingestion function. Defaults to False.
        enable_env_credentials (bool, optional): Whether to use environment credentials to run the setup. Defaults to False.
        enable_managed_identities (bool, optional): Whether to use Virtual Machine's managed identities to run the setup. Defaults to False.
    """
    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
    )
    logging.info("Starting setup.")
    if subscription_id is None:
        subscription_id = input("Enter subscription ID: ")
    if resource_group is None:
        resource_group = input("Enter function app resource group: ")
    if function_app_name is None:
        function_app_name = input("Enter chunking function app name: ")

    start_time = time.time()
    execute_setup(
        subscription_id,
        resource_group,
        function_app_name,
        search_principal_id,
        azure_search_use_mis,
        enable_managed_identities,
        enable_env_credentials,
    )
    response_time = time.time() - start_time
    logging.info(f"Finished setup in {format_timespan(response_time)}.")


if __name__ == "__main__":
    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
    )
    parser = argparse.ArgumentParser(
        description="Script to do the data ingestion setup for Azure AI Search.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    parser.add_argument("-s", "--subscription_id", help="Subscription ID")
    parser.add_argument("-r", "--resource_group", help="Resource group (Function App)")
    parser.add_argument("-f", "--function_app_name", help="Chunking function app name")
    parser.add_argument(
        "-a", "--search_principal_id", default="", help="Entra ID of the search service"
    )
    parser.add_argument(
        "-m",
        "--azure_search_use_mis",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="Use Search Service Managed Identity to Connect to data ingestion function",
    )
    parser.add_argument(
        "-i",
        "--enable_managed_identities",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="Use VM's managed identities for the setup",
    )
    parser.add_argument(
        "-e",
        "--enable_env_credentials",
        action=argparse.BooleanOptionalAction,
        default=False,
        help="Use environment credentials for the setup",
    )
    args = parser.parse_args()

    # Format `search_use_mis` to boolean
    search_use_mis = args.azure_search_use_mis

    # Log all arguments
    logging.info(f"[main] Subscription ID: {args.subscription_id}")
    logging.info(f"[main] Resource group: {args.resource_group}")
    logging.info(f"[main] Function app name: {args.function_app_name}")
    logging.info(f"[main] Search principal ID: {args.search_principal_id}")
    logging.info(f"[main] Azure AI Search use MIS: {search_use_mis}")
    logging.info(f"[main] Enable managed identities: {args.enable_managed_identities}")
    logging.info(
        f"[main] Enable environment credentials: {args.enable_env_credentials}"
    )

    main(
        subscription_id=args.subscription_id,
        resource_group=args.resource_group,
        function_app_name=args.function_app_name,
        search_principal_id=args.search_principal_id,
        azure_search_use_mis=search_use_mis,
        enable_managed_identities=args.enable_managed_identities,
        enable_env_credentials=args.enable_env_credentials,
    )