DatabricksにおけるRAGのデプロイとレビュー

Last updated at 2024-11-11Posted at 2024-11-10

こちらの続編です。

前回は、RAGを構成する機能をツールとして実装し、AI Playgroundで動作を確認しました。プロダクションに投入できる見通しが立ったらデプロイしましょう。ここでは、Mosaic AI Agent FrameworkとAgent Evaluationを活用します。

Agent Framework + Agent Evaluationと連携する機能がすでに用意されています。AI Playgroundの画面の上部にエクスポートボタンがあります。こちらをクリックすることで、選択されている言語モデル、ツール、AI Playground上でのリクエストやレスポンスを取り込んで、Pythonの実装ノートブックが生成されます。

以下のようなドライバーノートブックに遷移します。コメントを翻訳しながら実行していきます。

ドライバーノートブック

これはAIプレイグラウンドエクスポートによって自動生成されたノートブックです。同じフォルダーに3つのノートブックを生成しました：

agent: エージェントを構築するコードが含まれています。
config.yml: 設定が含まれています。
driver: エージェントのログ、評価、登録、およびデプロイを行います。

このノートブックは、agentノートブックで定義されたエージェントをデプロイするためにMosaic AIエージェントフレームワーク（AWS | Azure）を使用します。このノートブックは以下を行います：

エージェントをMLflowにログします
エージェントをUnity Catalogに登録します
エージェントをモデルサービングエンドポイントにデプロイします

前提条件

このノートブックのすべてのTODOを解決してください。
config.ymlの内容を確認してください。これはエージェントが利用できるツールとLLMエンドポイントを定義しています。
このフォルダーのagentノートブックを確認し、エージェントのコードを反復し、出力をテストしてください。

config.yml

このエージェントでは、独自のエンべディングモデルとベクトル検索インデックスを使用しています。これをエージェントの設定に含めるように、以下のconfig.ymlではvector_search_indexとembedding_endpointのエントリーを追加しています。

config.yml

agent_prompt: "あなたはDatabricksに関する質問に回答するためにツールを活用する有能なアシスタントです。提供されているツールを活用してください: Databricksに関する質問に回答するために、vector_searchとgenerate_responseを活用してください。"
llm_endpoint: "taka-openai-gpt-4o-mini"
vector_search_index: "takaakiyayoi_catalog.qiita_2023.qiita_2023_vs"
embedding_endpoint: "taka-japanese-embedding"
tools:
  uc_functions:
    - "takaakiyayoi_catalog.qiita_2023.generate_response"
    - "takaakiyayoi_catalog.qiita_2023.vector_search"

次のステップ

エージェントがデプロイされた後、AIプレイグラウンドでエージェントとチャットして追加のチェックを行ったり、組織内のSMEと共有してフィードバックを得たり、プロダクションアプリケーションに埋め込んだりすることができます。詳細はドキュメント（AWS | Azure）を参照してください。

%pip install -U -qqqq mlflow databricks-agents
dbutils.library.restartPython()

`agent`をMLflowモデルとしてログする

agentノートブックのコードをエージェントとしてログします。詳細はMLflow - Models from Codeを参照してください。

注意
config.ymlで定義したカスタムのエンべディングモデルサービングエンドポイントとベクトル検索インデックスをresourcesに追加しています。これらの設定を行わないと、デプロイされたエージェントからそれらのリソースにアクセスする際に、[REMOTE_FUNCTION_HTTP_FAILED_ERROR] The remote HTTP request failed with code 403, and error message 'HTTP request failed with status: {"error_code":"PERMISSION_DENIED","message":"Insufficient permissions for UC entity と言うような権限不足のエラーが発生します。

# モデルをMLflowにログする
import os
import mlflow

from mlflow.models import ModelConfig
from mlflow.models.signature import ModelSignature
from mlflow.models.rag_signatures import (
    ChatCompletionRequest,
    ChatCompletionResponse,
)
from mlflow.models.resources import (
    DatabricksFunction,
    DatabricksServingEndpoint,
    DatabricksVectorSearchIndex
)
from databricks.sdk import WorkspaceClient

w = WorkspaceClient()
config = ModelConfig(development_config="config.yml")
resources = [
    DatabricksServingEndpoint(endpoint_name=config.get("llm_endpoint")),
    DatabricksServingEndpoint(endpoint_name=config.get("embedding_endpoint")),
    DatabricksVectorSearchIndex(index_name=config.get("vector_search_index")),
]
uc_functions_to_expand = config.get("tools").get("uc_functions")
for func in uc_functions_to_expand:
    # 関数名が正規表現の場合、スキーマ内のすべての関数を取得
    if func.endswith("*"):
        catalog, schema, _ = func.split(".")
        expanded_functions = list(
            w.functions.list(catalog_name=catalog, schema_name=schema)
        )
        for expanded_function in expanded_functions:
            resources.append(
                DatabricksFunction(function_name=expanded_function.full_name)
            )
    # それ以外の場合は関数を追加
    else:
        resources.append(DatabricksFunction(function_name=func))

