openai_test_cost_usage.py

import openai
from util import get_message
openai.api_type = "azure"
openai.api_base = "https://hkust.azure-api.net"
openai.api_key = "6c25a1df10384b90afedc13130074359"
openai.api_version = "2023-05-15"
from tqdm import tqdm
from tenacity import retry,stop_after_attempt,stop_after_delay,wait_fixed

gpt_answer_path = r'./gpt_answer.txt'
sample_session_path = r'./data/uk_sample_1000.txt'
sample_item_path= r'./data/uk_item_sample_1000.txt'


pricing = {
    "prompt_tokens": 0.0015,
    "completion_tokens": 0.002,
}

# retry 5 times, interval 10s, waiting 2s 
# @retry(stop=(stop_after_delay(10) | stop_after_attempt(5)), wait=wait_fixed(2))
def get_message_from_api(message):
    response = openai.ChatCompletion.create(
            engine = "gpt-35-turbo",
            temperature = 0,
            messages = [{"role":"user","content":message}])
    a = response.get("choices")[0]["message"]["content"]

    usage = response.get("usage")
    completion_tokens = usage["completion_tokens"]
    prompt_tokens = usage["prompt_tokens"]

    return a, prompt_tokens, completion_tokens

def get_gpt_answer():
    f = open(gpt_answer_path,'w',encoding='utf-8')
    idx = 0
    for sessions, message in tqdm(get_message(sample_session_path,sample_item_path)):
        try:
            a, prompt_tokens, completion_tokens = get_message_from_api(message=message)


        except Exception as e:
            print(f'session {idx} have expection:')
            print(e)
            a = 'ERROR:'+str(idx)
            prompt_tokens = 0
            completion_tokens = 0

        f.write('Session: '+str(sessions)+'\n')
        f.write('Message: '+str(message)+'\n')
        f.write(a+'\n'+'-------------------------'+'\n')

        f.write('prompt_tokens: '+str(prompt_tokens) + "\n")
        f.write('completion_tokens: '+str(completion_tokens) + "\n")
        f.write("cost: "+str(pricing["prompt_tokens"]*prompt_tokens+pricing["completion_tokens"]*completion_tokens)+"\n")

        idx+=1

        if idx%100==0:
            f.flush()
        
        if idx > 10:
            break


if __name__ == "__main__":
    
    get_gpt_answer()