Create AI agent tools

Preview

This feature is in Public Preview.

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

Agent Framework helps developers create tools that AI agents can use to perform actions beyond language generation, such as retrieving structured or unstructured data or executing 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 tools and add them to agents with Mosaic AI Agent Framework, you can use any combination of the following methods:

• Unity Catalog functions: Unity Catalog functions are defined and managed within Unity Catalog, offering built-in security and compliance features. Writing your tool as a Unity Catalog function grants easier discoverability, governance, and reuse. Unity Catalog functions work especially well for applying transformations and aggregations on large datasets.

• Agent code tools: These tools are defined in the same code that defines the AI agent. This approach is useful when calling REST APIs, using arbitrary code or libraries, or executing low-latency tools. However, this approach lacks the built-in discoverability and governance provided by Unity Catalog functions.

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

To see examples of Unity Catalog function tools and agent code tools, see Agent tool examples

Improve tool-calling with documentation

Clear and detailed documentation helps AI agents understand when and how to use the tools you provide. When creating tools, document tool parameters and return values thoroughly to ensure that your AI agent uses the tools correctly and at the right time:

For Unity Catalog functions, use COMMENT to describe the tool and parameters.

Example of 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 of ineffective tool documentation

The following example shows ineffective COMMENT strings that miss key information, such as the return values.

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;

Agent tool examples

See the following articles for examples of agent tools:

• Code interpreter tools allow agents to execute arbitrary code such as Python.

• 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.

Add Unity Catalog tools to agents

Once you create the Unity Catalog tools, add them to your agent. LangChain agents can leverage the UCFunctionToolkit to incorporate UC tools.

Export tool-calling agents from the AI Playground

The AI Playground provides a convenient way to add Unity Catalog tools to an LLM, test the agent, and export its code.

To use the AI Playground to export agents, your workspace must meet the following requirements:

• Unity Catalog must be enabled.

• Serverless compute must be enabled.

• Either Pay-per-token foundation models or External models must be enabled.

Use the following steps to export tool-calling agents code:

1. From the AI Playground, select a model with the Tools enabled label.

   

2. Select Tools and click Add a tool.

3. In the dropdown menu, select a Unity Catalog function:

   

4. Use the Playground to chat and test the current combination of LLM, tools, and system prompt. Try variations to get a feel for how the current setup works.

   

   After adding tools, export the agent to Python notebooks:

5. Click Export to generate Python notebooks that define and deploy the agent.

   After exporting the agent code, you will see three files saved to your workspace:

   • agent notebook: Contains Python code defining your agent using LangChain.

   • driver notebook: Contains Python code to log, trace, register, and deploy the AI agent using Mosaic AI Agent Framework.

   • config.yml: Contains configuration information about your agent, including tool definitions.

6. Open the agent notebook to see the LangChain code defining your agent. Use this notebook to test and iterate on the agent programmatically, such as defining more tools.

7. When you’re happy with the agent’s outputs, run the driver notebook to log and deploy your agent to a Model Serving endpoint.

Next steps

• Trace an agent.