Create agent tools using Unity Catalog functions

Use Unity Catalog functions when your agent needs to run predefined logic, such as a SQL query whose parameters the agent supplies. For other tools, start with MCPs.

This guide uses Unity Catalog AI to create a small Python function, test it, and add it to an agent using your choice of framework. If you already have a function, install the libraries in step 1 and go to step 3.

Requirements

  • A Azure Databricks notebook with Python 3.10 or above. On classic compute, use Databricks Runtime 15.0 or above.
  • Serverless compute enabled in your workspace. The Function Client uses serverless general compute to run tools; a SQL warehouse alone is not sufficient.
  • USE CATALOG and USE SCHEMA on the function's parent catalog and schema. Creating a function also requires CREATE FUNCTION on the schema; calling an existing function requires EXECUTE on it. See Function privileges.
  • A model that supports tool calling. The LangGraph example uses a Azure Databricks model endpoint. The OpenAI, Anthropic, and LlamaIndex examples use provider API credentials, as described in their tabs.

Step 1: Install the libraries

Run the installation command for your framework in a notebook cell:

Framework Installation command
LangGraph %pip install --upgrade unitycatalog-langchain[databricks] databricks-langchain langgraph
OpenAI %pip install --upgrade unitycatalog-openai[databricks]
Anthropic %pip install --upgrade unitycatalog-anthropic[databricks]
LlamaIndex %pip install --upgrade unitycatalog-llamaindex[databricks] llama-index-llms-openai

For package source code, see the Unity Catalog AI repository.

Then restart Python in a separate cell:

dbutils.library.restartPython()

Step 2: Create and test a function

Use DatabricksFunctionClient to create and test the function. Replace my_catalog and my_schema with an existing catalog and schema where you can create functions. This example creates add_numbers; choose a different name if it already exists.

from unitycatalog.ai.core.databricks import DatabricksFunctionClient

client = DatabricksFunctionClient()
CATALOG = "my_catalog"
SCHEMA = "my_schema"


def add_numbers(number_1: float, number_2: float) -> float:
    """Add two numbers.

    Args:
        number_1: The first number to add.
        number_2: The second number to add.

    Returns:
        The sum of the two numbers.
    """
    return number_1 + number_2


client.create_python_function(func=add_numbers, catalog=CATALOG, schema=SCHEMA)

result = client.execute_function(
    function_name=f"{CATALOG}.{SCHEMA}.add_numbers",
    parameters={"number_1": 2.0, "number_2": 3.0},
)
if result.error:
    raise RuntimeError(result.error)
print(result.value)  # 5.0

Step 3: Use the function in your agent

Set function_name to your function's full name, then run the example for your framework. Each uses UCFunctionToolkit to call the function directly.

from unitycatalog.ai.core.databricks import DatabricksFunctionClient

function_name = "my_catalog.my_schema.add_numbers"
client = DatabricksFunctionClient()
question = "Use the tool to add 2.0 and 3.0."

LangGraph

This example uses the databricks-claude-sonnet-4-5 endpoint. Replace it with a tool-calling model available in your workspace.

from databricks_langchain import ChatDatabricks, UCFunctionToolkit
from langgraph.prebuilt import create_react_agent

toolkit = UCFunctionToolkit(function_names=[function_name], client=client)
agent = create_react_agent(
    ChatDatabricks(endpoint="databricks-claude-sonnet-4-5"),
    tools=toolkit.tools,
)
result = agent.invoke({"messages": [{"role": "user", "content": question}]})
print(result["messages"][-1].content)

See the LangChain integration reference.

OpenAI

Set the OPENAI_API_KEY and OPENAI_MODEL environment variables for your notebook process. Use a model that supports tool calling, and store the API key in a Azure Databricks secret.

import os
from openai import OpenAI
from unitycatalog.ai.openai.toolkit import UCFunctionToolkit
from unitycatalog.ai.openai.utils import generate_tool_call_messages

toolkit = UCFunctionToolkit(function_names=[function_name], client=client)
model_client = OpenAI()
messages = [{"role": "user", "content": question}]

for _ in range(10):
    response = model_client.chat.completions.create(
        model=os.environ["OPENAI_MODEL"],
        messages=messages,
        tools=toolkit.tools,
    )
    if not response.choices[0].message.tool_calls:
        print(response.choices[0].message.content)
        break
    messages.extend(generate_tool_call_messages(response=response, client=client))
