AI agent tools

Article
01/31/2025

Important

This article provides an overview of building AI agent tools using the Mosaic AI Agent Framework.

AI agent tools enable agents to perform tasks beyond language generation, such as retrieving structured or unstructured data and executing custom code.

For an introduction to AI agents, see What are compound AI systems and AI agents?.

Unity Catalog function tools vs. agent code tools

To create a tool with Mosaic AI Agent Framework, you can use any combination of the following methods:

Method	Description
Unity Catalog functions	- Defined and managed in Unity Catalog with built-in security and compliance features - Grants easier discoverability, governance, and reuse - Ideal for applying transformations and aggregations on large datasets
Agent code tools	- Defined in the AI agent’s code - Useful for calling REST APIs, using arbitrary code, or executing low-latency tools - Lacks built-in governance and discoverability of functions

Both methods are compatible with custom Python agents or agent-authoring libraries like LangGraph.

Create AI agent tools

Learn how to create AI agent tools that let your agents execute custom Python code. See Create custom AI agent tools with Unity Catalog functions.

Agent tool examples

See the following articles for examples of agent tools:

Structured data retrieval tools allow your agent to query structured data sources like SQL tables.
Unstructured data retrieval tools allow your agent to query unstructured data sources like a text corpora to perform retrieval augmented generation.
External connection tools connect agent tools to external services and APIs.

Add Unity Catalog tools to agents

Unlike agent code tools, which are defined in the agent’s code, Unity Catalog tools must be explicitly added to agents to make them available for use.

Databricks recommends using UCFunctionToolkit to integrate Unity Catalog tools with agent authoring frameworks and SDKs. See Create custom AI agent tools with Unity Catalog functions.

You can also use the AI Playground to quickly add Unity Catalog tools to agents to prototype behavior. See Prototype tool-calling agents in AI Playground.

Improve tool-calling with clear documentation

Well-documented tools help AI agents understand when and how to use tools effectively. Follow these best practices when documenting tool parameters and return values:

For Unity Catalog functions, use COMMENT to describe tool functionality and parameters.
Clearly define expected inputs and outputs.
Provide meaningful descriptions to improve usability.

Example: Effective tool documentation

The following example shows effective COMMENT strings for a Unity Catalog function tool that queries a structured table.

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 specific customer including their email and ID.'
RETURN SELECT CONCAT(
    'Customer ID: ', customer_id, ', ',
    'Customer Email: ', customer_email
  )
  FROM main.default.customer_data
  WHERE customer_name = customer_name
  LIMIT 1;

Example: Ineffective tool documentation

The following example lacks important details, making it harder for the AI agent to use the tool effectively:

CREATE OR REPLACE FUNCTION main.default.lookup_customer_info(
  customer_name STRING COMMENT 'Name of the customer.'
)
RETURNS STRING
COMMENT 'Returns info about a customer.'
RETURN SELECT CONCAT(
    'Customer ID: ', customer_id, ', ',
    'Customer Email: ', customer_email
  )
  FROM main.default.customer_data
  WHERE customer_name = customer_name
  LIMIT 1;

Share via