Databricksにおけるトレース機能搭載エージェントのデプロイ

Posted at 2025-10-10

引き続きMLflow 3の勉強中。こちらのサンプルノートブックを実行していきます。

Mosaic AI Agent Framework: シンプルなエージェントの作成、デプロイ、トレース

このノートブックでは、シンプルな生成AIエージェントの構築と管理方法を説明します：

MLflow 3の ResponsesAgent API を使ってシンプルな生成AIエージェントを作成
エージェントの手動テストとMLflowによるバッチ評価
Mosaic AI Agent Frameworkでエージェントを記録・デプロイ
エージェントのリアルタイムトレースと監視

このパターンは、どのAgent Frameworkエージェント（AWS | Azure | GCP）でも利用できます。

MLflow 3（AWS | Azure | GCP | OSS）は以下の観測機能を提供します：

品質や運用パフォーマンス（レイテンシ、リクエスト数、エラーなど）の追跡
Agent EvaluationのLLMジャッジを使った本番トラフィックの評価・ドリフト検知
個々のリクエストの詳細分析による応答改善
実際のログを評価セットに変換し、継続的な改善に活用

前提条件

このノートブックの TODO をすべて対応してから「すべて実行」をクリックしてください。

%pip install -U -qqqq backoff databricks-openai uv databricks-agents mlflow-skinny[databricks]
dbutils.library.restartPython()

エージェントのコード定義

下のセルでエージェントコードを定義します。これにより、%%writefile マジックコマンドを使ってローカルのPythonファイル agent.py に簡単に書き出し、後で記録やデプロイに利用できます。

%%writefile agent.py
import json
import warnings
from typing import Any, Generator

import backoff
import mlflow
import openai
from databricks.sdk import WorkspaceClient
from mlflow.entities import SpanType
from mlflow.pyfunc import ResponsesAgent
from mlflow.types.responses import (
    ResponsesAgentRequest,
    ResponsesAgentResponse,
    ResponsesAgentStreamEvent,
)
from openai import OpenAI

# TODO: あなたのモデルサービングエンドポイント名に置き換えてください
LLM_ENDPOINT_NAME = "databricks-claude-sonnet-4-5"

# TODO: システムプロンプトを日本語で更新してください
SYSTEM_PROMPT = """
あなたは親切なアシスタントです。簡潔で分かりやすい回答を提供してください。
"""


class SimpleChatAgent(ResponsesAgent):
    """
    Databricks OpenAIクライアントAPIを使ってLLMを呼び出すシンプルなチャットエージェントです。

    必要に応じて独自のエージェントに置き換えてください。
    @mlflow.trace デコレータはMLflowトレーシングでエージェントの呼び出しを追跡します。
    """

    def __init__(self):
        self.workspace_client = WorkspaceClient()
        self.client: OpenAI = self.workspace_client.serving_endpoints.get_open_ai_client()
        self.llm_endpoint = LLM_ENDPOINT_NAME
        self.SYSTEM_PROMPT = SYSTEM_PROMPT

    @backoff.on_exception(backoff.expo, openai.RateLimitError)
    @mlflow.trace(span_type=SpanType.LLM)
    def call_llm(self, messages: list[dict[str, Any]]) -> Generator[dict[str, Any], None, None]:
        with warnings.catch_warnings():
            warnings.filterwarnings("ignore", message="PydanticSerializationUnexpectedValue")
            for chunk in self.client.chat.completions.create(
                model=self.llm_endpoint,
                messages=self.prep_msgs_for_cc_llm(messages),
                stream=True,
            ):
                yield chunk.to_dict()

    # オートロギングを使う場合、ここに@mlflow.traceは不要ですが、span typeを上書きしたい場合は追加できます。
    def predict(self, request: ResponsesAgentRequest) -> ResponsesAgentResponse:
        outputs = [
            event.item
            for event in self.predict_stream(request)
            if event.type == "response.output_item.done"
        ]
        return ResponsesAgentResponse(output=outputs, custom_outputs=request.custom_inputs)

    # オートロギングを使う場合、ここに@mlflow.traceは不要ですが、span typeを上書きしたい場合は追加できます。
    def predict_stream(
        self, request: ResponsesAgentRequest
    ) -> Generator[ResponsesAgentStreamEvent, None, None]:
        messages = [{"role": "system", "content": SYSTEM_PROMPT}] + [
            i.model_dump() for i in request.input
        ]
        yield from self.output_to_responses_items_stream(chunks=self.call_llm(messages))


mlflow.openai.autolog()
AGENT = SimpleChatAgent()
mlflow.models.set_model(AGENT)

エージェントのテスト

エージェントと対話して出力をテストします。

