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

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`というファイル名で保存してください。

```python
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通信プログラムに以下の行を追加します：

```python
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`に追加してください。

```python
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通信プログラムに以下の行を追加します：

```python
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コードでデータを確認できます：

```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の使用状況を効果的に管理してください。