MLflow Study: Prompt Registry in Databricks

# はじめに

MLflow3がリリースされ、Databricks側でもDocumentの対応含めて様々な機能が実装されています。そのうちのひとつにMLflow Prompt Registryがあります。
今回は勉強として、実際に触りながら学んでみます。

https://docs.databricks.com/aws/ja/mlflow3/genai/prompt-version-mgmt/prompt-registry/create-and-edit-prompts

なお、上記の公式Documentに記載があるように、Databricksでは2025年6月現在**ベータ版**となっています。利用にはご注意ください。

検証はDatabricks Free Editionを使って行いました。

# MLflow Prompt Registryとは？

OSS版のドキュメントより。
以下にサイト内容を一部邦訳して抜粋します。

https://mlflow.org/docs/latest/genai/prompt-version-mgmt/prompt-registry/

> **MLflow Prompt Registry**は、生成AI（GenAI）アプリケーションにおけるプロンプトエンジニアリングと管理を効率化する強力なツールです。組織内でプロンプトをバージョン管理・追跡・再利用できるため、プロンプト開発における一貫性の維持やコラボレーションの向上に役立ちます。
> 
> - **再利用性**：プロンプトを一元管理し、複数のアプリケーションで再利用
> - **バージョン管理**：Git風のコミットベースのバージョン管理と、差分比較（diffハイライト）によるプロンプトバージョンの進化の追跡
> - **エイリアス**：柔軟なデプロイパイプラインを構築し、プロンプトバージョンをアプリケーション本体から分離。A/Bテストやロールバックも容易
> - **系譜（Lineage）**：MLflowの既存機能（モデル追跡や評価）とシームレスに連携し、GenAIライフサイクル全体を管理
> - **コラボレーション**：組織全体でプロンプトを共有し、チーム間でのプロンプト開発を促進
> 
> MLflow Prompt RegistryのDatabricks Unity Catalogとの連携は近日公開予定です。
> 
> ---
> 
> ## Unity Catalogとの統合
> 
> Prompt Registryは、Databricks Unity Catalogのガバナンスフレームワークを活用し、以下の機能を提供します。
> 
> - **一元化されたプロンプト管理**：プロンプトテンプレートを統制・検索可能なレジストリで管理
> - **バージョン管理**：プロンプトの変更履歴を完全な系譜とロールバック機能で追跡
> - **アクセス制御**：Unity Catalogの権限システムを利用し、プロンプトの閲覧・編集・利用を制御
> - **コラボレーション**：チーム間でのプロンプト開発の共有と協働を可能に
> - **ガバナンス**：分類やコンプライアンス要件を含むデータガバナンスポリシーをプロンプトテンプレートに適用
> 
> ---
> 
> ## 詳細なアクセス制御
> 
> - **個別・チーム・組織単位でのアクセス制御**
> - **データの系譜追跡**：プロンプトが異なるアプリケーションや実験でどのように利用されているかを追跡
> - **監査証跡**：プロンプトへのアクセスや変更を監視
> - **分類とタグ付け**：メタデータでプロンプトを整理し、ガバナンスポリシーを適用
> - **ワークスペース間での安全な共有**: 異なるDatabricksワークスペース感での安全な共有
> 
> ---
> 
> ## 今後の機能
> 
> 今後、Prompt Registryでは以下の機能がサポートされる予定です。
> 
> - **テンプレート管理**：変数置換機能付きでプロンプトテンプレートを作成・編集・バージョン管理
> - **テストと検証**：デプロイ前に評価データセットでプロンプトをテスト
> - **A/Bテスト**：異なるプロンプトバージョンを比較し、パフォーマンスを最適化
> - **MLflow Trackingとの連携**：プロンプトバージョンと実験・モデル実行を自動的に関連付け
> - **APIアクセス**：自動化ワークフローでプロンプトをプログラムから利用

というわけで生成AI用のプロンプトテンプレートを管理する機能です。
モデルと同様にテンプレートのバージョン管理ができ、将来的にはA/Bテストにも対応する模様。

プロンプトはテキスト表現なのでGit等でバージョン管理しやすいものではありますが、MLflowに専用の仕組みを設けることでさらに利便性が高まる・・・のかな。

もともとOSS版Prompt Registryはmlflow 2.xのときから存在していたのですが、MLflow3からDatabricks側でも対応された（まだベータですが）ようです。

では、実際にPrompt Registryを使ってみます。

# 基本的な利用

まずはノートブックを作成し、MLflowの最新版をインストールします。
ついでに```databricks-langchain```と```openai```パッケージもインストールしておきます。

```python
%pip install -U -q "mlflow[databricks]>=3.1.0" databricks-langchain openai
```

## プロンプトの登録

では、プロンプトテンプレートを作成してPrompt Registry上に登録してみます。

```python
import mlflow

catalog = "workspace"
schema = "default"
prompt_name = "sample_prompt"
uc_prompt_name = f"{catalog}.{schema}.{prompt_name}"

# Register a prompt template
prompt = mlflow.genai.register_prompt(
    name=uc_prompt_name,
    template="あなたは親切なアシスタントです。次の問いに答えてください: {{question}}",
    commit_message="Initial Commit"
)
print(f"Created version {prompt.version}")  # "Created version 1"
```

実行するとUnity Catalog上にプロンプトが登録されます。
・・・とは言っても、執筆時点ではカタログエクスプローラ上には表示されませんでした。

Databricks UI上のどこに表示されるかというと、「エクスペリメント」メニューから確認することができます。
上記を実行したノートブックのエクスペリメントを確認すると、「プロンプト」というタブがあり、そこから確認することができます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/5fb00a9c-4404-4f70-82ea-34cabd20db80.png)

