OpenAI APIを使用する際、トークン数を記録することは重要です。この記事では、既存のOpenAI通信プログラムに簡単に組み込める2つの方法を紹介します:
- APIレスポンスを利用する方法
- 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)
使用方法
- 既存の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)
使用方法
- 既存の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の使用状況を効果的に管理してください。