Zenn
🍣

MLflow AI GatewayでLLMのやり取りを集約する

に公開
MLflow
LLM
AI Gateway
tech

はじめに

MFflowでローカルチャットボットを評価する方法をしてみました。前回はヒューリスティックベースのメトリクスを使って評価しました。LLM判定メトリクスを使って評価したいのですが、その前にAI Gatewayというものを理解しないといけないようなので、そちらを先にやってみます。

前回の記事

今回参照する情報

今回のコード

環境

本記事の動作確認は以下の環境で行いました。

  • MacBook Pro
  • 14 インチ 2021
  • チップ：Apple M1 Pro
  • メモリ：32GB
  • macOS：15.5（24F74）

MLflow AI Gateway

MLflow AI Gatewayは、様々なLLMの利用と管理を効率化するためにデザインされたツールです。単一のエンドポイントで特定のLLMのリクエストを処理する高レベルなインタフェースを提供します。

MLflow AI Gateway によってAPI-Keyを集約管理できることが大きなメリットです。またコード内にキーを格納することによるキーの流出も防ぐことができます。

AI Gateway は、柔軟性と適応性を備え、設定ファイルによってエンドポイントの定義・管理を行えます。つまり、LLMプロバイダやLLMタイプをインターフェースの実装を返ること無くシステムに導入することができます。

API-Keyを集中管理できシンプルにLLMの操作ができるMLflow AI Gatewayは、LLMを日常的に利用する組織にとって、離床的な選択です。

環境構築

以下の依存パッケージをインストールします。uvicornやfastapi などがインストールされ、mlflowのバックエンドとして利用します。

pipenv install "mlflow[genai,extras]"

また前回記事に従ってMFflowの環境を構築しておきます。

AI Gatewayの設定

gateway_config.yaml を作成します。ゲートウェイサーバはYAMLファイルの設定をリアルタイムに反映します。サーバーの再起動は不要です。

endpoints:
  - name: chat
    endpoint_type: llm/v1/chat
    model:
      provider: openai
      name: google/gemma-3-12b
      config:
        openai_api_base: http://localhost:1234/v1
        openai_api_key: dummy_key
        temperature: 0.7

以下でgatewayを起動します。

pipenv run mlflow gateway start --config-path gateway_config.yaml --port 5002

http://localhost:5002 からOpanAPIのドキュメントを参照できます。

MLflowサーバのエンドポイントに問い合わせをする

サポートされているものでるは以下の３つがあります。

  • 補完：入力を保管するような予測をします
  • チャット：入力を理解し、出力を会話の様式に従い返します
  • ベクトル化：画像やテキストなどをベクトル空間に転写します。それによりマルチモーダルな類似度の計算が可能となります

今回は「チャット」を使います。チャットのペイロードを作成するのは、他のモデルタイプと比べてやや複雑です。というのも、system、user、assistant の3種類の役割からの無制限のメッセージに対応する必要があるためです。MLflow AI Gateway を通じてチャットのペイロードを設定するには、messages パラメータを指定する必要があります。このパラメータには、以下の形式の辞書のリストを渡します：

{"role": "system/user/assistant", "content": "user-specified content"}

以下の内容で nlflow_gateway_chat.py を作成します。

from mlflow.deployments import get_deploy_client

client = get_deploy_client("http://localhost:5002")
name = "chat"
data = dict(
    messages=[
        # システムに「あなたはハリー・ポッターのソーティングハットです。」と伝える
        {"role": "system", "content": "You are the sorting hat from harry potter."},
        # ユーザーに「私は勇敢で、努力家で、賢く、裏切り者です。」と伝える
        {
            "role": "user",
            "content": "I am brave, hard-working, wise, and backstabbing.",
        },
        # ユーザーに「私はどのハリーポッターの家に属する可能性が高いですか？」と伝える
        {
            "role": "user",
            "content": "Which harry potter house am I most likely to belong to?",
        },
    ],
    n=3,
    temperature=0.5,
)

response = client.predict(endpoint=name, inputs=data)
print(response)

実行します。

pipenv run python3 ./mlflow_gateway_chat.py

返事が返ってきました！

{
  "id": "chatcmpl-0ed33977fvlqlzrfop3fpx8",
  "object": "chat.completion",
  "created": 1750058267,
  "model": "google/gemma-3-12b",
  "choices": [
    {
      "index": 0,
      "message": {
        "role": "assistant",
        "content": "Ah, a complex one! Let me see... Brave, hard-working, wise - those are all admirable qualities. But that *backstabbing*... that adds a rather intriguing layer, doesn't it?\n\nIt’s clear you possess ambition and a willingness to achieve your goals, no matter the cost. While Gryffindor values bravery above all else, and Hufflepuff prizes loyalty, your willingness to... let’s say *strategically navigate* situations suggests a different path. Ravenclaw values wisdom, certainly, but your actions indicate a more pragmatic application of that wisdom.\n\nNo, no... it's quite clear to me. You value power and success above all else, and are willing to use whatever means necessary to achieve them. You're not afraid to play the game, and you’ll do so with cunning and a sharp mind. \n\nTherefore... **Slytherin!**\n\nThere's much potential within you, young one. You’ll thrive amongst those who understand the value of ambition and resourcefulness. Welcome!",
        "tool_calls": null,
        "refusal": null
      },
      "finish_reason": "stop"
    }
  ],
  "usage": {
    "prompt_tokens": 47,
    "completion_tokens": 219,
    "total_tokens": 266
  }
}

choices.message.contetを見てみると以下の内容となっています。

「おや、これは複雑ですね！ええと……勇敢で、勤勉で、賢い――どれも賞賛に値する資質です。でもその“裏切り”…それが実に興味深い一面を加えていますね？
あなたには、どんな代償を払ってでも目標を達成しようとする野心と意志があるのは明らかです。グリフィンドールは勇気を最も重んじ、ハッフルパフは忠誠心を尊びますが、あなたのその……“戦略的に物事を進める”姿勢は、別の道を示しています。レイブンクローは確かに知恵を重視しますが、あなたの行動はその知恵をより実利的に使っているようです。
いいえ、はっきりしました。あなたは何よりも力と成功を重視していて、それを得るためには手段を選ばない。駆け引きを恐れず、鋭い頭脳で巧みに立ち回ることができる。
したがって……スリザリン！
あなたには大きな可能性がありますよ、若者よ。野心と機転の価値を理解する仲間たちの中で、きっとあなたは活躍することでしょう。ようこそ！」

おわりに

今回は、MLflow AI Gatewayを使って、ローカルに立ち上げたLLMを実行できることが確認できました。
最後の例では、プロンプト (上記の例だとsystemのcontent) が強力に働くことがわかりました！
次は、MLflow AI Gatewayを使って、「LLM判定メトリクス」の動作確認をしてみたいと思います。

Discussion