Databricks Unity Catalog (UC)

This notebook shows how to use UC functions as LangChain tools, with both LangChain and LangGraph agent APIs.

See Databricks documentation (AWS|Azure|GCP) to learn how to create SQL or Python functions in UC. Do not skip function and parameter comments, which are critical for LLMs to call functions properly.

In this example notebook, we create a simple Python function that executes arbitrary code and use it as a LangChain tool:

CREATE FUNCTION main.tools.python_exec (
  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()
$$

It runs in a secure and isolated environment within a Databricks SQL warehouse.

%pip install --upgrade --quiet databricks-sdk langchain-community langgraph mlflow

Note: you may need to restart the kernel to use updated packages.

from langchain_community.chat_models.databricks import ChatDatabricks

llm = ChatDatabricks(endpoint="databricks-meta-llama-3-70b-instruct")

API Reference:ChatDatabricks

from langchain_community.tools.databricks import UCFunctionToolkit
from databricks.sdk import WorkspaceClient

tools = (
    UCFunctionToolkit(
        # You can find the SQL warehouse ID in its UI after creation.
        warehouse_id="xxxx123456789"
    )
    .include(
        # Include functions as tools using their qualified names.
        # You can use "{catalog_name}.{schema_name}.*" to get all functions in a schema.
        "main.tools.python_exec",
    )
    .get_tools()
)

API Reference:UCFunctionToolkit

LangGraph agent example

from langgraph.prebuilt import create_react_agent
agent = create_react_agent(llm, tools, state_modifier="You are a helpful assistant. Make sure to use tool for information.")
agent.invoke({"messages": [{"role": "user", "content": "36939 * 8922.4"}]})

{'messages': [HumanMessage(content='36939 * 8922.4', additional_kwargs={}, response_metadata={}, id='c189e3c1-f054-4950-bc62-75a133324ff9'),
  AIMessage(content='', additional_kwargs={'tool_calls': [{'id': 'call_779479d0-60c5-42b5-af84-bb0901beac7c', 'type': 'function', 'function': {'name': 'main__tools__python_exec', 'arguments': '{"code": "print(36939 * 8922.4)"}'}}]}, response_metadata={'prompt_tokens': 771, 'completion_tokens': 29, 'total_tokens': 800}, id='run-3d0b3639-a5f8-49cf-bd5e-ea2c354f804f-0', tool_calls=[{'name': 'main__tools__python_exec', 'args': {'code': 'print(36939 * 8922.4)'}, 'id': 'call_779479d0-60c5-42b5-af84-bb0901beac7c', 'type': 'tool_call'}]),
  ToolMessage(content='{"format": "SCALAR", "value": "329584533.59999996\\n", "truncated": false}', name='main__tools__python_exec', id='d6150904-3425-49a4-baac-05d2d5b4772b', tool_call_id='call_779479d0-60c5-42b5-af84-bb0901beac7c'),
  AIMessage(content='The result of the multiplication is:\n\n329584533.59999996', additional_kwargs={}, response_metadata={'prompt_tokens': 846, 'completion_tokens': 22, 'total_tokens': 868}, id='run-2e11918b-b624-484a-b4f9-fa952d903ca0-0')]}

LangChain agent example

from langchain.agents import AgentExecutor, create_tool_calling_agent
from langchain_core.prompts import ChatPromptTemplate

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)

API Reference:AgentExecutor | create_tool_calling_agent | ChatPromptTemplate

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

> Entering new AgentExecutor chain...

Invoking: `main__tools__python_exec` with `{'code': 'print(36939 * 8922.4)'}`


{"format": "SCALAR", "value": "329584533.59999996\n", "truncated": false}The result of the multiplication is:

329584533.59999996

> Finished chain.

{'input': '36939 * 8922.4',
 'output': 'The result of the multiplication is:\n\n329584533.59999996'}

Related

• Tool conceptual guide
• Tool how-to guides