こちらの続きです。
前回はツールを構築してAI Playgroundで動かすところまでをカバーしました。やり取りを通じてエージェントの目鼻顔立ちが見えてきたら、エージェントシステムとしてデプロイしたくなることでしょう。
その際に活用するのがエクスポートボタンです。
こちらにあるように、前のステップで定義したツールが組み込まれたエージェントシステムの実装を含むノートブックが生成されます。必要に応じてカスタマイズして本格運用に進むことができます。これも、Databricks AutoMLと同様のガラスボックスアプローチとなっています。
-
agentノートブック: を使用してエージェントを定義するPython LangChainコードが含まれています。 -
driverノートブック: エージェントフレームワークを使用してPython エージェントをログに記録、トレース、登録する デプロイするためのAI Mosaic AIコードが含まれています。 -
config.yml: ツール定義を含む、エージェントに関する構成情報が含まれます。
Databricksにおけるエージェント開発のアプローチ
ウォークスルーの前に、Databricksにおけるエージェント開発のアプローチに触れさせてください。RAGを構築された方はすでに体験されていると思いますが、このようなシステムにおいて本格運用に十分な品質を確保することは相当に難しいこととなっています。理由としては、以下のようなものが挙げられます。
- 評価手法自体が確立されているフェーズ
- 調整すべきコンポーネントが多岐にわたる
このような状況で闇雲にPOCを繰り返しても、望む結果を得ることができません。このため、Databricksでは開発時点から評価の仕組みを組み込む評価ドリブン開発ワークフローを採用しています。
このワークフローを実現しているのは、以下のDatabricksのコンポーネントです。
-
Mosaic AI Agent Framework
- エージェントを作成、MLflowで記録
- 1行のコードでエージェントを
- モデルサービングエンドポイントにデプロイ
- ステークホルダーからフィードバックを収集するためのレビューアプリを自動で作成
-
Mosaic AI Agent Evaluation
- エージェントを評価するためにプロプライエタリなLLMジャッジを活用
- レーテンシーやトークンのコストのようなパフォーマンスメトリクスを収集
-
AI Playground
- 様々なLLMに対するプロンプトエンジニアリング
- 複数のモデルの挙動を比較
- ツールの組み込み
前振りが長くなりましたが、最初に触れたエクスポートボタンを押すことで、上のコンポーネントとのインテグレーションが含まれたノートブックが作成されるということです。
エージェントのデプロイメントのウォークスルー
エクスポートされたノートブックは、ホームフォルダ配下のAgent databricks-meta-llama-3-1-70b-instruct 0006-10-21 08:59:08というようなフォルダが作成され、そちらに格納されます。
config.yml
使用するLLM、ツールが記載されています。
llm_endpoint: "databricks-meta-llama-3-1-70b-instruct"
tools:
uc_functions:
- "takaakiyayoi_catalog.cookie.*"
agent
エージェントの実装です。このノートブックは直接実行しません。後述のdriverノートブックからアクセスします。ローカルでテストする際には、ノートブックを実行することが可能です。
%pip install -U -qqqq mlflow-skinny
dbutils.library.restartPython()
import mlflow
from mlflow.models import ModelConfig
config = ModelConfig(development_config="config.yml")
import mlflow
import mlflow.deployments
from mlflow.models.rag_signatures import (
ChatCompletionRequest,
ChatCompletionResponse,
Message,
)
from dataclasses import asdict, is_dataclass
class ManagedAgentAPI(mlflow.pyfunc.PythonModel):
def predict(
self, context, model_input: ChatCompletionRequest
) -> ChatCompletionResponse:
from mlflow.entities import Trace
# Send a request to Databricks Model Serving with tools execution
# Convert input to dict if necessary for model serving compatibility
if is_dataclass(model_input):
req = asdict(model_input).copy()
else:
req = model_input.copy()
# Prepend system message if it exists
try:
system_message = config.get("agent_prompt")
req["messages"].insert(0, asdict(Message(role="system", content=system_message)))
except KeyError:
pass
# Add UC functions if they exist
try:
uc_functions = config.get("tools").get("uc_functions")
if uc_functions:
req["tools"] = [
{"type": "uc_function", "uc_function": {"name": name}}
for name in uc_functions
]
req["databricks_options"] = {"return_trace": True}
except KeyError:
req["tools"] = []
client = mlflow.deployments.get_deploy_client("databricks")
completion = client.predict(endpoint=config.get("llm_endpoint"), inputs=req)
# Add the trace from model serving API call to the active trace
if trace := completion.pop("databricks_output", {}).get("trace"):
trace = Trace.from_dict(trace)
mlflow.add_trace(trace)
return completion
agent = ManagedAgentAPI()
# TODO: replace this placeholder input example with an appropriate domain-specific example for your agent
agent.predict(
None,
ChatCompletionRequest(
messages=[Message(role="user", content="what is 23423 * 98023?")]
),
)
mlflow.models.set_model(agent)
driver
Mosaic AI Agent Frameworkを用いて、agentノートブックで実装されているエージェントの記録、評価、サービングエンドポイントへのデプロイなどをエンドツーエンドで実行します。
なお、AI Playgroundでやりとりした内容が評価データセットとして盛り込まれるようになっています。
Unity Catalogに登録するカタログ、スキーマ、モデル名を指定します。
mlflow.set_registry_uri("databricks-uc")
# TODO: define the catalog, schema, and model name for your UC model
catalog = "takaakiyayoi_catalog"
schema = "cookie"
model_name = "cookie_bot"
UC_MODEL_NAME = f"{catalog}.{schema}.{model_name}"
# register the model to UC
uc_registered_model_info = mlflow.register_model(model_uri=logged_agent_info.model_uri, name=UC_MODEL_NAME)
最後のコマンドを実行すると、モデルサービングエンドポイントにエージェントのモデルとレビューアプリがデプロイされます。
from databricks import agents
# Deploy the model to the review app and a model serving endpoint
agents.deploy(UC_MODEL_NAME, uc_registered_model_info.version)
モデルサービングエンドポイントがReadyになったら、レビューアプリにアクセスして挙動を確認します。
エージェントにおけるリクエストとレスポンス、レスポンスに対するフィードバックも推論テーブルに記録されるので、それらを活用して精度を改善することができます。