DatabricksでClaude 3.7 sonnetを使うAIエージェントをMosaic AI Agent Frameworkで実装評価！

# 導入

最近Databricksに関する様々なニュースが発表されており非常に面白いです。

その中でも、Anthropic社との提携とすぐにDatabricksネイティブでClaude 3.7 Sonnetを利用可能になったのは驚きでした。

https://qiita.com/taka_yayoi/items/9ccf1949933936a28264

Claude 3.7 Sonnetは仕事でも利用していますが、ビジネス利用においては非常にバランスが良く優秀なモデルだと感じてます。個人的にも現時点で最高に使いやすいモデルだと感じています。
（とか言っていたらGemini 2.5が出て来たりとすぐに意見変わりそうですが。。。同僚曰く、こちらもかなりすごいようですね）


加えて、いつの間にか**AWS東京リージョンにMosaic AI Agent Framework & Evaluationが来ていました！**
公式Docに記載されているので間違いないはず。

https://docs.databricks.com/aws/ja/resources/feature-region-support#ai%E3%81%8A%E3%82%88%E3%81%B3%E6%A9%9F%E6%A2%B0%E5%AD%A6%E7%BF%92%E6%A9%9F%E8%83%BD%E3%81%AE%E5%8F%AF%E7%94%A8%E6%80%A7

ようやく来たよ、~~俺の司が！~~

というわけで、遅まきながらMosaic AI Agent FrameworkとClaude 3.7 Sonnet両方を使ってみたかったので、公式Docなどを見ながらウォークスルーする内容となっています。

基本的には以下のDocや神の記事を参考にして実施しています。
あまり目新しい内容はありませんが、Claude 3.7 Sonnetを使ったエージェント作成の参考になれば幸いです。

https://docs.databricks.com/aws/ja/generative-ai/agent-framework/author-agent#chatagent-%E4%BE%8B

https://qiita.com/taka_yayoi/items/323220b34d97bc341959

また、本記事には手前みそですが以下で検証した内容も多く含んでいます。
よければ合わせて参照ください。

https://qiita.com/isanakamishiro2/items/ec6d21823d6f72ee2197

実行検証はDatabricks on AWS上で行いました。
ノートブックのクラスタはサーバレスです。
Mosaic AI Model ServingでClaude 3.7 sonnetを利用できるように環境設定されていることが前提となります。

# Step1. AIエージェントを作る

最初からクライマックス。

今回はユーザの依頼に対して「pythonコードを実行して回答を作る」エージェントを作成します。
雑に言えば、Databricks UnityCatalog上に標準で提供されている```"system.ai.python_exec```関数をツールとして利用できるエージェントです。

まず、ノートブックを作成して必要なパッケージをインストール。```langgraph```などエージェントを構成するフレームワークに加えて、Mosaic AI Agent Frameworkを利用するための```databricks-agents```パッケージをインストールします。

```python
%pip install -U -qqqq databricks-langchain databricks-agents>=0.16.0 mlflow-skinny[databricks] unitycatalog-langchain[databricks] langgraph==0.3.21 uv loguru rich

%restart_python
```