signature = ModelSignature(ChatCompletionRequest(), ChatCompletionResponse())

input_example = {
    "messages": [{"role": "user", "content": "find me franchises in the US"}]
}

with mlflow.start_run():
    logged_agent_info = mlflow.pyfunc.log_model(
        "agent",
        python_model=os.path.join(
            os.getcwd(),
            "agent",
        ),
        signature=signature,
        input_example=input_example,
        model_config="config.yml",
        resources=resources,
    )

エージェント評価を使用してエージェントを評価する

評価データセットのリクエストや期待される応答を編集し、エージェントを反復しながら評価を実行できます。mlflowを活用して計算された品質指標を追跡します。

import pandas as pd

eval_examples = [
    {
        "request": {
            "messages": [
                {
                    "role": "user",
                    "content": "DeltaとParquetの違い"
                }
            ]
        },
        "expected_response": None
    },
    {
        "request": {
            "messages": [
                {
                    "role": "user",
                    "content": "DeltaとParquetの違い"
                },
                {
                    "role": "assistant",
                    "content": "DeltaとParquetの違いについての情報を探すために、関連する文書を検索します。\n\n<uc_function_call>{\"id\":\"mFQnD8oosIYhAQNd60sr\",\"name\":\"takaakiyayoi_catalog.qiita_2023.vector_search\",\"arguments\":\"{\\\"query_string\\\":\\\"Delta Parquet 違い\\\"}\"}</uc_function_call>\n\n<uc_function_result>{\"id\":\"mFQnD8oosIYhAQNd60sr\",\"content\":\"{\\\"is_truncated\\\":false,\\\"columns\\\":[\\\"context\\\",\\\"url\\\"],\\\"rows\\\":[[\\\"パーティションは特に非常に大きなテーブルにおいては有効です。パーティショニングに関する数多くのパフォーマンス強化は、非常に大きなテーブル(数百TB以上)にフォーカスしています。  \\\\n多くのお客様は、ParquetベースのデータレイクからDelta Lakeに移行しています。CONVERT TO DELTA文によって、既存のデータを再度書き込むことなしに、ParquetベースのテーブルをDeltaテーブルに変換することができます。このようにして、多くのお客様は以前のパーティション戦略を継承する大規模なテーブルを手にすることができます。Databricksによって開発されるいくつかの最適化処理は、可能な限りこれらのパーティションを活用しようとし、Delta Lake向けに最適化されていないパーティション戦略のいくかの副作用を軽減します。  \\\\nDelta LakeとApache Sparkはオープンソーステクノロジーです。Databricksでは、パーティションへの依存度を削減する機能の導入を進めますが、オープンソースコミュニティでは、複雑性を増加させるような新機能を構築し続けるかもしれません。\\\",\\\"https://qiita.com/taka_yayoi/items/b50beddb8a687cd243c6\\\"]]}\"}</uc_function_result>\n\nDeltaとParquetの違いについて詳しく説明するために、関連する文書を参照して具体的な情報を提供します。\n\n<uc_function_call>{\"id\":\"NQc7v7niAToHXDWA05jN\",\"name\":\"takaakiyayoi_catalog.qiita_2023.generate_response\",\"arguments\":\"{\\\"query_string\\\":\\\"DeltaとParquetの違い\\\",\\\"context\\\":\\\"パーティションは特に非常に大きなテーブルにおいては有効です。パーティショニングに関する数多くのパフォーマンス強化は、非常に大きなテーブル(数百TB以上)にフォーカスしています。  \\\\n多くのお客様は、ParquetベースのデータレイクからDelta Lakeに移行しています。CONVERT TO DELTA文によって、既存のデータを再度書き込むことなしに、ParquetベースのテーブルをDeltaテーブルに変換することができます。このようにして、多くのお客様は以前のパーティション戦略を継承する大規模なテーブルを手にすることができます。Databricksによって開発されるいくつかの最適化処理は、可能な限りこれらのパーティションを活用しようとし、Delta Lake向けに最適化されていないパーティション戦略のいくかの副作用を軽減します。  \\\\nDelta LakeとApache Sparkはオープンソーステクノロジーです。Databricksでは、パーティションへの依存度を削減する機能の導入を進めますが、オープンソースコミュニティでは、複雑性を増加させるような新機能を構築し続けるかもしれません。\\\",\\\"url\\\":\\\"https://qiita.com/taka_yayoi/items/b50beddb8a687cd243c6\\\"}\"}</uc_function_call>\n\n<uc_function_result>{\"id\":\"NQc7v7niAToHXDWA05jN\",\"content\":\"{\\\"is_truncated\\\":false,\\\"columns\\\":[\\\"response\\\"],\\\"rows\\\":[[\\\"DeltaとParquetの主な違いは、Delta Lakeがトランザクション管理やスキーマの進化、データのバージョニングをサポートしている点です。特に、Delta Lakeはパーティショニングを活用してパフォーマンスを向上させる機能があり、非常に大きなテーブルにおいては特に効果的です。多くの顧客は、既存のParquetベースのデータレイクからDelta Lakeに移行する際、`CONVERT TO DELTA`文を使用して、再書き込みなしでテーブルを変換し、以前のパーティション戦略を維持しています。Databricksは、これらのパーティションを最大限に活用するための最適化を進めていますが、オープンソースコミュニティでは新たな機能が複雑性を増す可能性もあります。\\\"]]}\"}</uc_function_result>\n\nDeltaとParquetの主な違いは以下の通りです。\n\n1. **トランザクション管理**: Delta Lakeはトランザクション管理をサポートしており、データの整合性を保ちながら複数のユーザーが同時にデータにアクセスできるようにします。一方、Parquetはトランザクション管理の機能を持っていません。\n\n2. **スキーマの進化**: Delta Lakeはスキーマの進化をサポートしており、データのスキーマを変更する際に柔軟性があります。Parquetでは、スキーマの変更が難しい場合があります。\n\n3. **データのバージョニング**: Delta Lakeはデータのバージョニングを提供し、過去のデータに簡単にアクセスできるようにします。これにより、データの変更履歴を追跡することが可能です。Parquetはこの機能を持っていません。\n\n4. **パフォーマンスの最適化**: Delta Lakeはパーティショニングを活用してパフォーマンスを向上させる機能があり、特に非常に大きなテーブルにおいて効果的です。多くの顧客は、既存のParquetベースのデータレイクからDelta Lakeに移行する際、`CONVERT TO DELTA`文を使用して再書き込みなしでテーブルを変換し、以前のパーティション戦略を維持しています。\n\nこれらの違いにより、Delta Lakeはデータの管理や分析においてより高度な機能を提供します。詳しくは[こちらのリンク](https://qiita.com/taka_yayoi/items/b50beddb8a687cd243c6)をご覧ください。"
                },
                {
                    "role": "user",
                    "content": "Sparkの特徴を5つ挙げてください"
                }
            ]
        },
        "expected_response": None
    }
]

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

