-
Notifications
You must be signed in to change notification settings - Fork 3
/
openai-chat-endpoint-boilerplate.py
38 lines (33 loc) · 1.27 KB
/
openai-chat-endpoint-boilerplate.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
import os
import openai
import tiktoken
from dotenv import load_dotenv
load_dotenv()
openai.api_key = os.getenv('OPENAI_API_KEY')
enc = tiktoken.get_encoding("cl100k_base")
messages = []
def get_prompt(input):
context = []
messages.append(input)
for index, message in enumerate(messages):
if index % 2 == 0:
context.append({"role": "user", "content": message})
else:
context.append({"role": "assistant", "content": message})
return context
def create_chat_completion(input):
completion = openai.ChatCompletion.create(
model="gpt-3.5-turbo", # you can also go with gpt-4 here (cost will be different! https://openai.com/pricing )
messages=get_prompt(input),
)
messages.append(completion.choices[0].message.content)
return completion.choices[0].message.content
def count_used_tokens():
token_count = len(enc.encode(" ".join(messages)))
token_cost = token_count / 1000 * 0.002 # gpt-3.5-turbo cost
return "🟡 Used tokens this round: " + str(token_count) + " (" + format(token_cost, '.5f') + " USD)"
while True:
user_input = input("You: ")
completion = create_chat_completion(user_input)
print("Bot:", completion)
print(count_used_tokens())