unity catalog

GuideUse Unity CatalogFunctions and AI

Use catalog functions as LangChain tools

UCFunctionToolkit turns Unity Catalog functions into LangChain tools. Each call fetches the function from the catalog and runs it on the agent's machine.

Register Python functions in Unity Catalog, build LangChain tools from them with UCFunctionToolkit from the unitycatalog-langchain package, hand the tools to an agent, and trace its calls with MLflow. The catalog stores each function's signature, descriptions, and source code. The agent's machine runs the code, not the server. For where functions sit among the other catalog objects, see Core concepts.

Set up the client

Start Python with the LangChain integration, LangGraph, MLflow, and the MLflow chat model for LangChain:

EXECUTOR_MAX_MEMORY_LIMIT=4096 uv run --python 3.12 --with unitycatalog-langchain==0.4.0 --with langgraph-prebuilt==1.0.8 --with mlflow==3.10.0 --with langchain-community==0.4.1 python -m asyncio

Use Python 3.12, or 3.13.3 or later. On 3.13.0 to 3.13.2, a function that raises loses its traceback: the tool returns the generic "no output was produced" notice under error instead.

The function sandbox caps each call's virtual memory at EXECUTOR_MAX_MEMORY_LIMIT megabytes (100 by default), a limit that applies only on Linux. The command sets it to 4096, the value this page's scripts run with.

Create the function client:

tools.pyFull script
from unitycatalog.ai.core.client import (
    UnitycatalogClient,
    UnitycatalogFunctionClient,
)
from unitycatalog.client import ApiClient, CatalogsApi, Configuration

config = Configuration(host="http://localhost:8080/api/2.1/unity-catalog")
api = ApiClient(config)
client = UnitycatalogFunctionClient(api_client=api)

Register the functions the agent may call

The agent decides when to call a function from its description. The docstring's summary becomes the tool's description, and each Args: line becomes that argument's description.

Create a catalog and a schema for the functions:

tools.pyFull script
uc = UnitycatalogClient(api_client=api)
await uc.create_catalog_async(name="tools", comment="Functions for apps and agents")
await uc.create_schema_async(name="pricing", catalog_name="tools")

Save two functions as pricing.py, in the folder where the session runs:

pricing.py
def order_total(quantity: int, unit_price: float, discount_pct: float = 0.0) -> float:
    """
    Calculate the total price of an order line after a percentage discount.

    Args:
        quantity: Number of units ordered.
        unit_price: Price of one unit.
        discount_pct: Discount in percent, from 0 to 100.

    Returns:
        The discounted total, rounded to two decimals.
    """
    return round(quantity * unit_price * (1 - discount_pct / 100), 2)

pricing.py
def shipping_cost(weight_kg: float, express: bool) -> float:
    """
    Calculate the shipping cost of a parcel.

    Args:
        weight_kg: Weight of the parcel in kilograms, greater than zero.
        express: Whether to ship with next-day delivery.

    Returns:
        The shipping cost, rounded to two decimals.
    """
    if weight_kg <= 0:
        raise ValueError("weight_kg must be greater than zero")
    rate = 4.5 if express else 2.0
    return round(5.0 + rate * weight_kg, 2)

Then register them:

tools.pyFull script
from pricing import order_total, shipping_cost

for func in (order_total, shipping_cost):
    await client.create_python_function_async(
        func=func, catalog="tools", schema="pricing"
    )

Choose which functions become tools

List the functions by full name to build the toolkit:

tools.pyFull script
from unitycatalog.ai.langchain.toolkit import UCFunctionToolkit

toolkit = UCFunctionToolkit(
    function_names=[
        "tools.pricing.order_total",
        "tools.pricing.shipping_cost",
    ],
    client=client,
)
tools = toolkit.tools

Each tool's name is the function's full name with __ in place of the dots:

tools.pyFull script
print([tool.name for tool in tools])
# ['tools__pricing__order_total', 'tools__pricing__shipping_cost']

A name longer than 64 characters is cut to its last 64.

To include every function in a schema, use catalog.schema.*:

tools.pyFull script
toolkit = UCFunctionToolkit(function_names=["tools.pricing.*"], client=client)

Check what the model sees

Print a tool's description and its arguments:

tools.pyFull script
order_tool = tools[0]
print(order_tool.description)
print(order_tool.args)

The description is the docstring's summary, Calculate the total price of an order line after a percentage discount. Each argument has a JSON Schema type (integer for quantity, number for unit_price) and the description from its Args: line. A default is added to the description, as in Discount in percent, from 0 to 100. (Default: 0.0).

Call a tool directly

Call a tool before you hand it to an agent, to check the function and its arguments. A call returns a JSON string:

tools.pyFull script
output = order_tool.invoke(
    {"quantity": 3, "unit_price": 19.99, "discount_pct": 10.0}
)
print(output)  # {"format": "SCALAR", "value": 53.97}
print(json.loads(output)["value"])  # 53.97

The tool's schema converts an integer for a DOUBLE parameter, so the tool accepts 20 for unit_price, which a direct client call rejects. A direct invoke also turns every argument into JSON text (default=str) before the client checks its type. So DATE, DECIMAL, BINARY, INTERVAL, and STRUCT arguments fail here, even the ones the agent path accepts.

When the function raises, the call doesn't. It returns the traceback under error, which is what the agent sees:

