LambdaでClaudeのトークン数計算

備忘録

AntoropicのトークナイザーがBeta版しかないため、tiktokenを使用

Lambdaレイヤーを作成して、関数に紐づけ

• Zipファイル作成

mkdir python && cd $_
pip install tiktoken -t .
zip -r tiktoken.zip python

• Lambdaレイヤー作成

• Lambda関数に紐づけ

テスト用Lambda関数

import tiktoken

def lambda_handler(event, context):

    def calculate_token_length_claude(text, model="gpt-3.5-turbo"):
        # モデルに対応するトークナイザーを取得
        encoding = tiktoken.encoding_for_model(model)
        # テキストをエンコードし、トークン数を返す
        tokens = encoding.encode(text)
        return len(tokens)

    # テストデータ
    system_prompt = "This is the system prompt."
    user_content = "This is the user's content."

    # トークン数計算
    system_prompt_tokens = calculate_token_length_claude(system_prompt)
    user_content_tokens = calculate_token_length_claude(user_content)
    total_input_tokens = system_prompt_tokens + user_content_tokens

    # 結果を返す
    return {
        "system_prompt_tokens": system_prompt_tokens,
        "user_content_tokens": user_content_tokens,
        "total_input_tokens": total_input_tokens
    }