プロンプト名のリンクを開くと、各バージョンごとに登録されているプロンプトテンプレートを確認することができます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/48efe642-bb2f-4f98-b554-ba72f33fa966.png)

ここから登録したプロンプトを使用するサンプルコードを取得したり、プロンプトの削除を行うこともできます。

## プロンプトの利用

では、登録されたプロンプトを取得・利用してみます。
取得は```mlflow.genai.load_prompt```関数を利用します。
Unity Catalogに登録されたプロンプトは```"prompts:```から始まるURIになるようです。
ここで、パスとバージョン番号を指定することでプロンプトを読み込めます。

```python
prompt = mlflow.genai.load_prompt(name_or_uri=f"prompts:/{uc_prompt_name}/1")

prompt
```

```:出力結果
PromptVersion(name=workspace.default.sample_prompt, version=1, template="\u3042\u306a\u305f\u306f\u89a...)
```

結果は```PromptVersion```というクラスのインスタンスとして得られます。

ちなみに、```search_prompts```関数を使うと、すべてのプロンプトバージョン情報を得ることができます。

```python
prompts = mlflow.genai.search_prompts(filter_string=f"catalog='{catalog}' and schema='{schema}'")
prompts
```

```:出力結果
[<PromptInfo: name='workspace.default.sample_prompt', description='Initial Commit', tags={'PromptVersionCount': '1'}>]
```

こちらは、```PromptInfo```というクラスのインスタンスリストで返されます。
実際のプロンプトは、この結果を利用して```load_prompt```を実行する必要がありそう。

では、得られたプロンプトを使ってみましょう。

```python
import mlflow
from openai import OpenAI

# MLflowの自動ロギングを有効にして、アプリケーションにトレースを追加
mlflow.openai.autolog()

# 実行ノートブックと同じ資格情報を使用してOpenAIクライアント経由でDatabricks LLMに接続
mlflow_creds = mlflow.utils.databricks_utils.get_databricks_host_creds()
client = OpenAI(
    api_key=mlflow_creds.token,
    base_url=f"{mlflow_creds.host}/serving-endpoints"
)

# レジストリからプロンプトをロード
prompt = mlflow.genai.load_prompt(f"prompts:/{uc_prompt_name}/1")
formatted_prompt = prompt.format(question="Databricksとは何ですか?")

# LLMを呼び出す
response = client.chat.completions.create(
    model="databricks-llama-4-maverick",
    messages=[
        {
            "role": "user",
            "content": formatted_prompt,
        },
    ],
)
response.choices[0].message.content
```

実行すると、以下のようなTracing結果が得られます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/52a8f0a0-c28b-4b21-a20b-dec31e4653e9.png)

```PromptVersion```クラスは```format```メソッドを提供しており、そこにパラメータを渡すことで最終プロンプトを作成することができるようになっています。

同じようにLangChainでも利用してみましょう。

```python
import mlflow
from databricks_langchain import ChatDatabricks
from langchain_core.prompts import ChatPromptTemplate
import codecs

# MLflowの自動ロギングを有効化
mlflow.langchain.autolog()

prompt = mlflow.genai.load_prompt(name_or_uri=f"prompts:/{uc_prompt_name}/1")

# LangChainのプロンプトテンプレートをMLflowのPromptから作成する
langchain_prompt = ChatPromptTemplate.from_messages(
    [
        (
            "user",
            codecs.decode(prompt.to_single_brace_format(), "unicode-escape"),
        ),
    ]
)

llm = ChatDatabricks(model="databricks-llama-4-maverick")
chain = langchain_prompt | llm

# チェーンを実行
chain.invoke({"question": "Databricksの特徴を教えてください"})
```

実行すると、以下のようなTracing結果が得られます。

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/168727/6cda0fd0-5bd4-4f1f-97ad-9a4d11a2ed2c.png)

```to_single_brace_format```を使うことで、LangChainの```ChatPromptTemplate```に対応した文字列としてプロンプトを出力することができます。
（内部的には{{のような2重ブラケットを{という単一ブラケットに変換しているだけだと思います）


## プロンプトにエイリアスを設定する（未動作）

ドキュメントによれば、Prompt Registryに登録したプロンプトに対してエイリアスを設定することができるようです。

https://docs.databricks.com/aws/ja/mlflow3/genai/prompt-version-mgmt/prompt-registry/use-prompts-in-deployed-apps

ただ、執筆時点では正常に動作しませんでした。

コードとしては、以下のように```set_prompt_alias```関数を使うことで、特定のバージョンに対してエイリアスを設定することができます。

```python
import mlflow
mlflow.genai.set_prompt_alias(
    name=f"{uc_prompt_name}",
    alias="production",
    version=1
)
```

エイリアスを設定したプロンプトは、```prompts/catalog.schema.prompt_name@alias```というURIでプロンプトを指定できるようになり、デプロイしたコードを変更せずに改善したプロンプトを反映させることができるようになります。

上のコードで設定したエイリアスを使ってプロンプトを取り出すコードは以下のようになります。（執筆時点では、実行するとdoes not existエラーが発生しました）

```python
mlflow.genai.load_prompt(f"prompts:/{uc_prompt_name}@production")
# 執筆時点ではエラーが発生
```

エイリアスは有用な機能ですので、実装が待ち遠しいです。

# おわりに

簡単ですが、DatabricksのPrompt Registryを実際に使いながら学んでみました。
ベータ版ということでOSS版に比べて未実装な部分がまだ多いようです。
管理系だとプロンプト削除のメソッドが（おそらく）まだ無く、コードからではなく、UIからの削除に限られています。

LLMを利用する上でプロンプト管理は依然まだ重要だと思いますので、機能の充足が待ち遠しいです。