else:
    raise RuntimeError("The agent reached the tool-calling limit.")

The helper runs the requested functions and formats their results for the next model call. See the OpenAI integration reference.

Anthropic

Set the ANTHROPIC_API_KEY and ANTHROPIC_MODEL environment variables for your notebook process. Use a model that supports tool calling, and store the API key in a Azure Databricks secret.

import os
from anthropic import Anthropic
from unitycatalog.ai.anthropic.toolkit import UCFunctionToolkit
from unitycatalog.ai.anthropic.utils import generate_tool_call_messages

toolkit = UCFunctionToolkit(function_names=[function_name], client=client)
model_client = Anthropic()
messages = [{"role": "user", "content": question}]

for _ in range(10):
    response = model_client.messages.create(
        model=os.environ["ANTHROPIC_MODEL"],
        max_tokens=1024,
        tools=toolkit.tools,
        messages=messages,
    )
    if response.stop_reason != "tool_use":
        print("\n".join(block.text for block in response.content if block.type == "text"))
        break
    messages.extend(
        generate_tool_call_messages(
            response=response, client=client, conversation_history=[],
        )
    )
else:
    raise RuntimeError("The agent reached the tool-calling limit.")

The helper runs the requested functions. Append its messages to the conversation before the next model call. See the Anthropic integration reference.

LlamaIndex

This example uses an OpenAI model. Set the OPENAI_API_KEY and OPENAI_MODEL environment variables for your notebook process. Use a model that supports tool calling, and store the API key in a Azure Databricks secret.

import os
from llama_index.core.agent.workflow import FunctionAgent
from llama_index.llms.openai import OpenAI
from unitycatalog.ai.llama_index.toolkit import UCFunctionToolkit

toolkit = UCFunctionToolkit(function_names=[function_name], client=client)
agent = FunctionAgent(
    tools=toolkit.tools,
    llm=OpenAI(model=os.environ["OPENAI_MODEL"]),
)
result = await agent.run(user_msg=question)
print(result)

Run the last two lines in a notebook cell that supports await. See the LlamaIndex integration reference.

The agent calls add_numbers and returns an answer containing 5. If you use an existing function, change question to a task that function supports.

For deployment and authentication, see Build a custom agent and Agent authentication.

Local development

For testing trusted Python functions, initialize the client with execution_mode="local" and pass it to your framework's toolkit:

client = DatabricksFunctionClient(execution_mode="local")
toolkit = UCFunctionToolkit(function_names=[function_name], client=client)

Local mode retrieves the function from Unity Catalog and runs it in a local subprocess. It does not support SQL functions. Use the default serverless mode for production, and run only trusted code. See Execution modes.

For timeouts, memory limits, and serverless settings, see the local execution limits and Function Client environment variables.

Use the workspace MCP endpoint (legacy)

For existing MCP integrations, this LangGraph example connects to the same function through the legacy workspace endpoint. Enable the Managed MCP Servers preview and use an identity with access to the function.

Run %pip install --upgrade databricks-langchain langgraph "mcp>=1.24,<2" in a notebook cell, then restart Python. Replace my_catalog/my_schema/add_numbers with your function's path.

from databricks.sdk import WorkspaceClient
from databricks_langchain import (
    ChatDatabricks,
    DatabricksMCPServer,
    DatabricksMultiServerMCPClient,
)
from langchain_core.messages import convert_to_openai_messages
from langgraph.prebuilt import create_react_agent

workspace_client = WorkspaceClient()
function_path = "my_catalog/my_schema/add_numbers"
mcp_client = DatabricksMultiServerMCPClient([
    DatabricksMCPServer(
        name="uc-functions",
        url=f"{workspace_client.config.host}/api/2.0/mcp/functions/{function_path}",
        workspace_client=workspace_client,
    ),
])
agent = create_react_agent(
    ChatDatabricks(endpoint="databricks-claude-sonnet-4-5"),
    tools=await mcp_client.get_tools(),
    prompt=lambda state: convert_to_openai_messages(state["messages"]),
)
result = await agent.ainvoke({
    "messages": [{"role": "user", "content": "Use the tool to add 2.0 and 3.0."}],
})
print(result["messages"][-1].content)

For deployment, see Agent authentication. For a function that calls an external API with http_request, see Legacy REST API tools.