import mlflow
import pandas as pd

# 評価を開始するためにMLflowのランを開始
with mlflow.start_run(run_id=logged_agent_info.run_id):
    eval_results = mlflow.evaluate(
        f"runs:/{logged_agent_info.run_id}/agent",  # log_modelを呼び出す際に使用したartifact_pathに`chain`を置き換えます。
        data=eval_dataset,  # 評価データセット
        model_type="databricks-agent",  # Mosaic AI Agent Evaluationを有効にする
    )

# MLflow UIで評価結果を確認するか、以下のようにアクセスします:
display(eval_results.tables['eval_results'])

こちらでウォークスルーしたように、AI Playgroundでのやり取りを評価データセットとして評価します。これが、エージェントシステムを開発する際に最初から評価の仕組みを組み込んでおくという評価ドリブン開発ワークフローです。

モデルをUnity Catalogに登録

以下のcatalog、schema、model_nameを更新して、MLflowモデルをUnity Catalogに登録します。

mlflow.set_registry_uri("databricks-uc")

# TODO: UCモデルのカタログ、スキーマ、およびモデル名を定義します
catalog = "takaakiyayoi_catalog"
schema = "qiita_2023"
model_name = "databricks_rag"
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)

エージェントをデプロイ

from databricks import agents

# モデルをレビューアプリとモデルサービングエンドポイントにデプロイします
agents.deploy(UC_MODEL_NAME, uc_registered_model_info.version)

Successfully registered model 'takaakiyayoi_catalog.qiita_2023.feedback'.
Created version '1' of model 'takaakiyayoi_catalog.qiita_2023.feedback'.
2024/11/10 11:12:04 INFO mlflow.tracking._tracking_service.client: 🏃 View run feedback-model at: https://xxxxxx.databricks.com/ml/experiments/117686861228987/runs/8e00f1a5b8874c26a64369f0b1071151.
2024/11/10 11:12:04 INFO mlflow.tracking._tracking_service.client: 🧪 View experiment at: https://xxxxxx.databricks.com/ml/experiments/117686861228987.

    Deployment of takaakiyayoi_catalog.qiita_2023.databricks_rag 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://xxxxxx.databricks.com/ml/endpoints/agents_takaakiyayoi_catalog-qiita_2023-databricks_rag
    Review App: https://xxxxxx.databricks.com/ml/review/takaakiyayoi_catalog.qiita_2023.databricks_rag/1?o=1444828305810485

これでRAGエージェントとレビューアプリがデプロイされます。