0
0

【初心者向け】OpenAI APIのトークン数を簡単に記録する2つの方法(サンプルコード付き)

Posted at

OpenAI APIを使用する際、トークン数を記録することは重要です。この記事では、既存のOpenAI通信プログラムに簡単に組み込める2つの方法を紹介します:

  1. APIレスポンスを利用する方法
  2. tiktokenライブラリを使用する方法

前提条件

このコードを使用するには、以下のライブラリが必要です:

  • openai: OpenAI APIとの通信用
  • tiktoken: トークン数の計算用(方法2のみ)
  • sqlite3: データベース操作用

以下のコマンドでインストールできます:

pip install openai tiktoken

1. APIレスポンスを利用する方法

この方法では、OpenAIのAPIレスポンスに含まれるトークン数情報を利用します。

サンプルコード

以下のコードをコピーして、token_logger.pyというファイル名で保存してください。

import sqlite3
from datetime import datetime
from typing import Dict, Any

class TokenLogger:
    def __init__(self, db_name: str = 'token_usage.db'):
        self.db_name = db_name
        self._init_db()

    def _init_db(self):
        with sqlite3.connect(self.db_name) as conn:
            conn.execute('''CREATE TABLE IF NOT EXISTS token_usage
                            (id INTEGER PRIMARY KEY AUTOINCREMENT,
                             timestamp TEXT,
                             prompt_tokens INTEGER,
                             completion_tokens INTEGER,
                             total_tokens INTEGER)''')

    def log_usage(self, response: Dict[str, Any]):
        usage = response['usage']
        with sqlite3.connect(self.db_name) as conn:
            conn.execute(
                "INSERT INTO token_usage (timestamp, prompt_tokens, completion_tokens, total_tokens) VALUES (?, ?, ?, ?)",
                (datetime.now().isoformat(), usage['prompt_tokens'], usage['completion_tokens'], usage['total_tokens'])
            )

logger = TokenLogger()

def log_token_usage(response: Dict[str, Any]):
    logger.log_usage(response)

使用方法

  1. 既存のOpenAI通信プログラムに以下の行を追加します:
from token_logger import log_token_usage

# 既存のOpenAI APIコール
response = client.chat.completions.create(
    model="gpt-3.5-turbo",
    messages=messages
)

# この行を追加
log_token_usage(response)

# 以降は既存のコード

2. tiktokenライブラリを使用する方法

この方法では、OpenAIが提供するtiktokenライブラリを使用してトークン数を計算します。

サンプルコード

以下のコードをtoken_logger.pyに追加してください。

import tiktoken
from typing import List, Dict

class TiktokenLogger(TokenLogger):
    def count_tokens(self, text: str, model: str = "gpt-3.5-turbo") -> int:
        encoding = tiktoken.encoding_for_model(model)
        return len(encoding.encode(text))

    def log_usage(self, messages: List[Dict[str, str]], response_text: str, model: str = "gpt-3.5-turbo"):
        prompt_tokens = sum(self.count_tokens(m["content"], model) for m in messages)
        completion_tokens = self.count_tokens(response_text, model)
        total_tokens = prompt_tokens + completion_tokens

        with sqlite3.connect(self.db_name) as conn:
            conn.execute(
                "INSERT INTO token_usage (timestamp, prompt_tokens, completion_tokens, total_tokens) VALUES (?, ?, ?, ?)",
                (datetime.now().isoformat(), prompt_tokens, completion_tokens, total_tokens)
            )

tiktoken_logger = TiktokenLogger()

def log_tiktoken_usage(messages: List[Dict[str, str]], response_text: str, model: str = "gpt-3.5-turbo"):
    tiktoken_logger.log_usage(messages, response_text, model)

使用方法

  1. 既存のOpenAI通信プログラムに以下の行を追加します:
from token_logger import log_tiktoken_usage

# 既存のOpenAI APIコール
response = client.chat.completions.create(
    model="gpt-3.5-turbo",
    messages=messages
)

# この行を追加
log_tiktoken_usage(messages, response.choices[0].message.content)

# 以降は既存のコード

データの確認方法

記録されたデータは、SQLiteデータベース(token_usage.db)に保存されます。以下のPythonコードでデータを確認できます:

import sqlite3

def view_token_usage():
    with sqlite3.connect('token_usage.db') as conn:
        cursor = conn.execute("SELECT * FROM token_usage")
        for row in cursor:
            print(row)

view_token_usage()

まとめ

これら2つの方法を使うことで、既存のOpenAI通信プログラムにほとんど変更を加えずに、簡単にトークン使用量を記録できます。APIレスポンスを利用する方法は実際の使用量を正確に記録でき、tiktokenを使用する方法は事前にトークン数を推定するのにも役立ちます。

SQLiteを使用することで、データを簡単に保存し、後で分析することができます。目的に応じて適切な方法を選択し、APIの使用状況を効果的に管理してください。

0
0
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
0
0