azure-ai-search-agent-tool(Python)

Import Notebook

Create agent tools for Azure AI Search

This notebook creates agent tools that connect to and query Azure AI Search. These tools can retrieve information from your Azure AI Search index in Databricks.

To learn more about creating agent tools that connect to external services, see Databricks documentation (AWS | Azure).

Requirements

Before you begin, ensure the following are set up:

  1. Azure AI Search
    Set up an Azure AI Search service in your Azure account. For more details, refer to the official documentation: Azure AI Search Overview.

  2. Index Creation
    Create and populate an index in Azure AI Search that your agent tools will query. Follow the instructions in the Azure Search Get Started Guide.

Setup

Install the databricks-sdk

Update the following parameters with information for your catalog and provide a name for a new connection.

  • uc_catalog: Specifies the catalog where the external function will be created.
  • uc_schema: Specifies the schema in the catalog where the external function will be created.
  • connection_name: Specifies the connection name.
3

4

5

Create a Unity Catalog Connection

Create a Unity Catalog Connection by providing the host URL and token for authentication. To learn more about supported authentication methods, see Databricks documentation (AWS | Azure).

To securely manage your authentication details, we recommend using Databricks Secrets. Replace the placeholders <secret-scope> and <secret-key> with your specific secret scope and secret key.

If you prefer, you can also create a connection through the Catalog Explorer UI. See Databricks documentation (AWS | Azure).

7

Test the connection

Test the connection works by executing a search query:

  1. Ensure the path parameter includes the correct index and API version for your use case.
  2. Adjust the json payload to match the structure and fields required for your search query.

The sample query below searches for the term lake view in the Description and Tags fields of the specified index:

9

Create a Unity Catalog function

Now that you have a working connection create a Unity Catalog function that uses this connection.

  1. Update the function name (currently search_index) and comment as you see fit.
  2. Adjust the json payload to match the structure and fields required for your search query.
11

Call the function

Call the Unity Catalog function:

  1. Update the function name (currently search_index) to match your created function.
  2. Pass in the required parameters.
13

Create a post-processing Unity Catalog function

The previous function gets the necessary data but includes additional fields (such as "@odata.context") that is not be relevant to the final response. To fix this, create another Unity Catalog function that only extracts the hotel name and description from the response's value array.

TODO:

  • Replace <index-name> with the appropriate index name for your use case.
  • Update function and field names to ensure consistency with your earlier cells.
15

Test the new post-processing function.

17

Add the tool to an AI agent

Now that you have a Unity Catalog function, you can add it to an AI agent. See Databricks documentation (AWS | Azure).

Create a tool in agent code (optional)

As an alternative, you can also define this tool directly in agent code. To learn more about the differences, see Databricks documentation (AWS | Azure)

Rewrite the Unity Catalog function to search the index as a Python function. Remember to replace index-name.

20

21