ResponsesAgent 内のメソッドを手動でトレースしているため、各ステップのトレースを確認でき、OpenAI SDK経由のLLM呼び出しも自動でトレースされます。

このプレースホルダー入力を、あなたのエージェントに適したドメイン固有の例に置き換えてください。

dbutils.library.restartPython()

from agent import AGENT

AGENT.predict({"input": [{"role": "user", "content": "5+5は何ですか？"}]})

上のmlflow.openai.autolog()でオートロギングが有効化されているのでトレースが生成されます。

ノートブックエクスペリメントのトレースでも確認できます。

これはストリーミングの例です。

for event in AGENT.predict_stream(
    {"input": [{"role": "user", "content": "5+5は何ですか？"}]}
):
    print(event, "-----------\n")

type='response.output_text.delta' custom_outputs=None item_id='msg_bdrk_0172sRwMSMhLxJR4KoYB5iLQ' delta='5+5は**' -----------

type='response.output_text.delta' custom_outputs=None item_id='msg_bdrk_0172sRwMSMhLxJR4KoYB5iLQ' delta='10**です。' -----------

type='response.output_text.delta' custom_outputs=None item_id='msg_bdrk_0172sRwMSMhLxJR4KoYB5iLQ' delta='' -----------

type='response.output_item.done' custom_outputs=None item={'id': 'msg_bdrk_0172sRwMSMhLxJR4KoYB5iLQ', 'content': [{'text': '5+5は**10**です。', 'type': 'output_text'}], 'role': 'assistant', 'type': 'message'} -----------

ノートブックエクスペリメントのトレースにも追加されます。

エージェントをMLflowモデルとして記録し、Unity Catalogに登録

agent.py ファイルのコードをMLflowモデルとして記録します。詳細は MLflow - Models from Code を参照してください。

同じ記録処理で、モデルをUnity Catalogに登録できます。次のステップでエージェントをデプロイする際に必要です。Databricksのモデル管理については、AWS | Azure | GCP を参照してください。

import mlflow
from agent import LLM_ENDPOINT_NAME
from mlflow.models.resources import DatabricksServingEndpoint
from pkg_resources import get_distribution

# モデルレジストリはデフォルトでDatabricks Unity Catalogに設定されています。
# 必要に応じて下記でレジストリURIを変更できます。
mlflow.set_registry_uri("databricks-uc")

# TODO: UCに登録するモデルのカタログ名、スキーマ名、モデル名を指定してください。
catalog = "takaakiyayoi_catalog"
schema = "agents"
model_name = "simple_agent"
UC_MODEL_NAME = f"{catalog}.{schema}.{model_name}"

with mlflow.start_run():
    logged_agent_info = mlflow.pyfunc.log_model(
        # モデル/エージェント名は用途に合わせて分かりやすい名前に変更してください。
        name="agent",
        # 上で作成したPythonファイルを指定します。
        python_model="agent.py",
        # pip_requirementsではなくextra_pip_requirementsを指定する場合、
        # mlflowのバージョンがノートブック環境と一致していることを確認してください。
        extra_pip_requirements=[
            f"databricks-connect=={get_distribution('databricks-connect').version}",
        ],
        resources=[DatabricksServingEndpoint(endpoint_name=LLM_ENDPOINT_NAME)],
        # このオプションで記録と同時にモデルを登録できます。
        registered_model_name=UC_MODEL_NAME,
    )

🔗 View Logged Model at: https:/xxxx.cloud.databricks.com/ml/experiments/31285241155985/models/m-b3ebad8880d447a8a111af5c09159199?o=5099015744649857
2025/10/10 07:19:34 INFO mlflow.pyfunc: Predicting on input example to validate output
Successfully registered model 'takaakiyayoi_catalog.agents.simple_agent'.

🔗 Created version '1' of model 'takaakiyayoi_catalog.agents.simple_agent': https://xxxx.cloud.databricks.com/explore/data/models/takaakiyayoi_catalog/agents/simple_agent/version/1?o=5099015744649857

上でmlflow.start_runを呼び出しているので、MLflowランが作成されます。

ちなみに、MLflowランがエクスペリメントに登録された際にはジョブの実行タブが表示されるようになるんですね。

また、registered_model_nameを指定しているのでUnity Catalogにも登録されます。

デプロイ前のエージェント検証

デプロイ前に以下のチェックを行ってください。

手動チェック：mlflow.models.predict() API を利用。Databricksのドキュメント（AWS | Azure | GCP）も参照。
データセット評価：mlflow.genai.evaluate() API を利用。Databricksのドキュメント（AWS | Azure | GCP）も参照。

print(logged_agent_info.registered_model_version)

mlflow.models.predict(
    model_uri=f"runs:/{logged_agent_info.run_id}/agent",
    input_data={"input": [{"role": "user", "content": "こんにちは！"}]},
    env_manager="uv",
)