tools.pyFull script
shipping_tool = tools[1]
output = shipping_tool.invoke({"weight_kg": -1.0, "express": False})
print(json.loads(output)["error"])
# Traceback (most recent call last): ... ValueError: weight_kg must be greater than zero

Arguments the schema can't accept raise pydantic's ValidationError before the function runs. So do arguments the function doesn't declare (Extra inputs are not permitted):

tools.pyFull script
from pydantic import ValidationError

try:
    order_tool.invoke({"quantity": "three", "unit_price": 19.99})
except ValidationError as exc:
    print(exc)  # quantity: Input should be a valid integer ...

A function that returns a falsy value, such as 0, False, "", None, [], or {}, returns a notice instead of the value:

tools.pyFull script
output = order_tool.invoke({"quantity": 0, "unit_price": 19.99})
print(json.loads(output)["value"])
# The function execution has completed, but no output was produced. ...

When the model needs to see a zero or an empty result, return a value that isn't falsy, such as a string or a non-empty map.

Give the tools to an agent

LangChain's create_agent accepts the tools, but with the Unity Catalog function client they need an async adapter first.

Workaround: run the tools on the session's event loop

As built, toolkit.tools fail inside an agent. Running the agent with invoke fails with RuntimeError: … attached to a different loop, and with ainvoke it fails with RuntimeError: Timeout context manager should be used inside a task. LangGraph runs the tools' synchronous functions in worker threads, and the client's connection belongs to the session's event loop.

As a workaround, give each tool an asynchronous function that calls the client on the session's loop:

tools.pyFull script
def run_on_this_loop(tool):
    async def run(**arguments):
        result = await client.execute_function_async(
            tool.uc_function_name, parameters=arguments
        )
        return result.to_json()

    tool.coroutine = run
    return tool

agent_tools = [run_on_this_loop(tool) for tool in tools]

Run the agent

The agent needs a chat model in model. Set it first: point it at a gateway endpoint as in Route the model through the MLflow AI Gateway, or assign any LangChain chat model that supports tool calling. Then create the agent with the adapted tools and ask it a question:

tools.pyFull script
from langchain.agents import create_agent

agent = create_agent(model, tools=agent_tools)
result = await agent.ainvoke(
    {
        "messages": [
            {
                "role": "user",
                "content": "Price 3 units at 19.99 with 10% off, shipped express at 2 kg.",
            }
        ]
    }
)
for message in result["messages"]:
    print(message.type, message.content)

The messages are your question, each tool call the model made, a tool message with each function's JSON result, and the model's answer.

Trace the agent's tool calls with MLflow

MLflow's LangChain autologging records each agent run as a trace. It needs no MLflow server and no key: the traces go to mlflow.db, a SQLite file in the session's folder. Turn it on, then run the agent again:

tools.pyFull script
import mlflow

mlflow.set_tracking_uri("sqlite:///mlflow.db")
mlflow.langchain.autolog()

Print the spans of the latest trace:

tools.pyFull script
trace = mlflow.search_traces(max_results=1, return_type="list")[0]
for span in trace.data.spans:
    print(span.span_type, span.name)
# CHAIN LangGraph
# CHAIN model
# CHAT_MODEL ...
# CHAIN tools
# TOOL tools__pricing__order_total
# ...

The trace has one span for the agent's graph, one for each model turn, and one TOOL span for each function call. A tool span records the arguments the model sent and the tool message it got back.

Route the model through the MLflow AI Gateway

The MLflow AI Gateway holds the provider key, so the agent's process needs only the gateway's URL and an endpoint name.

Save the endpoint as gateway.yaml, with your provider, model, and base URL:

gateway.yaml
endpoints:
  - name: chat
    endpoint_type: llm/v1/chat
    model:
      provider: openai
      name: gpt-4o-mini
      config:
        # The gateway reads the key from this environment variable.
        openai_api_key: $OPENAI_API_KEY
        # Any OpenAI-compatible server works here, including a local one.
        openai_api_base: https://api.openai.com/v1

In the same folder, with your provider's key in OPENAI_API_KEY, start the gateway:

start_gateway.sh
uv run --python 3.12 --with "mlflow[gateway]==3.10.0" \
  mlflow gateway start --config-path gateway.yaml --port 5050

In the Python session, point model at the gateway's chat endpoint, then run the agent:

gateway_agent.py
from langchain_community.chat_models import ChatMlflow

model = ChatMlflow(target_uri="http://localhost:5050", endpoint="chat")

Control who can call a function

The toolkit fetches each function from the server when you build a tool, and the client fetches it again on every call (toolkit, client). With authorization enabled, the server returns a function to a caller who (FunctionService.java):

  • owns the metastore or the function's catalog,
  • owns the function's schema and has USE CATALOG on its catalog, or
  • has USE CATALOG on the catalog, USE SCHEMA on the schema, and EXECUTE on the function or owns it.

Registering a function needs USE CATALOG on the catalog (or its ownership), plus USE SCHEMA and CREATE FUNCTION on the schema (or its ownership). Creating the catalog needs CREATE CATALOG on the metastore (or its ownership), and creating the schema needs USE CATALOG and CREATE SCHEMA on the catalog (or its ownership). For how these grants inherit, see Core concepts.

The function then runs on the agent's machine. By default the client runs it in a sandboxed subprocess, which restricts imports and file access and caps CPU time and memory. It runs inside the agent's own process only if you create the client with execution_mode="local".

Next steps