Azure AI Search¶

Azure AI Search is the full-text and vector store behind retrieval-augmented generation on Azure. The AzureAISearchHook reads and writes the documents of one index, from an async task: every method is a coroutine, and async Python callables are supported since Airflow 3.2. Creating and changing indexes is left to the azure-search-documents SDK.

Prerequisite Tasks¶

To use these operators, you must do a few things:

Create necessary resources using AZURE PORTAL or AZURE CLI.
Install API libraries via pip.
pip install 'apache-airflow[azure]'
Detailed information is available Installation of Airflow®
Setup Connection.

The hook uses the azure_ai_search_default connection by default, see Microsoft Azure AI Search. The index comes from the hook (index=) or from the index extra of the connection.

Writing documents¶

upload inserts documents and replaces the ones whose key exists. The documents are sent in batches of batch_size (1000 by default, the limit of the service), and the hook raises AzureAISearchIndexingError with the keys and reasons when the service rejects some of them, from the per-document results of the write.

tests/system/microsoft/azure/example_azure_ai_search.py[source]

@task
async def upload_documents() -> int:
    """Insert the documents, replacing the ones whose key already exists."""
    hook = AzureAISearchHook(index=INDEX)
    return await hook.upload(DOCUMENTS)

Searching documents¶

search is an async generator over every matching document, following the pages of the service. Select the fields you need, and filter and sort with OData expressions.

The service pages the results with skip, which stops at 100,000 documents: past that many matches the search fails. To read a larger index, sort on a sortable and filterable key (order_by=["id"]) and filter on ranges of it (filter="id gt '<last key read>'"), one range per call, as described in Paging through a large number of results.

tests/system/microsoft/azure/example_azure_ai_search.py[source]

@task(retries=3)
async def find_reports() -> list[str]:
    """Return the keys of the reports; retries while the service still indexes the write."""
    hook = AzureAISearchHook(index=INDEX)
    reports = [
        document["id"]
        async for document in hook.search(select=["id"], filter="category eq 'report'", order_by=["id"])
    ]
    if len(reports) < 2:
        raise ValueError(f"Expected 2 reports, found {reports}")
    return reports

Updating and counting documents¶

merge changes the given fields of existing documents and merge_or_upload also inserts the documents whose key is new. count returns the number of documents of the index, or of those matching a filter.

tests/system/microsoft/azure/example_azure_ai_search.py[source]

@task
async def archive_reports(keys: list[str]) -> int:
    """Change one field of the reports, then count the documents that have it."""
    hook = AzureAISearchHook(index=INDEX)
    await hook.merge([{"id": key, "category": "archive"} for key in keys])
    return await hook.count(filter="category eq 'archive'")

Reading one document¶

get_document returns the document with a key, or None when the index has none.

tests/system/microsoft/azure/example_azure_ai_search.py[source]

@task
async def read_first_document(keys: list[str]) -> dict:
    hook = AzureAISearchHook(index=INDEX)
    document = await hook.get_document(keys[0], select=["id", "title"])
    if document is None:
        raise ValueError(f"Document {keys[0]!r} is missing")
    return document

Deleting documents¶

delete removes documents by key.

tests/system/microsoft/azure/example_azure_ai_search.py[source]

@task
async def delete_documents() -> int:
    """Delete the documents by key; only the key field is needed."""
    hook = AzureAISearchHook(index=INDEX)
    return await hook.delete({"id": document["id"]} for document in DOCUMENTS)

Reference¶

For further information, look at:

Was this entry helpful?