unitycatalog-langchain 0.1.0rc0

Support for Unity Catalog functions as LangChain tools

🦜🔗 Using Unity Catalog AI with Langchain

Integrate Unity Catalog AI package with Langchain to allow seamless usage of UC functions as tools in agents application.

Installation

Client Library

To use this package with Open Source Unity Catalog, you will need to install:

pip install unitycatalog-ai[oss]

To use this package with Databricks Unity Catalog, you will need to install:

pip install unitycatalog-ai[databricks]

Integration Library

With the appropriate client installed for the AI functionality for Unity Catalog, you can then install the LangChain (covers LangGraph as well) integration library:

pip install unitycatalog-langchain

Getting started

Open Source Unity Catalog

Creating a Client

To interact with OSS UC, initialize the UnitycatalogFunctionClient as shown below:

import asyncio
from unitycatalog.ai.core.oss import UnitycatalogFunctionClient
from unitycatalog.client import ApiClient, Configuration

# Configure the Unity Catalog API client
config = Configuration(
    host="http://localhost:8080/api/2.1/unity-catalog"  # Replace with your UC server URL
)

# Initialize the asynchronous ApiClient
api_client = ApiClient(configuration=config)

# Instantiate the UnitycatalogFunctionClient
uc_client = UnitycatalogFunctionClient(api_client=api_client)

# Example catalog and schema names
CATALOG = "my_catalog"
SCHEMA = "my_schema"

Creating a Function in UC OSS

You can create a UC function either by providing a Python callable or by submitting a FunctionInfo object. Below is an example (recommended) of using the create_python_function API that accepts a Python callable (function) as input.

To create a UC function from a Python function, define your function with appropriate type hints and a Google-style docstring:

def add_numbers(a: float, b: float) -> float:
    """
    Adds two numbers and returns the result.

    Args:
        a (float): First number.
        b (float): Second number.

    Returns:
        float: The sum of the two numbers.
    """
    return a + b

# Create the function within the Unity Catalog catalog and schema specified
function_info = uc_client.create_python_function(
    func=add_numbers,
    catalog=CATALOG,
    schema=SCHEMA,
    replace=False,  # Set to True to overwrite if the function already exists
)

print(function_info)

Databricks-managed Unity Catalog

To use Databricks-managed Unity Catalog with this package, follow the instructions to authenticate to your workspace and ensure that your access token has workspace-level privilege for managing UC functions.

Client setup

Initialize a client for managing UC functions in a Databricks workspace, and set it as the global client.

from unitycatalog.ai.core.client import set_uc_function_client
from unitycatalog.ai.core.databricks import DatabricksFunctionClient

client = DatabricksFunctionClient(
    warehouse_id="..." # replace with the warehouse_id
)

# sets the default uc function client
set_uc_function_client(client)

Create a function in UC

Create a python UDF in Unity Catalog with the client

# replace with your own catalog and schema
CATALOG = "catalog"
SCHEMA = "schema"

func_name = f"{CATALOG}.{SCHEMA}.python_exec"
# define the function body in SQL
sql_body = f"""CREATE OR REPLACE FUNCTION {func_name}(code STRING COMMENT 'Python code to execute. Remember to print the final result to stdout.')
RETURNS STRING
LANGUAGE PYTHON
COMMENT 'Executes Python code and returns its stdout.'
AS $$
    import sys
    from io import StringIO
    stdout = StringIO()
    sys.stdout = stdout
    exec(code)
    return stdout.getvalue()
$$
"""

client.create_function(sql_function_body=sql_body)

Now the function is created and stored in the corresponding catalog and schema.

Using the Function as a GenAI Tool

Create a UCFunctionToolkit instance

Langchain tools are utilities designed to be called by a model, and UCFunctionToolkit provides the ability to use UC functions as tools that are recognized natively by LangChain.

from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit

# create a UCFunctionToolkit that includes the above UC function
toolkit = UCFunctionToolkit(function_names=[f"{CATALOG}.{SCHEMA}.python_exec"])

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

# execute the tool directly
python_exec_tool.invoke({"code": "print(1)"})

Use the tools in a Langchain Agent

Now we create an agent and use the tools.

from langchain_community.chat_models.databricks import ChatDatabricks
from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain_core.prompts import ChatPromptTemplate

# Use Databricks foundation models
llm = ChatDatabricks(endpoint="databricks-meta-llama-3-1-70b-instruct")
prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "You are a helpful assistant. Make sure to use tool for information.",
        ),
        ("placeholder", "{chat_history}"),
        ("human", "{input}"),
        ("placeholder", "{agent_scratchpad}"),
    ]
)
agent = create_tool_calling_agent(llm, tools, prompt)

# Create the agent executor
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
agent_executor.invoke({"input": "36939 * 8922.4"})

Configurations for UC functions execution

We provide configurations for databricks client to control the function execution behaviors, check function execution arguments section.