利用Databricks Unity Catalog和LangChain实现Python代码执行工具

最新推荐文章于 2025-04-26 21:03:53 发布

fgayif

最新推荐文章于 2025-04-26 21:03:53 发布

阅读量255

点赞数 5

文章标签： unity langchain python

本文链接：https://blog.youkuaiyun.com/fgayif/article/details/146466547

版权

在现代数据分析和机器学习工作流中，灵活地集成各种工具和技术以简化复杂任务变得越来越重要。本文我们将展示如何利用Databricks Unity Catalog (UC) 创建一个Python函数，并将其作为LangChain工具进行使用。这个例子演示了如何在Databricks SQL仓库中安全地执行任意Python代码。

技术背景介绍

Databricks Unity Catalog提供了一种集中管理数据资产的方式，而LangChain是一种用于构建自然语言处理应用的工具链。结合这两者，我们可以在数据管理和自然语言处理项目中实现更强大的功能。

核心原理解析

我们的目标是创建一个Python函数并通过LangChain使用该函数。这个函数将允许我们在Databricks中执行任意Python代码，并返回执行结果。LangChain通过工具套件（Toolkits）使得这样的功能集成更加容易。

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

上面的SQL代码创建了一个能够执行Python代码的函数。该函数在Databricks SQL仓库中以安全、隔离的方式运行。

代码实现演示(重点)

下面我们将展示如何利用LangChain来调用上述函数。

# 安装必要的库
%pip install --upgrade --quiet databricks-sdk langchain-community mlflow

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

# 设置工具函数
tools = (
    UCFunctionToolkit(
        warehouse_id="xxxx123456789"  # 替换为实际的SQL仓库ID
    )
    .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)

# 调用工具并执行任务
agent_executor.invoke({"input": "36939 * 8922.4"})