from hackaprompt.completers import get_completer
import os

openai_api_key = os.getenv("OPENAI_API_KEY")

def test_count_tokens_gpt_35_turbo():

    completer = get_completer(model="gpt-3.5-turbo", openai_api_key=openai_api_key)

    prompt = "hello"
    expected_token_count = 1  # taken on https://platform.openai.com/tokenizer
    assert completer.get_token_count(prompt) == expected_token_count

    prompt = "This is a test to see how many tokens we have"
    expected_token_count = 11  # taken on https://platform.openai.com/tokenizer
    assert completer.get_token_count(prompt) == expected_token_count


def test_count_tokens_flanT5():

    completer = get_completer(model="FlanT5-XXL")

    prompt = "hello"
    expected_token_count = 2
    assert completer.get_token_count(prompt) == expected_token_count

    prompt = "This is a test to see how many tokens we have"
    expected_token_count = 14
    assert completer.get_token_count(prompt) == expected_token_count