次にコードでAIエージェントを実装します。
内容は[こちら](https://qiita.com/isanakamishiro2/items/ec6d21823d6f72ee2197)で解説したカスタムChatAgentとして実装しています。
上記記事との違いは設定するツールをUnityCatalog関数から指定していることと、LLMにClaude 3.7 Sonnetを利用していることです。

コードは```%%writefile```を使ってreact_agent.pyとして出力しています。

```python
%%writefile react_agent.py

from typing import Literal, Generator, List, Optional, Any, Dict, Mapping, Union
import uuid

import mlflow
from databricks_langchain import (
    ChatDatabricks,
    UCFunctionToolkit,
)
from langchain_core.tools import tool
from langgraph.graph.state import CompiledStateGraph
from langgraph.prebuilt import create_react_agent

from langchain_core.messages import BaseMessage, convert_to_openai_messages
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult
from databricks_langchain import ChatDatabricks
from databricks_langchain.uc_ai import (
    DatabricksFunctionClient,
    UCFunctionToolkit,
    set_uc_function_client,
)
from mlflow.pyfunc import ChatAgent
from mlflow.types.agent import (
    ChatAgentChunk,
    ChatAgentMessage,
    ChatAgentResponse,
    ChatContext,
)
from functools import reduce

mlflow.langchain.autolog()

# Databricks UnityCatalogから関数を取得するためのクライアントを設定
client = DatabricksFunctionClient()
set_uc_function_client(client)

class LangGraphChatAgent(ChatAgent):
    def __init__(self, agent: CompiledStateGraph):
        """ LangGraphのグラフを指定して初期化 """
        self.agent = agent

    def predict(
        self,
        messages: list[ChatAgentMessage],
        context: Optional[ChatContext] = None,
        custom_inputs: Optional[dict[str, Any]] = None,
    ) -> ChatAgentResponse:
        
        """
        指定されたチャットメッセージリストを使用して回答を生成する

        Args:
            messages (list[ChatAgentMessage]): チャットエージェントメッセージのリスト。
            context (Optional[ChatContext]): オプションのチャットコンテキスト。
            custom_inputs (Optional[dict[str, Any]]): カスタム入力のオプション辞書。

        Returns:
            ChatAgentResponse: 予測結果を含むChatAgentResponseオブジェクト。
        """

        request = {"messages": self._convert_messages_to_dict(messages)}

        messages = []
        usages = []
        for event in self.agent.stream(request, stream_mode="updates"):
            for node_data in event.values():
                messages.extend(
                    self._convert_lc_message_to_chat_message(msg)
                    for msg in node_data.get("messages", [])
                )
                usages.extend(
                    msg.response_metadata for msg in node_data.get("messages", [])
                )

        usage = self._sum_usages(usages)
        return ChatAgentResponse(messages=messages, usage=usage)

    def predict_stream(
        self,
        messages: list[ChatAgentMessage],
        context: Optional[ChatContext] = None,
        custom_inputs: Optional[dict[str, Any]] = None,
    ) -> Generator[ChatAgentChunk, None, None]:
        request = {"messages": self._convert_messages_to_dict(messages)}
        for event in self.agent.stream(request, stream_mode="updates"):
            for node_data in event.values():
                yield from (
                    ChatAgentChunk(
                        **{"delta": self._convert_lc_message_to_chat_message(msg)},
                        usage=msg.response_metadata
                    )
                    for msg in node_data["messages"]
                )

    def _convert_lc_message_to_chat_message(
        self, lc_message: BaseMessage
    ) -> ChatAgentMessage:
        """LangChainメッセージをChatAgentMessageに変換する。"""
        msg = convert_to_openai_messages(lc_message)
        if not "id" in msg:
            msg.update({"id": str(uuid.uuid4())})

        return ChatAgentMessage(**msg)

    def _sum_usages(self, usages: list[dict]) -> dict:
        """使用量のリストから使用量を合計する。"""

        def add_usages(a: dict, b: dict) -> dict:
            pt = "prompt_tokens"
            ct = "completion_tokens"
            tt = "total_tokens"
            return {
                pt: a.get(pt, 0) + b.get(pt, 0),
                ct: a.get(ct, 0) + b.get(ct, 0),
                tt: a.get(tt, 0) + b.get(tt, 0),
            }

        return reduce(add_usages, usages, {})


# DatabricksネイティブのClaude 3.7 SonnetをLLMとして利用
LLM_ENDPOINT_NAME = "databricks-claude-3-7-sonnet"
llm = ChatDatabricks(model=LLM_ENDPOINT_NAME)


# Databricks標準提供のツール system.ai.python_execをツールとして設定
tools = []
uc_tool_names = ["system.ai.python_exec"]
uc_toolkit = UCFunctionToolkit(function_names=uc_tool_names)
tools.extend(uc_toolkit.tools)

graph = create_react_agent(llm, tools=tools)
AGENT = LangGraphChatAgent(graph)
mlflow.models.set_model(AGENT)
```

試験実行してみましょう。

念のためPythonのセッションをリスタート。
```python
%restart_python
```

単純な計算を実行させます。

```python
from react_agent import AGENT
from rich import print

print(AGENT.predict({"messages": [{"role": "user", "content": "What is 5+5 in python"}]}))
```

```:出力結果
ChatAgentResponse(
    messages=[
        ChatAgentMessage(
            role='assistant',
            content='I can help you calculate 5+5 using Python. Let me execute this simple addition for you.',
            name=None,
            id='a2a39fd3-2d22-4076-b8df-340086bbdd56',
            tool_calls=[
                ToolCall(
                    id='toolu_bdrk_018VSV2nsZdNXw8TBiseetYe',
                    type='function',
                    function=Function(name='system__ai__python_exec', arguments='{"code": "print(5 + 5)"}')
                )
            ],
            tool_call_id=None,
            attachments=None
        ),
        ChatAgentMessage(
            role='tool',
            content='{"format": "SCALAR", "value": "10\\n"}',
            name='system__ai__python_exec',
            id='7b137fbf-990f-4ee6-8636-f83981cade9b',
            tool_calls=None,
            tool_call_id='toolu_bdrk_018VSV2nsZdNXw8TBiseetYe',
            attachments=None
        ),
        ChatAgentMessage(
            role='assistant',
            content='The result of 5+5 in Python is 10.',
            name=None,
            id='9d85fd3d-e47b-46cf-a40c-ac1e94722f3f',
            tool_calls=None,
            tool_call_id=None,
            attachments=None
        )
    ],
    finish_reason=None,
    custom_outputs=None,
    usage=ChatUsage(prompt_tokens=0, completion_tokens=0, total_tokens=0)
)
```

MLflow Tracingの画面。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/9ef2eac2-a2ae-4402-931e-42d87778e27a.png)

