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 systems 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.
Requirements
Databricks recommends installing the latest version of the MLflow Python client when developing agents.
Improve tool-calling with documentation
AI agents use your documentation to understand when and how to use the tools you provide. Provide clear and detailed documentation on the parameters and return values for each tool.
For Unity Catalog functions, use COMMENT
to describe the tool.
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;
Vector Search retriever
This example creates a retriever tool to query data from a Mosaic AI Vector Search index.
The Unity Catalog function databricks_docs_vector_search
queries a hypothetical vector search index containing Databricks documentation. It wraps the standard vector_search() function and uses the aliases page_content
and metadata
to match its output to the MLflow retriever schema.
Note
To conform to the MLflow retriever schema, any additional metadata columns must be added to the metadata
column using the SQL map function, rather than as top-level output keys.
Run the following code in a notebook or SQL editor.
CREATE OR REPLACE FUNCTION main.default.databricks_docs_vector_search (
-- The agent uses this comment to determine how to generate the query string parameter.
query STRING
COMMENT 'The query string for searching Databricks documentation.'
) RETURNS TABLE
-- The agent uses this comment to determine when to call this tool. It describes the types of documents and information contained within the index.
COMMENT 'Executes a search on Databricks documentation to retrieve text documents most relevant to the input query.' RETURN
SELECT
chunked_text as page_content,
map('doc_uri', url, 'chunk_id', chunk_id) as metadata
FROM
vector_search(
-- Specify your vector search index name here
index => 'catalog.schema.databricks_docs_index',
query => query,
num_results => 5
)
This retriever tool has the following caveats:
MLflow traces this Unity Catalog function as a
TOOL
span type rather than aRETRIEVER
span type. As a result, downstream Agent Framework applications like the agent review app and AI Playground will not show retriever-specific details such as links to chunks. For more information on span types, see MLflow Tracing Schema.SQL clients may limit the maximum number of rows or bytes returned. To prevent data truncation, you should truncate column values returned by the UDF. For example, you could use
substring(chunked_text, 0, 8192)
to reduce the size of large content columns and avoid row truncation during execution.Since this tool is a wrapper for the
vector_search()
function, it is subject to the same limitations as thevector_search()
function. See Limitations.
If this example is not suitable for your use case, create a vector search retriever using custom agent code instead. See Create a vector search retriever.
Assign Unity Catalog tools to agents
After you create Unity Catalog tools, assign them to your agent for use.
LangChain agents can use the UCFunctionToolkit to assign Unity Catalog tools.
The following describes how to export a LangChain tool-calling agent to code from 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 following requirements:
Unity Catalog enabled
Serverless compute enabled
Pay-per-token foundation models or External models enabled
From Playground, select a model with the Tools enabled label.
Select Tools and click Add a tool.
In the dropdown menu, select a Unity Catalog function:
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:
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.
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.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.