Κοινή χρήση μέσω

Create custom AI agent tools with Unity Catalog functions

  • Άρθρο

Important

This feature is in Public Preview.

Learn how to create AI agent tools using Unity Catalog functions to execute custom Python code for task-specific operations.

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.

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

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:

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.

Built-in Python executor tool

Azure Databricks also provides a built-in Unity Catalog function allowing AI agents to dynamically execute Python code within queries dynamically.

system.ai.python_exec is available by default and can be used like any other Unity Catalog function-based tool.

The following example shows you how to use the python_exec function in a SQL query:

SELECT python_exec("""
import random
numbers = [random.random() for _ in range(10)]
print(numbers)
""")