OpenAIのAPIをリレーしてトークン数を測定する

はじめに

OpenAIのAPIは非常に便利ですよね。多くの会社で活用されており、開発現場でも広く使用されていると思います。しかし、その中で問題になるのが、誰がどれくらいトークンを使用したかを把握することです。これが見えないため、個別のアカウントを持っている場合は分かりますが、共通の方法でトークン数を管理したいというニーズもあるでしょう。

そこで、今回はOpenAIのAPIをラップしてトークン数を測定する方法を模索しました。単純にAPIをラップして途中でトークン数を測定すれば、この課題を解決できるのではないかと考えましたので、コードを実装してみました。

やりたいこと(図）

一旦はこんなイメージ
今回はチャットの基本機能だけWrapします。Embeddingなどほしい方は、それぞれ欲しい機能のパスを追加してリレーして上げる必要があります。

どのようにユーザーを識別するか

ユーザーの識別はFastAPI側で払い出したAPIキーをもとに測定をする想定です。
※ただし、今回はいろんな認証方法があったりするので、認証はデモでtestというキーをユーザーから受けそれを認証とする事にします。

OpenAIのAPIをWrapする

以下に、OpenAIのWrap元があります。

こちらをそのままリクエストをリレーするようにします。ただし、Streamingなどはちょっとコツが要るので後述します。

まずは以下のパッケージをダウンロードします。

fastapi
uvicorn
python-dotenv

パスと機能を作成

まずは、元と同じように/chat/completionsというパスをPOSTメソッドで受け付けるようにします。

@app.post("/chat/completions")

次にAPI認証を噛ませ、認証OKであればリクエストをそのままリレーし、Requestsで送信します。

async def relay_openai_request(request: Request):
    # クライアントからのリクエストの認証ヘッダーを取得
    client_auth_header = request.headers.get("Authorization")
    if client_auth_header != "Bearer test":
        # 認証ヘッダーが一致しない場合、エラーを返す
        raise HTTPException(status_code=401, detail="Unauthorized")
    
    # JSONデータを取得
    data = await request.json()

    # OpenAIへのリクエストを中継
    response = requests.post(f"{OPENAI_URL}/chat/completions", json=data, headers=headers, stream=stream)

上記で一旦OPENAIのAPIに中継する事ができます。今度はResponseを加工してopenaiのライブラリなどでも使える形で変換します。

    # 通常のJSONレスポンスを返す
    response_data = response.json()
    total_tokens = response_data.get("usage", {}).get("total_tokens")
    if total_tokens is not None:
        print(f"使用したトークン数: {total_tokens}")
    else:
        print("トークン数情報が見つかりません")

    return response_data

こちらで、標準的なレスポンスの実装は完了です。
上記のprintの部分で使用したトークン数などを見ることができます。

StreamがTrueであればストリーミングで返すようにする

stream = data.get("stream", False)

こちらが記載されていればレスポンスをストリーミングで返すようにする必要があります。
ただしこちらの場合は、ストリーミングなので通常レスポンスのように返せませんので、FastAPIのStreamingResponseを使ってストリーミングで返す実装をします。

# ストリーミングオプションが有効の場合、ストリーミングレスポンスを返す
    if stream:
        total_tokens_limit = int(response.headers.get("x-ratelimit-limit-tokens", 0))
        total_tokens_remaining = int(response.headers.get("x-ratelimit-remaining-tokens", 0))
        used_tokens = total_tokens_limit - total_tokens_remaining

        if used_tokens > 0:
            print(f"使用したトークン数: {used_tokens}")
        else:
            print("トークン数情報が見つかりません")
        # print(response.headers)
        def generate():
            for chunk in response.iter_content(chunk_size=8192):
                yield chunk.decode('utf-8')

        # ストリーミングとしてテキストを返す
        return StreamingResponse(generate(), media_type="text/event-stream")

ストリーミングの場合、残念なことに標準レスポンスのときのようなtotal_tokensは拾ってこれませんので、使える合計のトークンと、残りのトークンを減算することで出しています。

ヘッダーx-ratelimit-limit-tokens（合計トークン数）
x-ratelimit-remaining-tokens（残りトークン数）

また、注意点としてはStreamingResponseを返す際は、text/event-streamで返す必要があります。こちらで返さない場合、Pythonのopenaiのライブラリ側でエラーが発生します。

最終的に完成したAPIコード

from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import StreamingResponse
import requests
import os
from dotenv import load_dotenv

# .envファイルから環境変数を読み込む
load_dotenv()

app = FastAPI(
    title="PrivateGPT",
    description="ただのWrapper",
    version="0.0.1"
)

# OpenAIのエンドポイントとAPIキーの定義
OPENAI_URL = "https://api.openai.com/v1"
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

# リクエストのヘッダーを定義
headers = {
    "Authorization": f"Bearer {OPENAI_API_KEY}",
    "Content-Type": "application/json",
    "User-Agent": "OpenAI FastAPI Wrapper"
}

@app.post("/chat/completions")
async def relay_openai_request(request: Request):
    # クライアントからのリクエストの認証ヘッダーを取得
    client_auth_header = request.headers.get("Authorization")
    if client_auth_header != "Bearer test":
        # 認証ヘッダーが一致しない場合、エラーを返す
        raise HTTPException(status_code=401, detail="Unauthorized")
    
    # JSONデータを取得
    data = await request.json()
    # ストリーミングオプションが設定されているか確認
    stream = data.get("stream", False)

    # OpenAIへのリクエストを中継
    response = requests.post(f"{OPENAI_URL}/chat/completions", json=data, headers=headers, stream=stream)

    # レスポンスのステータスコードを確認
    if response.status_code != 200:
        raise HTTPException(status_code=response.status_code, detail=response.text)

    # ストリーミングオプションが有効の場合、ストリーミングレスポンスを返す
    if stream:
        """
        レスポンスヘッダーから使用したトークン数を取得するには、
        ヘッダーx-ratelimit-limit-tokens（合計トークン数）と
        x-ratelimit-remaining-tokens（残りトークン数）を使用して計算することができます。
        """
        total_tokens_limit = int(response.headers.get("x-ratelimit-limit-tokens", 0))
        total_tokens_remaining = int(response.headers.get("x-ratelimit-remaining-tokens", 0))
        used_tokens = total_tokens_limit - total_tokens_remaining

        if used_tokens > 0:
            print(f"使用したトークン数: {used_tokens}")
        else:
            print("トークン数情報が見つかりません")
        # print(response.headers)
        def generate():
            for chunk in response.iter_content(chunk_size=8192):
                yield chunk.decode('utf-8')

        # ストリーミングとしてテキストを返す
        return StreamingResponse(generate(), media_type="text/event-stream")

    # 通常のJSONレスポンスを返す
    response_data = response.json()
    total_tokens = response_data.get("usage", {}).get("total_tokens")
    if total_tokens is not None:
        print(f"使用したトークン数: {total_tokens}")
    else:
        print("トークン数情報が見つかりません")

    return response_data

if __name__ == "__main__":
    # FastAPIのアプリを実行
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8080)

テスト

API側のWrapperが完成したのでテストします。
openaiのライブラリの中にあるChatCompletionクラスを使ってテストします。

標準レスポンステスト

import openai

openai.api_base = "http://127.0.0.1:8080"
openai.api_key = "test"

completion = openai.ChatCompletion.create(
  model="gpt-3.5-turbo",
  messages=[{"role":"system","content":"作文を300文字で作って"}],
  temperature=0.5,
)

response = completion.choices[0].message['content']
print(response)

結果

ストリーミングテスト

import os
import openai
from openai import util
import json

openai.api_base = "http://127.0.0.1:8080"
openai.api_key = "test"


# チャンクごとの応答を生成します
for chunk in openai.ChatCompletion.create(
    model='gpt-3.5-turbo',
    messages=[
        {
            "role": "user",
            "content": "作文を300文字で作って"
        }
    ],
    temperature=0,
    stream=True  
):
    content = chunk["choices"][0].get("delta", {}).get("content")
    if content is not None:
        print(content, end='', flush=True)  # flush=Trueを追加

結果

ん～～なんかトークン数少ないような・・・
ストリーミングだと正確なトークン数が拾ってこれないようです。なのでStreamが来た際はエラーで返してやる必要があるかもしれません。