というわけで、きちんとツール(pythonコードの実行)を呼んで5+5が計算されました。

# Step2. エージェントのロギング

MLflowに実装したエージェントをロギングします。
ポイントは[こちら](https://qiita.com/isanakamishiro2/items/ec6d21823d6f72ee2197)で検証したように、Databricks上で利用するマネージドサービスをリソースとして指定すること。
（今回のケースで言えば、Claude 3.7 SonnetのエンドポイントとUnityCatalog関数をresoucesに含めています）

```python
import mlflow
from react_agent import tools, LLM_ENDPOINT_NAME
from mlflow.models.resources import DatabricksFunction, DatabricksServingEndpoint
from unitycatalog.ai.langchain.toolkit import UnityCatalogTool
from rich import print

resources = [DatabricksServingEndpoint(endpoint_name=LLM_ENDPOINT_NAME)]
for tool in tools:
    if isinstance(tool, UnityCatalogTool):
        resources.append(DatabricksFunction(function_name=tool.uc_function_name))
print(resources)

input_example = {
    "messages": [
        {
            "role": "user",
            "content": "15番目のフィボナッチ数を計算する"
        }
    ]
}
print(input_example)

with mlflow.start_run():
    logged_agent_info = mlflow.pyfunc.log_model(
        artifact_path="agent",
        python_model="react_agent.py",
        input_example=input_example,
        pip_requirements=[
            "mlflow",
            "langgraph==0.3.21",
            "databricks-langchain==0.4.1",
            "unitycatalog-langchain[databricks]",
        ],
        resources=resources,
    )
```

これでエージェントを再利用したりAgent Evaluation Frameworkから利用できるようになりました。

# Step3. Agent Evaluationを使ったエージェントの評価

以下のDocにあるように、DatabricksのMosaic AI Agent Frameworkはエージェント評価の機能を提供しています。

https://docs.databricks.com/aws/ja/generative-ai/agent-evaluation

Mosaic AI Agent Frameworkが東京リージョンで利用できるようになった結果、この機能も利用できるようになりました。さっそくこの機能を使った簡易評価を行ってみます。

まず、質問・指示と想定回答のペアをPandas DFとして作成。

```python
import pandas as pd

eval_examples = [
    {
        "request": {
            "messages": [
                {
                    "role": "user",
                    "content": "15番目のフィボナッチ数を計算する"
                }
            ]
        },
        "expected_response": "フィボナッチ数列の15番目の数は610です"
    },
    {
        "request": {
            "messages": [
                {
                    "role": "user",
                    "content": "日本の総理大臣をpython標準モジュールだけを使ってWebから入手して回答して"
                }
            ]
        },
        "expected_response": "日本の総理大臣は石破茂です"
    }

]

eval_dataset = pd.DataFrame(eval_examples)
display(eval_dataset)
```

以下のような評価セットのデータフレームが作成されます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/f310d323-5b8f-43a4-9529-070a3634b6ea.png)

