使用 Databricks SQL 代理查詢

使用 Python SDK 的 Databricks

Avi Steinberg

Senior Software Engineer

Databricks Unity Catalog

在 Databricks 儲存資料的主要方式
內含 schema，可包含多個資料表
從多種來源在 schema 中建立資料表
samples catalog 中提供多個開源 schema

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 Warehouse ID

Databricks SQL Warehouse 連線詳細資料

將 LangChain 驗證到 Databricks Workspace

匯出環境變數：

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

Databricks SQL 代理

向 AI 代理詢問 Databricks catalog 中某個 schema 的資料
AI 代理能回答需要 SQL 的資料問題
指定 verbose=True 可顯示 AI 代理的推理過程

LangChain Query output

使用 Databricks SQL 代理查詢 Catalog

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 Catalog 建立 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)

The average trip takes approximately 15 minutes.

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

來查詢 Databricks 資料吧！

使用 Python SDK 的 Databricks