token_usage_tracker.py

"""
Token Usage Tracker

This script allows users to retrieve and display the total number of tokens generated by a specific model on a given date. It interacts with the OpenAI API to fetch usage data, including the number of tokens generated and used in context.

The script requires an environment file (.env) containing the OpenAI API key. It uses the 'requests' library to make API requests and 'datetime' to handle dates.

Usage:
    1. Ensure a .env file with the necessary API key is present in the same directory as this script.
    2. Run the script directly. The script is set to query data for the current date but can be modified for other dates.

Author: Hit the Code Labs
Date: 2023-12-08
"""

import requests
import datetime
import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Retrieve API key from environment variable
api_key = os.environ.get("OPENAI_API_KEY")

# Headers for API request
headers = {'Authorization': f'Bearer {api_key}'}

# API endpoint for retrieving usage data
url = 'https://api.openai.com/v1/usage'

# Set the date for usage data retrieval
date = datetime.date(2023, 12, 8)

# Parameters for the API request, formatted as 'YYYY-MM-DD'
params = {'date': date.strftime('%Y-%m-%d')}

# Sending the API request and storing the response
response = requests.get(url, headers=headers, params=params)
usage_data = response.json()

# Process and display the usage data if available
if "data" in usage_data:
    usage_data2 = usage_data['data']
    M = {}  # Dictionary to store total tokens used by each model

    # Iterate through the data to calculate total number of tokens for each model
    for data in usage_data2:
        model_name = data['snapshot_id']
        n_generated_tokens_total = data['n_generated_tokens_total']
        n_context_tokens_total = data['n_context_tokens_total']
        total_tokens = n_generated_tokens_total + n_context_tokens_total

        if model_name not in M:
            M[model_name] = total_tokens
        else:
            M[model_name] += total_tokens

    print(M)  # Output the total tokens used for each model

# Sample output: {'gpt-4-1106-preview': 480468}