【Azure OpenAI Service】asyncioでAPI呼び出し並行処理

Azure OpenAI Serviceを使っていますが、特にGPT-4では応答に時間がかかります。

そこで非同期でAPIを呼び出し、並行でcompletionを実行することで、全体の処理時間短縮を図りました。

コード

必要なライブラリをインストールします。

pip install openai tenacity

openai
AzureOpenAIのAPIにアクセスするために必要です。
tenacity
リトライロジックを実装するために使用します。非同期処理のリトライにも対応しています。

今回動作確認を行ったバージョンは以下のとおりです。

openai: 1.13.3
tenacity: 8.2.3

プロンプトのリストをもとに、複数のプロンプトでのAPI実行を並列で行うサンプルコードです。

import sys
import os
import asyncio
from openai import AsyncAzureOpenAI
from tenacity import retry, wait_fixed, stop_after_attempt

azure_deployment = os.getenv("AZURE_DEPLOYMENT")
azure_endpoint = os.getenv("AZURE_OPENAI_ENDPOINT")
api_key = os.getenv("OPENAI_API_KEY")
api_version = os.getenv("OPENAI_API_VERSION")


@retry(wait=wait_fixed(3), stop=stop_after_attempt(10))
async def async_aoai_get_response(
        system_prompt: str,
        user_prompt: str,
        temperature: float = 0.0
):
    client = AsyncAzureOpenAI(
        azure_endpoint=azure_endpoint,
        api_key=api_key,
        api_version=api_version,
    )

    res = await client.chat.completions.create(
        model=azure_deployment,
        temperature=temperature,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt}
        ]
    )

    return res.choices[0].message.content


async def main(
        system_prompt_list: list,
        user_prompt_list: list
):
    tasks = []
    for system_prompt, user_prompt in zip(system_prompt_list, user_prompt_list):
        task = async_aoai_get_response(system_prompt, user_prompt)
        tasks.append(task)
    responses = await asyncio.gather(*tasks)
    return responses


if __name__ == '__main__':
    system_prompt_list = ["System prompts here"]  # ブランクでもOK
    user_prompt_list = ["Your prompts here"]

    # Windowsの場合のみSelectorEventLoopを使用
    if sys.platform == 'win32':
        asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())

    result = asyncio.run(main(system_prompt_list, user_prompt_list))
    print(result)

処理の流れは以下のとおりです。

環境変数の読み込み
環境変数からAzure OpenAIサービスの設定（デプロイメント、APIキー、APIバージョン、エンドポイント）を読み込みます。
AsyncAzureOpenAIクライアントの初期化
読み込んだ設定を使用して、非同期通信に対応するAsyncAzureOpenAIクライアントインスタンスを初期化します。
非同期関数の定義
- async_aoai_get_response: 非同期関数で、指定されたシステムプロンプトとユーザープロンプトを使用してAzure OpenAIサービスにリクエストを送信し、AIによる応答を取得します。この関数はretryデコレータを使用しており、リクエストが失敗した場合は最大10回まで3秒おきにリトライします。
- main: 与えられたプロンプトリストに対してasync_aoai_get_response関数を非同期に実行するタスクを生成し、asyncio.gatherを使用してこれらのタスクを並列に実行します。この関数は、全てのプロンプトに対するAI応答をリストとして返します。
メイン処理の実行
スクリプトの実行時には、main関数が呼び出され、事前に定義されたプロンプトリストとシステムプロンプトリストを引数として渡します。その後、非同期にAPIリクエストが行われ、結果が集約されて出力されます。

同時に実行しすぎると、リトライが多発しますのでご注意ください。

実行結果

試しに実行してみます。

以下のとおりプロンプトのリストを作成し、実行します。

    system_prompt_list = ["あなたはデータアナリストです。", "あなたはデータアナリストです。"]
    user_prompt_list = ["相関係数について簡単に教えて。", "重回帰分析について簡単に教えて。"]

実行結果は以下のとおりです。

[
‘相関係数は、2つの変数間の関係の強さと方向を測定する統計的な指標です。相関係数の値は-1から1までの範囲で、-1は完全な負の相関（一方の変数が増加すると、もう一方の変数が減少する）、0は相関がない（2つの変数間に関連性がない）、1は完全な正の相関（一方の変数が増加すると、もう一方の変数も増加する）を示します。\n\n相関係数は、変数間の関係を理解するための重要なツールであり、予測モデルの作成や変数の選択にも使用されます。ただし、相関係数が因果関係を示すわけではないことに注意が必要です。つまり、ある変数がもう一方の変数の変動を引き起こすとは限らないということです。’,
‘重回帰分析は、統計学における予測モデルの一つで、複数の説明変数から目的変数を予測するための手法です。これは、一つの説明変数だけでなく、複数の説明変数を使用して目的変数の変動を説明しようとするものです。\n\n例えば、家の価格を予測する場合、家の大きさだけでなく、立地や築年数なども考慮に入れることができます。これらの複数の要素を同時に考慮することで、より正確な予測が可能になります。\n\n重回帰分析の数学的な表現は以下のようになります。\n\nY = a + b1*X1 + b2*X2 + … + bn*Xn + e\n\nここで、Yは目的変数、X1からXnは説明変数、aは切片、b1からbnは各説明変数の係数（影響力）、eは誤差項を表します。\n\nこのモデルを用いて、各説明変数の係数を推定し、それらが目的変数にどの程度影響を与えるかを解析します。また、説明変数間の相互作用や、非線形の関係も考慮することができます。’
]

複数のAPI応答を取得できています。