では、この評価セットを使用して、先ほどロギングしたエージェントの評価を行ってみましょう。

```python
import mlflow

run_id = logged_agent_info.run_id
with mlflow.start_run(run_id=run_id):
    eval_results = mlflow.evaluate(
        f"runs:/{run_id}/agent",
        data=eval_dataset,  # 評価データセット
        model_type="databricks-agent",  # Mosaic AI Agent Evaluationを有効にする
    )

# MLFlow UIで評価結果を確認するか（コンソール出力を参照）、その場でアクセスする:
display(eval_results.tables['eval_results'])
```

このコードを実行すると評価が行われたのちにノートブック上に評価結果が表示されますが、せっかくなのでMLflowのUIで確認してみます。

実行すると以下のように「View evaluation results」のボタンが表示されます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/81e4fe92-d504-47b5-b0dc-fea0dcb7cff9.png)

こちらを押すと、ロギングした評価結果がUIで表示されます。


![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/f3f2fa49-35a7-45cb-ae3f-45745c300593.png)

最下部の2行が最新の評価です。
Correctness(正確性)とSafety(安全性)がどちらもPassしていますね。

青字のinputをクリックすると評価結果の詳細が表示されます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/9eedd496-fbc4-4b1c-8b30-e3fa16cd7e64.png)

評価理由含めて表示されるので、非常にわかりやすい。
ちなみに「See detailed trace view」ボタンを押すと、MLflow Tracingの結果も確認できます。
エージェントの評価内容や処理結果が明確に追えるので、評価内容に問題があったときの原因解決に有用です。

さて、単純評価ではうまくいきましたので、これをDatabricks上にデプロイしましょう。

# Step4. Agent Frameworkを使ったエージェントの公開

作成したエージェントをDatabricks Mosaic AI Model Servingで公開します。
この際にMosaic AI Agent Frameworkを使うと簡単にデプロイかつ評価用アプリから利用できるようになります。

まず、デプロイ前の検証。デプロイ時の依存関係チェックなど、事前にデプロイ時エラーが起こらないか確認できます。

```python
mlflow.models.predict(
    model_uri=f"runs:/{run_id}/agent",
    input_data={"messages": [{"role": "user", "content": "Hello!"}]},
    env_manager="uv",
)
```

問題が起きなければ、ロギングしたエージェントをUnityCatalog上に登録します。

```python
mlflow.set_registry_uri("databricks-uc")

catalog = "training"
schema = "llm"
model_name = "chat_react_agent"
UC_MODEL_NAME = f"{catalog}.{schema}.{model_name}"

# Unity Catalogへの登録
uc_registered_model_info = mlflow.register_model(
    model_uri=f"runs:/{run_id}/agent", name=UC_MODEL_NAME
)
```

最後にAgent Frameworkのデプロイ機能を使ってエージェントをMosaic AI Model Servingにデプロイする処理が実行されます。

