Databricks SQLエージェントによるクエリ

Databricks と Python SDK

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エージェント

DatabricksカタログスキーマのデータについてAIエージェントに質問できる
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: 応答の最大トークン数を指定する

¹ 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データをクエリしてみましょう！

Databricks と Python SDK