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.