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.

1. Install the Databricks Unity Catalog integration package for LlamaIndex.

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

2. Create an instance of the Unity Catalog functions client.

   from unitycatalog.ai.core.base import get_uc_function_client
   
   client = get_uc_function_client()
   

3. Create a Unity Catalog function written in 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
   )
   

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

   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"}
   

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

   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)`")