watsonx Orchestrate上のAgentで複数Geminiモデルでロードバランシングしてみる

watsonx Orchestrateでカスタム・モデルがサポートされました

watsonx Orchestrate上で動作するモデルとしてgranite,llama,mistralを提供していますが、ADK1.4.2より外部のモデルをカスタム・モデルとして利用することが可能になりました。また、ADK1.5.0より、Model Policyがサポートされ、複数モデル間のロード・バランシングやフォール・バックが可能になりました。この記事ではこれらの機能を利用して複数のGeminiモデルを動的に切り替える構成を検証します。

カスタム・モデルの定義

カスタム・モデルのプロバイダーとしては以下のものをサポートしています。

• OpenAI
• Anthropic
• Google
• watsonx.ai
• Mistral
• OpenRouter
• Ollama

パラメータの指定

ADK1.4.2ではenvファイルを用いて各LLMに必要なパラメータを定義する必要がありましたが、1.5.0より、Connectionsを用いて定義することが可能になっています。（ファイルによる詳細の指定も引き続き利用可能です。）例えば、Geminiの場合の必須パラメータはapi_keyですが、以下の様に指定します。

orchestrate connections add -a google_creds
orchestrate connections configure -a google_creds --env draft -k key_value -t team
orchestrate connections set-credentials -a google_creds --env draft -e "api_key=my_api_key"

それぞれのプロバイダー毎に指定可能なパラメータは下記ページを参照してください。
https://developer.watson-orchestrate.ibm.com/llm/managing_llm

モデルの追加

モデルの追加は、models add コマンドを使用します。

orchestrate models add --name google/gemini-2.0-flash-lite --app-id google_creds

2つのモデルを追加し、追加されたモデルを確認してみます。

orchestrate models list

model policyの定義

model policyを用いることでモデルの振る舞いを指定することが可能です。具体的には

• loadbalance
• fallback
• single
  の3つの指定が可能です。
  今回は以下のコマンドを使用して、先ほど登録した2つのモデルを使ったloadbalanceの指定をしてみましょう。

orchestrate models policy add --name loadbalance_policy --model virtual-model/google/gemini-2.0-flash-lite --model virtual-model/google/gemini-2.0-flash --strategy loadbalance --strategy-on-code 500 --retry-on-code 503  --retry-attempts 1

(※strategy-on-code、--retry-on-code、--retry-attemptsは論理的には必須ではないのですが、必須パラメータになっているようで指定しないとエラーになったため、とりあえず指定しています。)
loadbalanceは複数のモデル間でロード・バランスするストラテジーです。デフォルトでは1:1にバランシングされるため、2つのモデルが交互に呼び出されるはずです。
なお、作製したPolicyは、virtual-policy/loadbalance_policyという名前のモデルとして登録されていました。

テスト

Agentのモデルとして、作成したポリシーを指定してテストをしてみましょう。今回はこちらの記事で紹介したlangfuseを使って使用されているモデルを確認します。
チャットからAgentとやりとりをするたびにLLMが呼び出されますが、以下の様に2つのモデルが使用されていることが分かります。

fallbackを試してみる

次にfallbackについても試してみましょう。例えば、Geminiはサーバー側の負荷が高い場合に503のエラーが返ってくることがあるのですが、その場合に別のプロモデルを使用するようなポリシーを実装することが可能です。503エラーを再現するのが難しいため、今回はapi_keyを誤ったもので設定し、エラーの場合に別のモデルを使用するように構成してみたいと思います。

先ほどと同じようにカスタム・モデルを構成しますが、gemini-2.0-flash側だけ別のapp-idを指定して誤ったapi_keyが使われるようにします。次に、以下のコマンドでfallback設定を持つポリシーを定義します。

orchestrate models policy add --name fallback_policy --model virtual-model/google/gemini-2.0-flash --model virtual-model/google/gemini-2.0-flash-lite --strategy fallback --strategy-on-code 400 --retry-on-code 503  --retry-attempts 1

このポリシーを使用したAgentを作成してチャットから試してみたところ、エラーなく動作しました。何度か順番にapi_keyを変更してみてlangfuse上で確認もしましたが、複数モデルを定義した場合、先に定義したものから使われるようです。また、残念ながらfallbackの動作は、トレースからは判断できないようです。おそらくライブラリ内でハンドリングしてしまっているのだと思いますが、これについては改善をしてほしいところです。
その他気になった点としては、strategy-on-codeについては指定したもの以外の場合であってもエラーの場合には常にfallbackが発動するように見えました。(確認したところ、この辺りの挙動については既知のバグのようです。）

まとめ

watsonx Orchestrateでは、これまで利用可能だったgranite/llama/mistralに加えて、OpenAI/Googleなどの外部のモデルをカスタム・モデルとして使用してAgentを動作させることが可能になりました。また、ポリシーを定義することによって、外部LLMを使用した際のエラー時の挙動などを細かく指定し、可用性を高めることが可能です。