litellmを利用してAzure OpenAIとAmazon BedrockのチャットAPIを共通的に呼び出す

背景

• 社内で AI チャットサービスを展開する際、GPT-4 と GPT-3.5 などでモデル切り替えを可能にすることが多い
• モデルの切り替えを、Azure OpenAI と Amazon Bedrock などベンダーをまたぐ場合に API 呼び出し分ける必要があり手間になる
• litellm というライブラリがその問題を解消できそうなので試してみる
  • https://github.com/BerriAI/litellm
  • Open Interpreter でも採用されている
    • https://github.com/KillianLucas/open-interpreter/#change-your-language-model

ライブラリインストール

以下ライブラリをインストール。pip や pipenv などで

litellm
jinja2
boto3

※ litellm のバージョンは 0.1.813 より新しいものを使う。Bedrock が使えないバグが修正されている

認証

今回は pipenv を利用するので .env ファイルを準備。
必要な認証情報を選択

AZURE_API_KEY=<Azure API キー>
AZURE_API_BASE=<APIエンドポイント>
AZURE_API_VERSION=2023-05-15

AWS_REGION_NAME=<リージョン>
AWS_ACCESS_KEY_ID=<アクセスキーID>
AWS_SECRET_ACCESS_KEY=<シークレットアクセスキー>
AWS_SESSION_TOKEN=<セッショントークン>

スクリプト

以下サンプル。呼び出し処理やレスポンスが共通化されている

from litellm import completion

# azure call
response = completion(
    "azure/デプロイしたgpt4系のモデル", messages=[{"content": "日本の初代総理大臣は誰ですか。", "role": "user"}]
)
print(response["choices"][0]["message"]["content"])

# bedrock call
response2 = completion(
    model="bedrock/anthropic.claude-v2",
    messages=[{"content": "日本の初代総理大臣は誰ですか。", "role": "user"}],
)
print(response2["choices"][0]["message"]["content"])

実行結果

日本の初代総理大臣は伊藤博文（いとう ひろぶみ）です。彼は1885年に日本の初代内閣制度が開始されると、そのトップとして初代内閣総理大臣に就任しました。

日本の初代総理大臣は伊藤博文です。

1885年に内閣制度が導入され、翌1886年に第一回内閣が成立した際に、伊藤博文が初代の内閣総理大臣に就任しました。以後、内閣総理大臣は日本の政府の長としての役割を果たすことになります。