使用 Databricks SQL 代理进行查询

使用 Python SDK 的 Databricks

Avi Steinberg

Senior Software Engineer

Databricks Unity Catalog

在 Databricks 中存储数据的主要方式
包含架构，架构可含多张表
可从多种来源在架构中建表
samples 目录包含多个开源架构

Databricks Unity Catalog

¹ https://docs.databricks.com/aws/en/catalogs/

LangChain

构建与部署 LLM 应用的库
可创建使用 SQL 查询数据的 AI 代理

Python LangChain 依赖项

安装 Python 依赖：

pip install --upgrade databricks-langchain langchain-community langchain 
                      databricks-sql-connector databricks-sqlalchemy

导入库：

from langchain_community.agent_toolkits import create_sql_agent
from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langchain_community.utilities import SQLDatabase
from databricks_langchain import ChatDatabricks

Databricks SQL 仓库 ID

Databricks SQL 仓库连接详情

将 LangChain 认证到 Databricks 工作区

导出环境变量：

os.environ["DATABRICKS_TOKEN"] = "<Your-Access-Token>"
os.environ["DATABRICKS_HOST"] = "<your-workspace-id>.cloud.databricks.com"

Databricks SQL 代理

向 AI 代理就 Databricks 目录的架构数据提问
AI 代理可回答需要 SQL 的数据问题
指定 verbose=True 可显示代理的推理过程

LangChain 查询输出

使用 Databricks SQL 代理查询目录

db = SQLDatabase.from_databricks(
    catalog="samples", 
    schema="nyctaxi",
    warehouse_id=warehouse_id)
llm = ChatDatabricks(
    endpoint="databricks-meta-llama-3-3-70b-instruct",
    temperature=0.1,
    max_tokens=100)
toolkit = SQLDatabaseToolkit(db=db, llm=llm)
agent = create_sql_agent(llm=llm, toolkit=toolkit, verbose=True)
# Query the Databricks SQL Agent
result = agent.run("What's the time and distance of the longest trip?")
print(result)

¹ https://docs.databricks.com/aws/en/large-language-models/langchain#databricks-sql-agent

从 Databricks 目录创建 LangChain SQL 数据库

db = SQLDatabase.from_databricks(
    catalog="samples", 
    schema="nyctaxi",
    warehouse_id="<your-warehouse-id>"
)

¹ api.python.langchain.com/en/latest/utilities/langchain_community.utilities.sql_database.SQLDatabase.html

使用基础模型创建 LangChain LLM

llm = ChatDatabricks(

    endpoint="databricks-meta-llama-3-3-70b-instruct",

    temperature=0.1,
    max_tokens=100,
)

temperature：0 到 1 的浮点数，用于控制回答的随机性
max_tokens：限定回答的最大 token 数

¹ https://docs.databricks.com/aws/en/large-language-models/langchain

创建 LangChain SQL 数据库工具包

toolkit = SQLDatabaseToolkit(db=db, llm=llm)

¹ https://docs.databricks.com/aws/en/large-language-models/langchain

创建 Databricks SQL 代理

agent = create_sql_agent(llm=llm, toolkit=toolkit, verbose=True)

result = agent.run("What's the time and distance of the longest trip?")
display(result)

平均行程约 15 分钟。

¹ https://docs.databricks.com/aws/en/large-language-models/langchain

让我们查询 Databricks 数据！

使用 Python SDK 的 Databricks