Databricks Mosaic AI Agent Framework & EvaluationがAWS東京リージョンにやってきました！

https://qiita.com/isanakamishiro2/items/e88f6f1e655388c51ea4#_reference-70db93113be50edcbbe7

@isanakamishiro2 さんはいつも週末にも投稿しているなーと読んでいたところ

> 加えて、いつの間にか**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)します。

![Screenshot 2025-03-30 at 17.35.36.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/011e3252-a1b9-473a-abf1-ec55992096f9.png)

本当だ。

![Screenshot 2025-03-30 at 17.30.38.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/95347c77-b91c-4db8-ab22-0956e4530a06.png)

AI Playgroundにも**エクスポート**ボタンが。なんのこっちゃと思われるかもしれませんが、このボタンがあることで、AI Playgroundで検証したエージェントをそのまま本番運用まで持っていくことができるのこのボタンなのです。詳しくは[こちら](https://docs.databricks.com/aws/ja/generative-ai/agent-framework/ai-playground-agent)。

# Mosaic AI Agent Framework & Evaluationとは

日本語で言えば、**Mosaic AIエージェントフレームワークおよび評価**となりますが、個人的には英語表記が好きです。

https://www.databricks.com/jp/product/machine-learning/retrieval-augmented-generation

RAGを含む生成AIエージェントシステムの開発、運用をシンプルにしてくれるフレームワークです。エージェントシステムの開発には様々な要素が関わるので、このフレームワークがカバーする範囲も以下のように広範です。

- エージェントのプロとタイピングと作成
- エージェントのトレース
- エージェントのサービング
- エージェントの評価
- エージェントの改善

# 動かしてみる

こちらを少し変えて動かします。

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

Mosaic AI Vector Searchのベクトル検索インデックスを作成しておきます。

![Screenshot 2025-03-30 at 17.51.51.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/dd72141e-d03f-407b-9a03-fbc5d134c204.png)

以前はコメントアウトしていた部分を設定します。

```py
# TODO: ベクトル検索インデックスを追加
VECTOR_SEARCH_TOOLS.append(
     VectorSearchRetrieverTool(
         index_name="takaakiyayoi_catalog.vector_seach.vs_qiita"
     )
)
```

## 動作確認

ベクトルストアを検索するような質問をします。

```py
AGENT.predict({"messages": [{"role": "user", "content": "DatabricksとJupyter notebookの違いをQiitaの記事から教えてください"}]})
```

ベクトルストアを検索して回答を返してくれてます！しかも、検索結果もわかりやすく確認できるようになってる。

![Screenshot 2025-03-30 at 18.09.25.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/160f03ff-322d-48fe-aba2-050fd165c0fd.png)

## エージェントの評価

以降のステップを進めていくと、評価用データセットを用いたエージェントの評価も行えます。

```py
import pandas as pd

eval_examples = [
    {
        "request": {"messages": [{"role": "user", "content": "LLMエージェントとは何ですか？"}]},
        "expected_response": None,
    }
]

eval_dataset = pd.DataFrame(eval_examples)
display(eval_dataset)
```
```py
import mlflow

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

# MLflow UIで評価結果を確認するか、以下でアクセス:
display(eval_results.tables["eval_results"])
```

以下の結果を見ると、エージェント出現前の記事しか検索されていないので結果としてはいけてないということがわかります。このように評価のステップを含めることで、エージェント改善の糸口を見つけることができます。しかし、この辺のUIも前見た時から改善されてますね。

![Screenshot 2025-03-30 at 18.11.45.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/134c2e15-9df1-453f-a880-9bffe45765fe.png)
![Screenshot 2025-03-30 at 18.12.33.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/58ac1a8d-b945-4f7b-b9a3-d5f50982d648.png)

## エージェントの登録

ここまで開発したエージェントは非常に大切なものです。Unity Catalogが提供するモデルレジストリを活用することで、エージェントのバージョン管理、共有、モデルサービングエンドポイントへのデプロイが可能となります。

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

# TODO: UCモデルのカタログ、スキーマ、およびモデル名を定義する
catalog = "takaakiyayoi_catalog"
schema = "vector_seach"
model_name = "openai_responses_agent"
UC_MODEL_NAME = f"{catalog}.{schema}.{model_name}"

# モデルをUCに登録する
uc_registered_model_info = mlflow.register_model(model_uri=logged_agent_info.model_uri, name=UC_MODEL_NAME)
```
![Screenshot 2025-03-30 at 18.16.36.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/ab82636a-c94d-42b4-8577-abe52261b1be.png)

## エージェントのデプロイ

エージェントはユーザーが使ってなんぼです。[モデルサービングエンドポイント](https://docs.databricks.com/aws/ja/machine-learning/model-serving/)にエージェントをデプロイすることで、REST API経由でさまざまなインタフェースから利用できるようになります。

以下を実行する際には、ベクトル検索インデックスやエンベディングモデルのサービングエンドポイントのアクセス権を設定するようにしてください。

```py
from databricks import agents

agents.deploy(
    UC_MODEL_NAME,
    uc_registered_model_info.version,
    tags={"endpointSource": "docs"},
    environment_vars={
        "OPENAI_API_KEY": f"{{{{secrets/{secret_scope_name}/{secret_key_name}}}}}"
    },
)
```

これで、REST API経由でエージェントにアクセスできるようになります。

![Screenshot 2025-03-30 at 19.28.01.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/8f3fb4fa-f2eb-460c-a491-5d176245a1cb.png)

AI Playgroundでも動作確認できます。

![Screenshot 2025-03-30 at 19.34.06.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/fdf5b42b-e1a4-4605-868a-037950887aee.png)

そして、これらのデータやエージェントの関係性は[リネージ](https://docs.databricks.com/aws/ja/data-governance/unity-catalog/data-lineage)として管理されます。

![Screenshot 2025-03-30 at 18.36.53.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/0c88a756-142d-4658-988d-75e74ff7751d.png)

エージェント開発に活用ください！

:::note
**注意**
サービングしたエージェントは、ベクトルストアの検索の際にエラーが起きて期待した通りに動作しなかったので後日動作確認します。
:::


### はじめてのDatabricks

[はじめてのDatabricks](https://qiita.com/taka_yayoi/items/8dc72d083edb879a5e5d)

### Databricks無料トライアル

[Databricks無料トライアル](https://databricks.com/jp/try-databricks)