Create custom AI agent tools with Unity Catalog functions

Preview

This feature is in Public Preview.

Learn how to create AI agent tools using functions to execute custom logic and perform specific tasks.

What is UCFunctionToolkit

Databricks recommends using UCFunctionToolkit to create and use agent tools defined with Unity Catalog functions.

UCFunctionToolkit is a key component of the Unity Catalog AI Core Library that simplifies the process of creating and using Unity Catalog functions as tools in AI frameworks like LangChain. It serves as a bridge between Unity Catalog functions and AI agents.

Once set up, you can create Unity Catalog functions, wrap them with UCFunctionToolkit, and use them as tools in your AI agents. This abstraction allows for consistent tool creation across authoring libraries and enables features like auto-tracing of retriever functions, making it easier to view retrieved documents and related data.

See Unity Catalog’s AI Core library quickstart for more information and examples.

Define a Unity Catalog function

To use UCFunctionToolkit, you must install the Unity Catalog AI Core Library and any necessary integration packages.

Python

## Install UCToolkit for LangChain with the Databricks extra
# Install the core package for databricks
%pip install unitycatalog-ai[databricks]

# Install the LangChain integration
%pip install unitycatalog-langchain[databricks]

note

This example uses LangChain, you must install the correct package for your agent library. For example, OpenAI agents should use %pip install unitycatalog-openai[databricks]. See Unity Catalog AI integrations.

Create a Unity Catalog function containing the tool’s logic.

Python

from unitycatalog.ai.core.databricks import DatabricksFunctionClient

CATALOG = "my_catalog"
SCHEMA = "my_schema"

def add_numbers(number_1: float, number_2: float) -> float:
  """
  A function that accepts two floating point numbers adds them,
  and returns the resulting sum as a float.

  Args:
      number_1 (float): The first of the two numbers to add.
      number_2 (float): The second of the two numbers to add.

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

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

Use the function as a tool

Once you’ve created a function, you can make it accessible to AI authoring libraries by wrapping it into a toolkit. The following example uses LangChain, but a similar approach can be applied to other frameworks:

Python

from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit

# Define the UC function to be used as a tool
func_name = f"{CATALOG}.{SCHEMA}.add_numbers"

# Create a toolkit with the UC function
toolkit = UCFunctionToolkit(function_names=[func_name], client=client)

tools = toolkit.tools

For more information on UCFunctionToolKit, see Unity Catalog documentation.