Get OpenAI Usage / Spend information programmatically, broken down by API key, model, input and output

openai_usage_tool.py

import argparse
from datetime import date, timedelta
from typing import TypedDict

import requests
from tabulate import tabulate


class OpenAiUsageTool:
    '''
    Get OpenAI Usage / Spend information broken down by API key, model, input, output, and total cost.
    '''
    ACTIVITY_URL = 'https://api.openai.com/v1/dashboard/activity'

    MODEL_TO_COST_PER_M = {
        # Completion models.
        'gpt-4-0125-preview': (10, 30),
        'gpt-3.5-turbo-0125': (0.50, 1.50),
        'gpt-3.5-turbo-1106': (0.50, 1.50),
        'gpt-3.5-turbo-0613': (0.50, 1.50),

        # Embedding models.
        # We set the output cost for embeddings to a large number to catch any bugs in the tool.
        'text-embedding-3-small': (0.02, 1000),
        'text-embedding-3-large': (0.13, 1000),
        'text-embedding-ada-002-v2': (0.10, 1000),
    }

    def get_openai_spend(self, auth_key: str | None, time_period: str):
        '''
        Get the OpenAI spend per API key.
        :param api_key: Your OpenAI API key
        :return: None, prints the spend per API key
        '''
        if not auth_key:
            auth_key = input('Go to https://platform.openai.com/usage and paste Authorization key from the "activity" request: ')
            print(f'Authorization key: "{auth_key}"')
        if not auth_key.startswith('sess-'):
            print('Authorization key must start with "sess-"')
            return

        today = date.today()
        start_date = today
        if time_period == 'week':
            start_date = today - timedelta(days=today.weekday())
        elif time_period == 'month':
            start_date = today.replace(day=1)
        end_date = today + timedelta(days=1)
        print(f'Getting usage from {start_date.isoformat()} to {end_date.isoformat()}')

        response = requests.get(
            f'{self.ACTIVITY_URL}?start_date={start_date.isoformat()}&end_date={end_date.isoformat()}',
            headers={'Authorization': f'Bearer {auth_key}'},
        ).json()

        key_name_model_dict: dict[tuple[str, str], tuple[int, int, int]] = {}
        for object in response['data']:
            key_name = object['api_key_name']
            model = object['snapshot_id']
            num_context_tokens = object['n_context_tokens_total']
            num_generated_tokens = object['n_generated_tokens_total']
            context_cost = self.MODEL_TO_COST_PER_M[model][0] * num_context_tokens / 1e6
            generated_cost = self.MODEL_TO_COST_PER_M[model][1] * num_generated_tokens / 1e6
            total_cost = context_cost + generated_cost
            if (key_name, model) not in key_name_model_dict:
                key_name_model_dict[(key_name, model)] = (0, 0, 0)
            prev_context_cost, prev_generated_cost, prev_total_cost = key_name_model_dict[(key_name, model)]
            key_name_model_dict[(key_name, model)] = (prev_context_cost + context_cost, prev_generated_cost + generated_cost, prev_total_cost + total_cost)

        class KeyNameModel(TypedDict):
            key: str
            model: str
            context_cost: int
            generated_cost: int
            total_cost: int

        key_name_models: list[KeyNameModel] = []
        organization_context_cost = 0
        organization_generated_cost = 0
        organization_total_cost = 0
        for (key_name, model), (context_cost, generated_cost, total_cost) in key_name_model_dict.items():
            organization_context_cost += context_cost
            organization_generated_cost += generated_cost
            organization_total_cost += total_cost
            key_name_models.append({
                'key': key_name,
                'model': model,
                'context_cost': context_cost,
                'generated_cost': generated_cost,
                'total_cost': total_cost,
            })

        key_name_models.sort(key=lambda x: x['total_cost'], reverse=True)

        class Data(TypedDict):
            key: list[str]
            model: list[str]
            context_cost: list[int]
            generated_cost: list[int]
            total_cost: list[int]
        data: Data = {
            'key': ['total_organization'],
            'model': ['all_models'],
            'context_cost': [organization_context_cost],
            'generated_cost': [organization_generated_cost],
            'total_cost': [organization_total_cost],
        }
        for key_name_model in key_name_models:
            data['key'].append(key_name_model['key'])
            data['model'].append(key_name_model['model'])
            data['context_cost'].append(key_name_model['context_cost'])
            data['generated_cost'].append(key_name_model['generated_cost'])
            data['total_cost'].append(key_name_model['total_cost'])
        print(tabulate(data, headers=list(data.keys()), tablefmt='plain', maxcolwidths=100))


if __name__ == '__main__':
    parser = argparse.ArgumentParser(prog='openai_usage_tool', allow_abbrev=False)
    parser.add_argument('--auth_key', help='API authorization session key')
    parser.add_argument('--time_period', choices=['week', 'month'], default='month', help='The time period to get usage for')
    args = parser.parse_args()

    OpenAiUsageTool().get_openai_spend(args.auth_key, args.time_period)