😒

[OpenAI API] ストリーミングレスポンスの構造を理解してトークン数を取得する

2024/05/27に公開

はじめに

OpenAI APIで出力をストリーミングさせたときのtoken数取得についての記事を書きます。

OpenAI APIは従量課金なのでリクエストごとに使用するtoken数が気になりますが、ストリーミング時は長らくtoken数が表示されない状態となっていました。

しかし2024年5月、ようやくストリーミングでも使用token数が表示されるように機能追加されました。

https://cookbook.openai.com/examples/how_to_stream_completions#4-how-to-get-token-usage-data-for-streamed-chat-completion-response

これで一応出力させることができるようになりましたが、ストリーミング時はレスポンスがチャンクごとに返されることもあり、チャンクの構造を理解していないといまだにちょっと複雑なのでまとめておきます。

ストリーミングではないときのtoken数取得

おさらいまでに、ストリーミングではないときのtoken数取得についても書いておきます。
と言ってもストリーミングではないときは何もしなくてもtoken数が出力されます。

以下のようなコードで「こんにちは」とリクエストしてみます。今回はレスポンスの全体を見たいのでresponseをjson型に変換して出力しています。

from openai import OpenAI
import json

client = OpenAI()

response = client.chat.completions.create(
    model="gpt-4o",
    messages=[
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "こんにちは"},
    ],
)
print(json.dumps(response.to_dict(), ensure_ascii=False, indent=4))

すると以下のようなレスポンスが返ってきます。

{
    "id": "chatcmpl-9T8B9bYY41N1kMW4gfQLF425fil7B",
    "choices": [
        {
            "finish_reason": "stop",
            "index": 0,
            "logprobs": null,
            "message": {
                "content": "こんにちは!今日はどんなお手伝いをしましょうか?",
                "role": "assistant",
                "function_call": null,
                "tool_calls": null
            }
        }
    ],
    "created": 1716729979,
    "model": "gpt-4o-2024-05-13",
    "object": "chat.completion",
    "system_fingerprint": "fp_43dfabdef1",
    "usage": {
        "completion_tokens": 14,
        "prompt_tokens": 18,
        "total_tokens": 32
    }
}

response.choices[0].message.contentに生成されたテキストが入っています。

そして、response.usageの部分にtoken数が出ています。

prompt_tokensが入力で使ったtoken、completion_tokensが出力、total_tokensが合計です。使ったモデルとこの値があれば、ここからコストを算出できます。

ストリーミングのときのtoken数取得

リクエストについて

ここからがストリーミングのときです。
ストリーミングのときは以下のようなコードでリクエストを送ります。今回もレスポンスをjson型に変換して出力しています。

from openai import OpenAI
import json

client = OpenAI()

stream = client.chat.completions.create(
    model="gpt-4o",
    messages=[{"role": "user", "content": "東京の魅力を50文字で教えて"}],
    stream=True,
    stream_options={"include_usage": True},
)
end_flag = False
full_txt=""
for chunk in stream:
    print(json.dumps(chunk.to_dict(), ensure_ascii=False, indent=4))

ストリーミングするときは、stream=True を設定すればよいです。

そして、レスポンスにtoken数を含めるにはさらにstream_options={"include_usage": True}を含めればよいです。これが2024年5月の機能追加です。

個人的にはストリーミングのレスポンスにtoken数が含まれないのはバグでは?くらいに思っていたので、オプションなしで常に出る形がよかったのになぁと思います。

ただ、実装のされ方をみるに何か難しい事情があったかなぁという印象です。

しかし、何はともあれ出力されるようになったのでこれは嬉しい変更です。

そして、「なんだ簡単じゃん」と思われるかもしれませんが、実は私はこの後これをprintするのに少し詰まりました。

というのも、ストリーミング時はレスポンスがチャンクごとに返ってきますが、チャンクごとに特定のスキーマがnullだったり存在しなかったりと構造に差分があります。
これを理解しないで無邪気にprintすると、out of indexなどつまらないエラーを引いてしまうのでレスポンスをまとめておきたいと思ったのがこの記事のきっかけです。

レスポンスについて

前提としてストリーミングレスポンスはストリーミングではないときとスキーマが異なります。そして、そのスキーマがチャンクごとに連続して返ってきます。

