MLflowのPrompt Registryを使ってみる！ #mlflow #databricks - CCCMKホールディングス TECH LABの Tech Blog

こんにちは、CCCMKホールディングスAIエンジニアの三浦です。

生成AI周りの機能がアップデートされたMLflowに、ベータ版の機能で"Prompt Registry"が搭載されていました。この機能はプロンプトエンジニアリングの効率化を目的とした機能で、生成AIアプリに搭載するプロンプトの編集、バージョン管理、タグ付け等を行うことが出来ます。Prompt Registryに登録されたプロンプトはコマンド一つでダウンロードすることが可能です。

今回の記事では、現在ベータ版であるPrompt Registryを触って調べてみたことをまとめたいと思います。おそらく今後様々なアップデートが行われ、現在の仕様から変更される可能性がありますので、その点をご了承ください。

また、今回はAzure Databricksの"eastus"リージョンで試しています。

Prompt Registryで出来ること

プロンプトの管理

Prompt Registryを使うと、生成AIアプリで使用するプロンプトを保存し、管理することが出来ます。MLflowのExperimentsに紐づいて管理され、その実体はDatabricksだとUnity CatalogのSchema下に格納されます。NotebookからプロンプトをPrompt Registryに登録するには、次のようなコードで実現することが出来ます。

import mlflow

# 登録するプロンプト
supervisor_prompt = \
"""あなたはユーザーからの命令を達成するため、複数のAgentにタスクを渡す管理者Agentです。あなたの管理下のAgentと彼らに与えられた指示は以下の通りです。

- planner: {{ planner_prompt }}
- wikipedia: {{ wikipedia_prompt }}
- reporter: {{ reporter_prompt }}

あなたは管理下のAgentに対し1つずつタスクを与えます。同時に複数のAgentにタスクを渡してはいけません。
また、あなた自身はタスクを実行したり、ユーザーへの回答を生成してはいけません。"""

# 登録
mlflow.genai.register_prompt(
    # 登録先のcatalog_name, schema_name, prompt_name
    name="catalog_name.schema_name.prompt_name", 
    template=supervisor_prompt,
    commit_message="Initial commit",
    tags={
        "language": "ja",
    },
)

Unity Catalog環境のDatabricksでは、登録する際の名前nameは"カタログ名.スキーマ名.プロンプト名"の形式にする必要があります。また、登録するプロンプトにおいて変数になる部分は"{{ }}"で囲みます。LangChainだと"{ }"を使うので、少し注意が必要です。

プロンプトの編集

プロンプトが登録されると、現在実行中のNotebookに紐づいたExperimentのPrompt Registryで確認することが出来るようになります。

登録されたプロンプト一覧からプロンプトを選ぶと、以下のようにこれまでのバージョンとその内容を見ることが出来ます。

この画面から、"Create new version"をクリックすると、プロンプトを作成する画面に移ります。

私はNotebookやソースコード上でプロンプトを書くのが苦手なので、専用のプロンプトエディタがあるとうれしいな、と思います。現在の編集画面はプロンプトの新規作成しか出来ず、機能もシンプルですが、今後以前のバージョンのプロンプトを元に新しいバージョンのプロンプトを作ったり、編集機能も充実していくといいな、と思いました。

プロンプトの比較

バージョン間でプロンプトの差分を確認することが出来ます。

プロンプトへのタグ、エイリアスの登録

プロンプトにはタグやエイリアスを付与することが出来ます。エイリアスは例えば"beta", "staging", "production"といったように開発のステージに応じた文字列を付けておくと、このバージョンのプロンプトは"production"環境で使う、このバージョンは"staging"環境で使う、といった切り分けが可能になります。

プロンプトをダウンロードして使う

Prompt Registryのプロンプトをダウンロードしてアプリ内で使うことが出来ます。次のように、ダウンロードするプロンプト名とバージョンを指定します。.templateでプロンプトを取得できます。

mlflow.genai.load_prompt(
    "prompts:/catalog_name.schema_name.prompt_name/3" #3はver番号
).template

結果

'"あなたはユーザーからの命令を達成するため、複数のAgentにタスクを渡す管理者Agentです。あなたの管理下のAgentと彼らに与えられた指示は以下の通りです。\\n\\n- planner: {{ planner_prompt }}\\n- wikipedia_searcher: {{ wikipedia_prompt }}\\n- reporter: {{ reporter_prompt }}\\n\\nあなたは管理下のAgentに対し1つずつタスクを与えます。同時に複数のAgentにタスクを渡してはいけません。\\nまた、あなた自身はタスクを実行したり、ユーザーへの回答を生成してはいけません。"'

LangChainで読み込む際は変数名を"{{ }}"ではなく"{ }"で囲む必要がありますが、次のようにto_single_brace_format()を呼び出すと"{{ }}"を"{ }"に変換したプロンプトを取得することが出来ます。

mlflow.genai.load_prompt(
    "prompts:/catalog_name.schema_name.prompt_name/3"
).to_single_brace_format()

結果

'"あなたはユーザーからの命令を達成するため、複数のAgentにタスクを渡す管理者Agentです。あなたの管理下のAgentと彼らに与えられた指示は以下の通りです。\\n\\n- planner: {planner_prompt}\\n- wikipedia_searcher: {wikipedia_prompt}\\n- reporter: {reporter_prompt}\\n\\nあなたは管理下のAgentに対し1つずつタスクを与えます。同時に複数のAgentにタスクを渡してはいけません。\\nまた、あなた自身はタスクを実行したり、ユーザーへの回答を生成してはいけません。"'

また、バージョン番号でなくエイリアス指定でプロンプトを取得することも可能です。こちらの方が実用性があるかもしれません。

mlflow.genai.load_prompt(
    "prompts:/catalog_name.schema_name.prompt_name@staging"
).to_single_brace_format()

まとめ

今回はMLflowに最近追加されたPrompt Registryを触ってみて調べたことをまとめてみました。実はもう少し色々と試していたのですが、上手くいかないことや調べ切れなかったことが多く、今回は基本的な内容にとどめました。MLflowの生成AI周りの機能はまだドキュメントも少な目で、色々と詰まってしまうところも多いのですが、上手く使えると便利だろうな・・・と思っています。今後はAgentをAPIで提供できるModel Servingを活用していきたいなと考えており、このPrompt Registryも絡めていけるか調べていこうと思います。