-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathmain.py
133 lines (82 loc) · 4.26 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
from parse_convo import process_conversations
import tiktoken
from utils import handle_args, diplay_cost_all_time_frames, plot_costs_for_all_time_frames, diplay_cost_for_time_frame
def count_conversation_tokens(encoding, conversations):
"""
Iterate over the conversations
determine the total number of tokens sent by the user and chatgpt in the last 30 days
"""
user_tokens = 0
chatgpt_tokens = 0
for conversation in conversations:
for message in conversation['messages']:
if message['author'] == 'user':
user_tokens += len(encoding.encode(message['text']))
elif message['author'] == 'ChatGPT':
chatgpt_tokens += len(encoding.encode(message['text']))
return user_tokens, chatgpt_tokens
def calculate_cost(user_tokens, chatgpt_tokens, user_token_price=0.00500, chatgpt_token_price=0.01500):
"""
Calculate the cost of user and chatgpt tokens in dollars.
Args:
user_tokens (int): The number of tokens sent by the user.
chatgpt_tokens (int): The number of tokens sent by ChatGPT.
user_token_price (float, optional): The price of 1k user tokens. Defaults to 0.00500.
chatgpt_token_price (float, optional): The price of 1k ChatGPT tokens. Defaults to 0.01500.
Returns:
tuple: A tuple containing the total cost of user tokens, chatgpt tokens, and total cost.
"""
user_cost = (user_tokens / 1000) * user_token_price
chatgpt_cost = (chatgpt_tokens / 1000) * chatgpt_token_price
return user_cost, chatgpt_cost, user_cost + chatgpt_cost
def calculate_month(encoding, conversation_summary, time_frame):
"""
Calculate the cost of user and chatgpt tokens in dollars for a specific month.
Args:
encoding (tiktoken.Encoding): The encoding of the model.
conversation_summary (dict): The conversation summary.
time_frame (str): The time frame to calculate the cost for.
Returns:
tuple: A tuple containing the total cost of user tokens, chatgpt tokens, and total cost.
"""
user_tokens, chatgpt_tokens = count_conversation_tokens(encoding, conversation_summary[time_frame])
user_cost, chatgpt_cost, total_cost = calculate_cost(user_tokens, chatgpt_tokens)
return user_cost, chatgpt_cost, total_cost
def calculate_cost_for_all_time_frames(encoding, conversation_summary):
"""Calculate the cost for all time frames."""
final_total_cost = 0
final_user_cost = 0
final_chatgpt_cost = 0
cost_data = {}
for time_frame in conversation_summary.keys():
print(f"--- {time_frame.upper()} ---")
user_cost, chatgpt_cost, total_cost = calculate_month(encoding, conversation_summary, time_frame)
cost_data[time_frame] = (user_cost, chatgpt_cost, total_cost)
final_total_cost += total_cost
final_user_cost += user_cost
final_chatgpt_cost += chatgpt_cost
print(f"User Tokens Cost: ${user_cost:.4f}")
print(f"ChatGPT Tokens Cost: ${chatgpt_cost:.4f}")
print(f"Total Tokens Cost: ${total_cost:.4f}")
print("-" * 40)
# Visualize costs for all time frames
plot_costs_for_all_time_frames(cost_data, conversation_summary.keys())
return final_total_cost, final_user_cost, final_chatgpt_cost
if __name__ == "__main__":
# Should create save/load mechanism for conversation_summary.
# right now, we process every time
args = handle_args()
output = process_conversations(args)
conversation_summary = output if args.return_only else output[1]
if not args.return_only:
for info in output[0]:
print(f"Created {info['file']} in directory {info['directory']}")
encoding = tiktoken.encoding_for_model(args.model_name)
if args.time_frame == 'all':
final_total_cost, final_user_cost, final_chatgpt_cost = calculate_cost_for_all_time_frames(encoding, conversation_summary)
diplay_cost_all_time_frames(final_total_cost, final_user_cost, final_chatgpt_cost, conversation_summary)
else:
time_frame = args.time_frame
user_cost, chatgpt_cost, total_cost = calculate_month(encoding, conversation_summary, time_frame)
diplay_cost_for_time_frame(
user_cost, chatgpt_cost, total_cost, time_frame)