Display cost usage based on converstations from chatgpt

calc.py

import json
from rich.console import Console
from rich.table import Table

with open("conversations.json", 'r') as f:
    j_data = json.loads(f.read())

prompts = []
completion = []

for element in j_data:
    mapping = element['mapping']
    for key, value in mapping.items():
        if value['message'] is None:
            continue

        author = value['message']['author']['role']
        parts = value['message']['content']['parts'] if value['message'] is not None else None

        if parts and author == "user":
            prompts.extend(parts)
        elif parts and author == "system" or author == "assistant":
            completion.extend(parts)


def get_total_count(messages):
    total_count = 0
    for message in messages:
        for i in range(0, len(message), 4): # A tokeen is 4 characters long
            total_count += 1
    return total_count

console = Console()

total_prompts = get_total_count(prompts) 
total_completion = get_total_count(completion)
total_prompts_K = total_prompts / 1000 # / 1000 to get in K
total_completion_K = total_completion / 1000 # / 1000 to get in K

table = Table(show_header=True, header_style="bold magenta", title="Data")

table.add_column("GPT Messages", justify="right", style="cyan", no_wrap=True)
table.add_column("User Messages", justify="right", style="cyan", no_wrap=True)
table.add_column("Prompt Tokens", justify="right", style="cyan", no_wrap=True)
table.add_column("Completion Tokens", justify="right", style="cyan", no_wrap=True)
table.add_column("Prompt Tokens (K)", justify="right", style="cyan", no_wrap=True)
table.add_column("Completion Tokens (K)", justify="right", style="cyan", no_wrap=True)
table.add_row(f"{len(prompts)}", f"{len(completion)}", f"{total_prompts}", f"{total_completion}", f"{total_prompts_K}", f"{total_completion_K}")

console.print(table)

table = Table(show_header=True, header_style="bold magenta", title="GPT4")

table.add_column("Model", justify="right", style="cyan", no_wrap=True)
table.add_column("Prompt", justify="right", style="cyan", no_wrap=True)
table.add_column("Completion", justify="right", style="cyan", no_wrap=True)
table.add_column("Total", justify="right", style="cyan", no_wrap=True)

prompt_cost = total_prompts_K * 0.03
completion_cost = total_completion_K * 0.06
str_prompts = "{:.2f}".format(prompt_cost)
str_completion = "{:.2f}".format(completion_cost)
str_total = "{:.2f}".format(prompt_cost + completion_cost)
table.add_row("8K context", f"${str_prompts}", f"${str_completion}", f"${str_total}")

prompt_cost = total_prompts_K * 0.06
completion_cost = total_completion_K * 0.12
str_prompts = "{:.2f}".format(prompt_cost)
str_completion = "{:.2f}".format(completion_cost)
str_total = "{:.2f}".format(prompt_cost + completion_cost)
table.add_row("32K context", f"${str_prompts}", f"${str_completion}", f"${str_total}")

console.print(table)

table = Table(show_header=True, header_style="bold magenta", title="Chat")
table.add_column("Model", justify="right", style="cyan", no_wrap=True)
table.add_column("Usage", justify="right", style="cyan", no_wrap=True)
table.add_column("Total", justify="right", style="cyan", no_wrap=True)

usage_cost = (total_prompts_K + total_completion_K) * 0.002
str_total = "{:.2f}".format(usage_cost)
table.add_row("gpt-3.5-turbo", f"${str_total}", f"${str_total}")

console.print(table)