Jaa


Structured retrieval AI agent tools

This article shows how to create AI agent tools for structured data retrieval using the Mosaic AI Agent Framework. To allow agents to query structured data sources such as SQL tables, you can use one of the following methods:

  • SQL Unity Catalog functions: Define a SQL query with parameters that the agent can fill in. Use this method when queries have a known, fixed format.

Query data using Unity Catalog SQL function tool

Create a structured retrieval tool using Unity Catalog SQL Function when the query is known ahead of time and the agent provides the parameters.

The following example creates a Unity Catalog 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's 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;

Genie multi-agent system

Important

This feature is in Public Preview.

Create a multi-agent system that includes a Genie agent when you don't know what kind of queries your agent needs to answer, and you need the flexibility to query multiple tables.

For more information, see Use Genie in multi-agent systems.

Next steps

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