これもトレースとして記録されます。

バッチ評価

MLflowを使ってエージェントをバッチトレースで評価する方法を紹介します。詳細はDatabricksのドキュメント（AWS | Azure | GCP）を参照してください。

評価対象のトレースを収集
実行するスコアラー（LLMジャッジ）を選択
評価を実行し、メトリクスを算出

traces = mlflow.search_traces(max_results=10)
traces

これまで記録されたトレースにアクセスできます。

ここでのスコアラーは、こちらの記事で書いた一般的な評価概念としてのスコアラーです。

from mlflow.genai.scorers import (
    RelevanceToQuery,  # ユーザーのリクエストに対する関連性を評価
    Safety,            # 有害または不適切な内容が含まれていないか評価
    Guidelines,        # カスタムガイドライン評価
)

scorers = [
  RelevanceToQuery(),  # メールがユーザーの要望に応えているかをチェック
  Safety(),  # 有害または不適切な内容が含まれていないかをチェック
  # 以下はカスタムガイドライン:
  Guidelines(
      name="簡潔なコミュニケーション",
      guidelines="回答は必ず簡潔かつ要点を押さえてください。",
  ),
]

# 上記で選択したスコアラーを使って評価を実行
eval_results = mlflow.genai.evaluate(
    data=traces,
    model_id=logged_agent_info.model_id,
    scorers=scorers,
)

表示されるView evaluation results. をクリックして、3つのトレースの評価結果を確認します。

バージョンタブのLoggedModelにも評価結果が紐づいています。

エージェントのデプロイ

Agent Framework（AWS | Azure | GCP）を使ってエージェントをデプロイします。デフォルトでは、デプロイしたエージェントのトレースが現在の実験にリアルタイムで記録され、推論テーブル（有効化時）にも記録されます。

from databricks import agents

agents.deploy(UC_MODEL_NAME, model_version=logged_agent_info.registered_model_version, tags={"endpointSource": "docs"})


    Deployment of takaakiyayoi_catalog.agents.simple_agent version 1 initiated.  This can take up to 15 minutes and the Review App & Query Endpoint will not work until this deployment finishes.

    View status: https://xxxx.cloud.databricks.com/ml/endpoints/agents_takaakiyayoi_catalog-agents-simple_agent/?o=5099015744649857
    Review App: https://xxxx.cloud.databricks.com/ml/review-v2/5c793570088c4d4b8d0dcec747607964/chat?o=5099015744649857
    Monitor: https://xxxx.cloud.databricks.com/ml/experiments/31285241155985?compareRunsMode=TRACES&o=5099015744649857

You can refer back to the links above from the endpoint detail page at https://xxxx.cloud.databricks.com/ml/endpoints/agents_takaakiyayoi_catalog-agents-simple_agent/?o=5099015744649857.
Deployment(model_name='takaakiyayoi_catalog.agents.simple_agent', model_version='1', endpoint_name='agents_takaakiyayoi_catalog-agents-simple_agent', served_entity_name='takaakiyayoi_catalog-agents-simple_agent_1', query_endpoint='https://xxxx.cloud.databricks.com/serving-endpoints/agents_takaakiyayoi_catalog-agents-simple_agent/served-models/takaakiyayoi_catalog-agents-simple_agent_1/invocations?o=5099015744649857', endpoint_url='https://xxxx.cloud.databricks.com/ml/endpoints/agents_takaakiyayoi_catalog-agents-simple_agent/?o=5099015744649857', review_app_url='https://xxxx.cloud.databricks.com/ml/review-v2/5c793570088c4d4b8d0dcec747607964/chat?o=5099015744649857')

数分待つとモデルサービングエンドポイントの準備ができます。

今日気づいたんですけど、この時点でトレースタブがモデルサービングエンドポイントにも追加されているんですよね。

エンドポイントからのリアルタイムトレースの表示

デプロイしたエージェントは、デフォルトでMLflowエクスペリメントにリアルタイムでトレースを記録します。トレースを記録するMLflowエクスペリメントを変更したい場合は、agents.deploy(...) の前に mlflow.set_experiment(...) を呼び出してください。

本番監視を有効化すると、MLflowエクスペリメントからDeltaテーブルにトレースをコピーできます（AWS | GCP）。監視を有効化した場合、MLflowエクスペリメントの Monitoring タブで本番トレースに対して品質スコアラーを更新できます。

print(f"\nリアルタイムでエンドポイントのトレースをMLflowエクスペリメントで確認できます: https://{mlflow.utils.databricks_utils.get_browser_hostname()}/ml/experiments/{mlflow.get_experiment_by_name(mlflow.utils.databricks_utils.get_notebook_path()).experiment_id}/traces")