1. DatabricksでAIエージェントを作ってみたい
「Agent Bricks」とやらが発表された！…けど
絶賛Databricks勉強中のせっとです。
実際に触ったり情報集めたりとしているのですが、先週のサミットで以下の発表がありました。
私もまだ理解が追い付いていない箇所がたくさんありますが、簡単に言うと
「エージェントを簡単に作成できるプラットフォーム"Agent Brick"の提供を始めたよ！」
とのことです。
このAgent Bricksを使うメリットとして、
- 開発コスト削減
- デプロイまでのスピード感と運用効率
- アクセス制御や監査が容易
などが大きく主張されています。
他にもサミットの動画や記事をいくつか拝見しましたが、
モデル精度やスコアにこだわっているというよりは、
とにかく 「ビジネス現場での利便性」 にこだわっている印象を受けました。
今回このAgent Bricksを使用してみたかったのですが、
まだFree Editionだとツールバーに表示がありませんでした。
Beta版だからでしょうか？正確な情報は見つけ切れませんでした。
左下のところ。本当はこんな感じで表示があるらしい。
ない。
「Mosaic AI エージェント」ならできそう
ただ、Databricksには他にもエージェントを構築できるフレームワークがあるようで、
それがこの 「Mosaic AIエージェントフレームワーク」
昨年にリリースされているのですかね。勉強不足でした。ざっくり一言でいうと
「組織専用のChatGPTみたいなものを、安全に・正確に動かすためのフレームワーク」
のようです。
海外が中心ですが、導入事例もいくつかありました。
セキュリティの不安を解消してくれるのは嬉しいですね。
今回はこのMosaic AIエージェントフレームワークを使用して、プチエージェントを構築してみたいと思います。
本記事でできるようになること
結論、こんな感じ。
どちらもdatabricks上のplaygroundで動かしているものですが、
このエージェントをAzureやAWSで簡単にデプロイできる。
ただの計算ではなくPythonを実行しています。
こちらは普通の自然言語処理
「でもこの会話なら普通のチャットボットと変わらんやん」と言われそうですが、実際そうでして…。
今回は下記クイックスタートガイドを参考に実行したのですが、このチュートリアルだけですと、
MLflowとChatAgentを最小構成で使った「チャットボット」 に近い状態でした。
そのため「AIエージェント感（＝判断・外部情報・動作の連鎖）」はほぼありません。
ただ「AIエージェント感」は薄くても、DatabricksにおけるAIエージェントの実運用に必要なポイントはいくつかあるのかなと思いました。
一旦は上記ガイドを実践しながら調べながらでいろいろ書いてみます。
2. ざっくり全体構成
まず全体構成をざっくり見るとこんな感じです。
[1] LLM選定 → [2] エージェント定義 → [3] モデル登録 → [4] 本番デプロイ
(使えるモデル確認) (ChatAgent実装) (MLflow + Unity Catalog) (REST API化)
改めて特徴を挙げるとこんな感じかと思います。
- エージェントの定義→登録→運用を すべてDatabricksで完結できる
- MLflowとUnity Catalogによる 再現性とガバナンスの担保
- エンドポイント経由で社内エージェントを ノーコードで 呼び出せる
3. 実装ステップ
では実際にここからはガイドを実行してみたいと思います。
(1) LLMエンドポイントの選択
最初に、Databricks上に登録されたLLMの中から今使えるLLMを自動的に選ぶ処理を行います。
なんだか不思議な処理な気がしたのですが、どうやらDatabricksでは、
- LLMは「モデル」ではなく「エンドポイント名」で指定し、
- 候補リストを順に試して、最初に成功したものを採用する
という仕組みがよく使われるようです。
「本番環境ではモデルを入れ替えたい」や「実験的に複数LLMを比較したい」ということが容易にできるのですね。
（折りたたみ）LLMエンドポイントの選択
# 最初に使えるLLMエンドポイントを自動で選ぶための準備
LLM_ENDPOINT_NAME = None
from databricks.sdk import WorkspaceClient
# 指定したエンドポイントが使えるかどうかをチェックする関数
def is_endpoint_available(endpoint_name):
try:
# DatabricksのOpenAI互換クライアントを取得
client = WorkspaceClient().serving_endpoints.get_open_ai_client()
# テスト用のメッセージを送って、応答が返ってくるか確認
client.chat.completions.create(
model=endpoint_name,
messages=[{"role": "user", "content": "AIとは何ですか？"}]
)
return True # 応答があればOK
except Exception:
return False # エラーが出たら使えない
# Databricksクライアントを初期化
client = WorkspaceClient()
# 候補のエンドポイントを順番にチェック
for candidate_endpoint_name in [
"databricks-claude-3-7-sonnet",
"databricks-meta-llama-3-3-70b-instruct"
]:
if is_endpoint_available(candidate_endpoint_name):
# 最初に使えたエンドポイントを選ぶ
LLM_ENDPOINT_NAME = candidate_endpoint_name
# どれも使えなかったらエラーを出す
assert LLM_ENDPOINT_NAME is not None, "LLM_ENDPOINT_NAMEを指定してください"
(2) エージェントの定義
次は 「Databricks上で動作する生成AIエージェントのロジック」 を構築します。
これがAIエージェントの中枢（思考＋実行）に当たる部分で、
今回の実装の中でも特に重要なところです。
Mosaic AI エージェントフレームワークで提供されている抽象クラスChatAgentを拡張して、
- ユーザーとの対話を受け取り、
- LLMを使って応答を生成し、
- 応答を構造化された形式で返す
という処理を可能にしています。
（折りたたみ）エージェントの定義
import uuid
import mlflow
from typing import Any, Optional
from mlflow.pyfunc import ChatAgent # MLflowのチャットエージェントのベースクラス
from mlflow.types.agent import ChatAgentMessage, ChatAgentResponse, ChatContext
mlflow.openai.autolog() # LLMの呼び出しやツール実行をMLflowに自動ログ
# ChatAgentを継承して、独自のエージェントを定義
class QuickstartAgent(ChatAgent):
def predict(
self,
messages: list[ChatAgentMessage], # ユーザーとの会話履歴
context: Optional[ChatContext] = None, # 会話の文脈（使わない場合もある）
custom_inputs: Optional[dict[str, Any]] = None, # カスタム入力（オプション）
) -> ChatAgentResponse:
# 1. 最後のユーザーメッセージ（プロンプト）を取り出す
prompt = messages[-1].content
# 2. run_agent関数を呼び出して、LLMの応答を取得
raw_msgs = run_agent(prompt)
# 3. 応答メッセージをChatAgentMessage形式に変換して返す
out = []
for m in raw_msgs:
out.append(ChatAgentMessage(id=uuid.uuid4().hex, **m))
return ChatAgentResponse(messages=out)
ここでは上記のような処理しかできていませんが、
predictメソッドの中に「単にLLMを呼び出す」だけでなく、
外部ツールの実行、状況に応じた処理分岐、他システムとの連携といった処理を組み込む ことで、
いわゆる AIエージェントらしい “ふるまい” を実現できます。
class QuickstartAgent(ChatAgent):
def predict(...):
...
たとえば、ユーザーの質問に応じて自動的に社内のデータベースを検索したり、Pythonで計算やグラフを生成して返したりするような処理は、すべて predict() の中で制御できます。
こうした応用は一旦今回はスキップして次回チャレンジしたいと思います。
(3) Unity CatalogにMLflowモデルとして保存
次は、Databricks上で構築したAIエージェントを MLflowモデルとして登録・保存する処理 です。
単にPythonファイルを保存するだけでなく、
使用リソース（LLMやPython関数）ごとエージェントをパッケージ化して、組織的に管理・再利用可能にする
というのがポイントです。
（折りたたみ）Unity CatalogにMLflowモデルとして保存
#必要なライブラリをインポート
import mlflow
from mlflow.models.resources import DatabricksFunction, DatabricksServingEndpoint
from pkg_resources import get_distribution
from quickstart_agent import LLM_ENDPOINT_NAME # 使うLLMのエンドポイント名を取得
# モデルを登録するカタログとスキーマを指定（デフォルトは current_catalog() と "default"）
catalog_name = spark.sql("SELECT current_catalog()").collect()[0][0]
schema_name = "default"
registered_model_name = f"{catalog_name}.{schema_name}.quickstart_agent"
# エージェントが使うリソース（LLMエンドポイントとPython実行ツール）を指定
resources = [
DatabricksServingEndpoint(endpoint_name=LLM_ENDPOINT_NAME),
DatabricksFunction(function_name="system.ai.python_exec"),
]
# モデルの保存先をUnity Catalogに設定
mlflow.set_registry_uri("databricks-uc")
# MLflowでモデルを保存・登録
with mlflow.start_run():
logged_agent_info = mlflow.pyfunc.log_model(
artifact_path="agent", # モデルの保存先（相対パス）
python_model="quickstart_agent.py", # エージェントのPythonファイル
extra_pip_requirements=[
f"databricks-connect=={get_distribution('databricks-connect').version}"
], # 必要なライブラリを指定
resources=resources, # エージェントが使うリソースを指定
registered_model_name=registered_model_name, # Unity Catalogに登録するモデル名
)
この登録処理の後は、MLflow UIやAPIからこのエージェントを呼び出して使うことができます。
また、ここではリソースとして、以下2つを指定しています。
-
DatabricksServingEndpoint：LLMを呼び出すためのエンドポイント -
DatabricksFunction： 任意のPythonコードを実行できる関数（今回はsystem.ai.python_exec）
resources = [
DatabricksServingEndpoint(endpoint_name=LLM_ENDPOINT_NAME),
DatabricksFunction(function_name="system.ai.python_exec"),
]
このようにLLMと関数を組み合わせることで、
エージェントに「外部と連携する力」を持たせる ことができます。
さらにここに検索用の関数や、社内APIとの連携機能などを追加すれば、
より高度で自律的なAIエージェントを構築することも可能です。
ちなみにCatalogの画面はこんな感じ。
一旦何もないですが、Unity Catalogを使うことで、
モデル単位で「部門別」「用途別」に権限管理 が可能となり、
「A部門のエージェントはA部門しか呼び出せない」といった制御もできるようになります。
(4) デプロイ
最後に、Databricks上に構築・登録されたAIエージェントを「REST API」としてデプロイします。
従来の「MLモデルのサービング」と違うのは、
チャット形式で応答できるAIエージェントが即座にAPIとして使える状態になる という点かと思います。
from databricks import agents
deployment_info = agents.deploy(
model_name=registered_model_name, # Unity Catalogに登録したモデル名
model_version=logged_agent_info.registered_model_version, # 登録されたモデルのバージョン
)
上記実行後、数分待って「サービング」を見てみるとこんな感じ。
ここで右上の「使用」⇒「Playgroundで試す」で冒頭のチャット画面に飛ぶことができます。
4. まとめ：DatabricksのAIエージェントでできそうなこと
今回は Databricks上で動くAIエージェントのプチ構築を体験してみました。
ここまで紹介してきたように、Mosaic AI エージェントフレームワーク に含まれる
ChatAgent クラスをベースに、LLMを組み込んだAIエージェントを比較的簡単に構築できます。
ただし、本質的な価値は今回のようなただのチャット応答ではなく、
その先にある 「実業務への応用」 にあります。
今回のコードを出発点に、実務で役立つエージェントに発展させる具体的な方向性はいくつも存在していると思っています。（実装難易度はさておき…）
-
議事録からのToDo抽出＋タスク登録エージェント
→predict()に要約・ToDo抽出処理を組み込み、DatabricksFunctionでAsanaやJira APIを呼び出してチケットを作成 -
CSVアップロードで即集計・レポート生成エージェント
→custom_inputs経由でCSVを受け取り、DatabricksでPython実行 → グラフ付きレポートを出力・保存 -
Slack経由で回答する社内FAQボット
→predict()でRAG（検索拡張生成）を組み込み、Slack連携で「チャットで聞けば答えてくれる」社内ナレッジBotを実装
このように、DatabricksのAIエージェントは単なるチャット応答に留まらず、
業務フローに入り込むアシスタントとして活用できます。
コードの再利用性やMLflowやUnity Catalogによる一元管理も、企業導入に適した強みです。
PoCから始めて段階的に育てていく開発スタイルとも相性が良いと感じました。
Discussion