ストリーミング時のレスポンス(チャンク)は大きく4つのフェーズに分かれているようです。

テキスト生成スタート

まずは最初のチャンクが返ってきます。ここには生成されたテキストは含まれていません。
ちなみにこの後のチャンクではchoices[0].delta.content に生成されたテキストが入ってきますが、ここでは何も入っていません。

{
    "id": "chatcmpl-9TUEaNpZcyR6n2WuOqFdlQuQzo9Zs",
    "choices": [
        {
            "delta": {
                "content": "",
                "role": "assistant"
            },
            "finish_reason": null,
            "index": 0,
            "logprobs": null
        }
    ],
    "created": 1716814760,
    "model": "gpt-4o-2024-05-13",
    "object": "chat.completion.chunk",
    "system_fingerprint": "fp_3196d36131",
    "usage": null
}

生成テキスト出力

生成されたテキストがチャンクごとに返ってきます。choices[0].delta.contentに最初のチャンクである'東京'が入っています。

{
    "id": "chatcmpl-9TUEaNpZcyR6n2WuOqFdlQuQzo9Zs",
    "choices": [
        {
            "delta": {
                "content": "東京"
            },
            "finish_reason": null,
            "index": 0,
            "logprobs": null
        }
    ],
    "created": 1716814760,
    "model": "gpt-4o-2024-05-13",
    "object": "chat.completion.chunk",
    "system_fingerprint": "fp_3196d36131",
    "usage": null
}

ここから生成されたテキストがこの形式で連続して返ってきます。

テキスト生成フィニッシュ

生成されたテキストがひととおり返ってくると、以下のレスポンスが返ってきます。

{
    "id": "chatcmpl-9TUEaNpZcyR6n2WuOqFdlQuQzo9Zs",
    "choices": [
        {
            "delta": {},
            "finish_reason": "stop",
            "index": 0,
            "logprobs": null
        }
    ],
    "created": 1716814760,
    "model": "gpt-4o-2024-05-13",
    "object": "chat.completion.chunk",
    "system_fingerprint": "fp_3196d36131",
    "usage": null
}

今まで生成されたテキストの入っていたchoices[0].delta.contentはdeltaのレイヤーでnullになっています。

代わりにfinish_reasonstopになっています。これが生成が終了したことを示しています。

token出力

最後にtoken数が返ってきます。ここではchoices[]になり、usageにトークン情報が入っています。

{
    "id": "chatcmpl-9TUEaNpZcyR6n2WuOqFdlQuQzo9Zs",
    "choices": [],
    "created": 1716814760,
    "model": "gpt-4o-2024-05-13",
    "object": "chat.completion.chunk",
    "system_fingerprint": "fp_3196d36131",
    "usage": {
        "completion_tokens": 40,
        "prompt_tokens": 17,
        "total_tokens": 57
    }
}

ストリーミングで生成されたテキストを出力するときはprint(chunk.choices[0].delta.content)と書くと思いますが、stream_options={"include_usage": True}を設定するとここでchoicesnullなので[0]が見つからずエラーになります。

これがちょっと癖のある仕様だなと思います。

ちなみに今回は以下のようなテキストが生成されました。
東京は伝統と近未来が融合する都市。豊富な食文化、多彩なエンタメ、歴史的建造物と現代建築が魅力です。

コード

最後に生成されたテキストとtoken数を出力するコードを書きました。
基本的にはchunk.choices[0].delta.contentを出力し、finish_reasonstopになったらchunk.usageのトークン情報を出力して終了します。

from openai import OpenAI
import json

client = OpenAI()

stream = client.chat.completions.create(
    model="gpt-4o",
    messages=[{"role": "user", "content": "東京の魅力を100文字で教えて"}],
    stream=True,
    stream_options={"include_usage": True},
)

end_flag = False

for chunk in stream:
    if end_flag:
        print(chunk.usage.to_dict())
    else:
        if chunk.choices and chunk.choices[0].delta.content is not None:
            print(chunk.choices[0].delta.content, end="")
        elif chunk.choices[0].finish_reason == "stop":
            end_flag = True
            print("")

以下のように出力されます。
標準出力サンプル

おわりに

簡単に出力できるかと思いきや、レスポンスの構造を理解していなくてちょっとはまりました。
みなさんお気を付けください。

関連情報

Discussion