Create AI agent tools

Preview

This feature is in Public Preview.

This article shows you how to create tools for AI agents using the Mosaic AI Agent Framework. To learn more, see What are compound AI system and AI agents?.

AI agents use tools to perform actions besides language generation, such as retrieving structured or unstructured data, executing code, or talking to remote services (for example, sending an email or Slack message).

To provide tools to an agent with Mosaic AI Agent Framework, you can use any combination of the following methods:

• Create tools with Unity Catalog functions: Unity Catalog functions provide easy discovery, governance, and sharing. Unity Catalog functions work especially well for applying transformations and aggregations on large datasets.

• Create tools with agent code: This approach is useful when making calls to REST APIs, using arbitrary code or libraries, or executing tools with very low latency. However, this approach lacks the built-in discoverability and governance provided by Unity Catalog functions. Weigh this tradeoff when building your agent to determine which approach is best.

Both approaches work for agents written in custom Python code or using an agent-authoring library like LangGraph.

Improve tool-calling with documentation

AI agents use your documentation to understand when and how to use the tools you provide. So, when defining tools, be sure to add clear and detailed documentation for the tool, its parameters, and its return values.

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

Requirements

• Databricks recommends installing the latest version of the MLflow Python client when developing agents. See Authentication for dependent resources for information on mlflow version requirements.

Create tools with Unity Catalog functions

The following examples show you how to create agent tools with Unity Catalog functions.

Python executor tool

This example creates a tool to execute Python code.

An LLM can use this tool to execute Python code provided by a user or written by the LLM.

Run the following code in a notebook cell. It uses the %sql notebook magic to create a Unity Catalog function called python_exec.

%sql
CREATE OR REPLACE FUNCTION
main.default.python_exec (
code STRING COMMENT 'Python code to execute. Remember to print the final result to stdout.'
)
RETURNS STRING
LANGUAGE PYTHON
DETERMINISTIC
COMMENT 'Executes Python code in the sandboxed environment and returns its stdout. The runtime is stateless, and you can not read the output of the previous tool executions. No such variables, "rows", or "observations" were defined. Calling another tool inside a Python code is NOT allowed. Use standard Python libraries only.'
AS $$
import sys
from io import StringIO
sys_stdout = sys.stdout
redirected_output = StringIO()
sys.stdout = redirected_output
exec(code)
sys.stdout = sys_stdout
return redirected_output.getvalue()
$$

Table query tool

This example creates a tool to query customer data from a Unity Catalog table.

It creates a Unity Catalog function called lookup_customer_info that an LLM can use 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 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;

Assign Unity Catalog tools to agents

After you create the Unity Catalog tools, assign the tools to your agent so it can use them.

This section describes how to create a tool-calling agent using the AI playground.

Export tool-calling agents from the AI Playground

Use the AI Playground to assign Unity Catalog tools to an LLM, test the agent, and then export the code that defines the tool-calling agent.

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

• Unity Catalog enabled

• Serverless compute enabled

• Pay-per-token foundation models or External models enabled

1. From 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:

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

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

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

Create tools with agent code

You can also create tools in the agent’s code rather than in Unity Catalog functions.

This approach is useful when calling REST APIs, using arbitrary code or libraries, or executing low-latency tools. However, defining tools in the agent’s code lacks the discoverability and governance provided by Unity Catalog functions.

For an example of a vector search retrieval tool defined in agent code, see Create a vector search retriever tool.

Next steps

• Create an agent.