Integrate LlamaIndex with Databricks Unity Catalog tools

Use Databricks Unity Catalog to integrate SQL and Python functions as tools in LlamaIndex workflows. This integration combines Unity Catalog governance with LlamaIndex's capabilities to index and query large datasets for LLMs.

Requirements

Install Python 3.10 or above.

Integrate Unity Catalog tools with LlamaIndex

Run the following code in a notebook or Python script to create a Unity Catalog tool and use it in a LlamaIndex agent.

Install the Databricks Unity Catalog integration package for LlamaIndex.

Python
%pip install unitycatalog-llamaindex[databricks]
dbutils.library.restartPython()

Create an instance of the Unity Catalog functions client.

Python
from unitycatalog.ai.core.base import get_uc_function_client

client = get_uc_function_client()

Create a Unity Catalog function written in Python.

Python
CATALOG = "your_catalog"
SCHEMA = "your_schema"

func_name = f"{CATALOG}.{SCHEMA}.code_function"

def code_function(code: str) -> str:
  """
  Runs Python code.

  Args:
    code (str): The Python code to run.
  Returns:
    str: The result of running the Python code.
  """
  import sys
  from io import StringIO
  stdout = StringIO()
  sys.stdout = stdout
  exec(code)
  return stdout.getvalue()

client.create_python_function(
  func=code_function,
  catalog=CATALOG,
  schema=SCHEMA,
  replace=True
)

Create an instance of theUnity Catalog function as a toolkit, and run it to verify that the tool behaves properly.

Python
from unitycatalog.ai.llama_index.toolkit import UCFunctionToolkit
import mlflow

# Enable traces
mlflow.llama_index.autolog()

# Create a UCFunctionToolkit that includes the UC function
toolkit = UCFunctionToolkit(function_names=[func_name])

# Fetch the tools stored in the toolkit
tools = toolkit.tools
python_exec_tool = tools[0]

# Run the tool directly
result = python_exec_tool.call(code="print(1 + 1)")
print(result)  # Outputs: {"format": "SCALAR", "value": "2\n"}

Use the tool in a LlamaIndex ReActAgent by defining the Unity Catalog function as part of a LlamaIndex tool collection. Then verify that the agent behaves properly by calling the LlamaIndex tool collection.

Python
from llama_index.llms.openai import OpenAI
from llama_index.core.agent import ReActAgent

llm = OpenAI()

agent = ReActAgent.from_tools(tools, llm=llm, verbose=True)

agent.chat("Please run the following python code: `print(1 + 1)`")

Requirements​

Integrate Unity Catalog tools with LlamaIndex​

Requirements

Integrate Unity Catalog tools with LlamaIndex