Structured retrieval AI agent tools

Preview

This feature is in Public Preview.

This article shows how to create AI agents for structured data retrieval using the Mosaic AI Agent Framework. Structured retrievers enable agents to query structured data sources such as SQL tables.

To learn more about agent tools, see Create AI agent tools.

Table query tool

The following example creates a tool that allows an agent to query structured customer data from a Unity Catalog table.

It defines a UC function called lookup_customer_info, which allows an AI agent to retrieve structured data from a hypothetical customer_data table.

Run the following code in a SQL editor.

CREATE OR REPLACE FUNCTION main.default.lookup_customer_info(
  customer_name STRING COMMENT 'Name of the customer whose info to look up'
)
RETURNS STRING
COMMENT 'Returns metadata about a particular customer given the customer name, including the customer's email and ID. The
customer ID can be used for other queries.'
RETURN SELECT CONCAT(
    'Customer ID: ', customer_id, ', ',
    'Customer Email: ', customer_email
  )
  FROM main.default.customer_data
  WHERE customer_name = customer_name
  LIMIT 1;

Next steps

After you create an agent tool, add the tool to an AI agent. See Add Unity Catalog tools to agents.