```python
from databricks import agents
agents.deploy(UC_MODEL_NAME, uc_registered_model_info.version)
```

処理状況や結果は「サービング」メニューから確認できます。
概ね15分以内にはデプロイが終わって利用可能な状態となります。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/b8680c65-1dc4-44fe-b904-40a05cfdd58c.png)

これで作成したエージェントをAPIとして利用できるようになりました。

例えば、Playgroundからこのエージェントを指定してチャットしてみます。
簡単な計算問題を入力しました。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/dd2aeecc-bfef-4921-ad2a-631b00d9812d.png)

実行結果はこちら。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/6c9dd496-8711-489d-b431-2f1ffbe1fd4e.png)

また、専用のレビューアプリを使えるようになります。
該当エージェントのサービングメニューから「レビューアプリを開く」でアプリを開くことが出来ます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/cd820b8f-1e87-41f9-bb07-f116b8912ebb.png)


Playgroundと同様にチャットを入力すると回答を得られ、それに対するフィードバックを入力できます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/141c92ab-81b0-49c7-a967-125529174909.png)

レビューアプリについては、ユーザへの評価依頼タスクを設定することもできます。詳しくは以下のドキュメントを参照ください。

https://docs.databricks.com/aws/ja/generative-ai/agent-evaluation/review-app

# まとめ

エージェントの作成から評価・デプロイまでをMosaic AI Agent Frameworkを使って実施してみました。
また、モデルとしてはDatabricksネイティブのClaude 3.7 Sonnetを利用しています。
コードの生成に長けたモデルであり、今回のようなコード実行を行うエージェントにおいても強力な性能を発揮します。

Databricks単体で強力なLLM＋AIエージェント構築のフレームワーク全てが揃っているというこの状況は個人的にゾクゾクしています。昨年のDATA+AI SUMMITの発表から1年弱で日本リージョンの状況もかなりAIエージェントの展開に適した状況へ揃ってきた感じがしますね。
あとはDatabricks Appsも東京リージョンに来ると個人的には非常にありがたい。
[AI Builder](https://qiita.com/taka_yayoi/items/7bd37827ffff11207d14)も楽しみ。

これからの未来も期待しかありません。

# 補足：Claude 3.7 Sonnetの価格はいくら？

:::note warn
以下の内容はサイト等を見て個人が判断した内容です。
正確な情報はDatabricks社の営業等から入手されることをお薦めします。

内容に誤りがあれば指摘ください。
:::


Databricks Claude 3.7 SonnetをPay-per-tokenで利用するにあたって利用料が気になったので確認した結果も記載しておきます。

まずDBU単位の消費量はサービングメニューのモデル説明を見ると以下の通り。
(Databricks AWSの場合)

2025年3月時点では1DBUあたり$0.070かと思うので、入力1Mトークンあたり約3ドル、出力1Mトークンあたり約15ドルとなります。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/55772d58-f530-4672-bc64-00390a703e1b.png)

ただPricingのページを見ると、どうも2025年10月1日まではプロモーション価格が適用されそうに見えます。そのため少し金額感が違うかも。

https://www.databricks.com/product/pricing/foundation-model-serving

以下に上記サイトからAWS+Claude 3.7 Sonnetを指定した場合の料金表示を抜粋。
(Azureなどはまた違うので注意してください)

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/e5defe42-cb82-4319-9645-dac8faf4a03c.png)

下表でもPay-Per-TokenのDBUレートがサービングメニューの料金説明と異なり、Llama 3.1 405BBと近しくなっていたますね。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/bbe5edee-69a1-4c57-ac1e-edfa0a3982c6.png)

期間限定ではありますが、API料金は本家よりたぶん安い？
加えてDatabricksの様々な機能(AIゲートウェイなど)を考えるとお得だなあ。
Claude 3.7 Sonnetを利用するためだけにDatabricksのアカウントを持ってもいいんじゃないでしょうか（良くない）。