DatabricksにおけるAIエージェントのデプロイメントおよび評価

こちらの続きです。

https://qiita.com/taka_yayoi/items/1e15e5dd0a191352eced

前回はツールを構築してAI Playgroundで動かすところまでをカバーしました。やり取りを通じてエージェントの目鼻顔立ちが見えてきたら、エージェントシステムとしてデプロイしたくなることでしょう。

その際に活用するのが**エクスポート**ボタンです。

![Screenshot 2024-10-21 at 9.58.37.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/73d3f361-8eb1-bd54-d425-3f631bd8586d.png)

[こちら](https://docs.databricks.com/ja/generative-ai/create-log-agent.html#export-and-deploy-ai-playground-agents)にあるように、前のステップで定義したツールが組み込まれたエージェントシステムの実装を含むノートブックが生成されます。必要に応じてカスタマイズして本格運用に進むことができます。これも、Databricks AutoMLと同様の[ガラスボックスアプローチ](https://www.databricks.com/jp/product/automl)となっています。

- `agent` ノートブック: を使用してエージェントを定義するPython LangChainコードが含まれています。
- `driver` ノートブック: エージェントフレームワークを使用してPython エージェントをログに記録、トレース、登録する デプロイするためのAI Mosaic AIコードが含まれています。
- `config.yml`: ツール定義を含む、エージェントに関する構成情報が含まれます。

# Databricksにおけるエージェント開発のアプローチ

ウォークスルーの前に、Databricksにおけるエージェント開発のアプローチに触れさせてください。RAGを構築された方はすでに体験されていると思いますが、このようなシステムにおいて本格運用に十分な品質を確保することは相当に難しいこととなっています。理由としては、以下のようなものが挙げられます。

- 評価手法自体が確立されているフェーズ
- 調整すべきコンポーネントが多岐にわたる

![Screenshot 2024-10-21 at 10.08.44.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/99277167-61b3-3131-d4ed-ebceecc35427.png)

このような状況で闇雲にPOCを繰り返しても、望む結果を得ることができません。このため、Databricksでは開発時点から評価の仕組みを組み込む**評価ドリブン開発ワークフロー**を採用しています。

![Screenshot 2024-10-21 at 10.10.53.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/26bf49da-76a7-f720-8522-7532a71f88e9.png)

このワークフローを実現しているのは、以下のDatabricksのコンポーネントです。

- [Mosaic AI Agent Framework](https://docs.databricks.com/ja/generative-ai/retrieval-augmented-generation.html)
    - エージェントを作成、MLflowで記録
    - 1行のコードでエージェントを
    - モデルサービングエンドポイントにデプロイ
    - ステークホルダーからフィードバックを収集するためのレビューアプリを自動で作成
- [Mosaic AI Agent Evaluation](https://docs.databricks.com/ja/generative-ai/agent-evaluation/evaluate-agent.html)
    - エージェントを評価するためにプロプライエタリなLLMジャッジを活用
    - レーテンシーやトークンのコストのようなパフォーマンスメトリクスを収集
- [AI Playground](https://docs.databricks.com/ja/large-language-models/ai-playground.html)
    - 様々なLLMに対するプロンプトエンジニアリング
    - 複数のモデルの挙動を比較
    - ツールの組み込み

前振りが長くなりましたが、最初に触れた**エクスポート**ボタンを押すことで、上のコンポーネントとのインテグレーションが含まれたノートブックが作成されるということです。

# エージェントのデプロイメントのウォークスルー

エクスポートされたノートブックは、ホームフォルダ配下の`Agent databricks-meta-llama-3-1-70b-instruct 0006-10-21 08:59:08`というようなフォルダが作成され、そちらに格納されます。

**config.yml**

使用するLLM、ツールが記載されています。

```yaml:config.yml
llm_endpoint: "databricks-meta-llama-3-1-70b-instruct"
tools:
  uc_functions:
    - "takaakiyayoi_catalog.cookie.*"
```

**agent**

エージェントの実装です。このノートブックは直接実行しません。後述の`driver`ノートブックからアクセスします。ローカルでテストする際には、ノートブックを実行することが可能です。

```py
%pip install -U -qqqq mlflow-skinny
dbutils.library.restartPython()
```

```py
import mlflow
from mlflow.models import ModelConfig

config = ModelConfig(development_config="config.yml")
```

```py
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()
```

```py
# 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?")]
    ),
)
```

```py
mlflow.models.set_model(agent)
```

**driver**

Mosaic AI Agent Frameworkを用いて、`agent`ノートブックで実装されているエージェントの記録、評価、サービングエンドポイントへのデプロイなどをエンドツーエンドで実行します。

なお、AI Playgroundでやりとりした内容が評価データセットとして盛り込まれるようになっています。

Unity Catalogに登録するカタログ、スキーマ、モデル名を指定します。

```py
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)
```

最後のコマンドを実行すると、モデルサービングエンドポイントにエージェントのモデルとレビューアプリがデプロイされます。

```py
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)
```

![Screenshot 2024-10-21 at 9.56.20.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/6fa690ba-97e9-84ae-0be7-5aba33d86a62.png)

モデルサービングエンドポイントがReadyになったら、レビューアプリにアクセスして挙動を確認します。

![Screenshot 2024-10-21 at 9.49.26.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1168882/95523725-f8f4-b96a-1fee-414e5ee5d943.png)

エージェントにおけるリクエストとレスポンス、レスポンスに対するフィードバックも[推論テーブル](https://docs.databricks.com/ja/machine-learning/model-serving/inference-tables.html)に記録されるので、それらを活用して精度を改善することができます。


### はじめてのDatabricks

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

### Databricks無料トライアル

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