スポンサーリンク

【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)

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

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

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

スポンサーリンク

実行結果

試しに実行してみます。

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

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

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

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

スポンサーリンク
スポンサーリンク
スポンサーリンク
AzureGenerative AIPython
著者SNS
タイトルとURLをコピーしました