build_index/pricing.py

# -*-coding:utf-8 -*-
import tiktoken


def num_tokens_from_string(string: str, encoding_name: str):
    # Function to convert string to tokens and estimate user cost.
    encoding = tiktoken.get_encoding(encoding_name)
    num_tokens = len(encoding.encode(string))
    total_price = ((num_tokens / 1000) * 0.0004)
    return num_tokens, total_price


def check_price(docs):
    docs_content = ""
    for doc in docs:
        docs_content += doc.page_content

    tokens, total_price = num_tokens_from_string(string=docs_content, encoding_name="cl100k_base")

    print(f"Number of Tokens = {format(tokens, ',d')}")
    print(f"Approx Cost = ${format(total_price, ',.2f')}")
    user_input = input("Price Okay? (Y/N) \n").upper() == 'Y'
    return user_input