Update Anthropic token counting

[the-praxs]

Closes #83.

[areibman]

Changes make sense, but we're unfortunately failing a bunch of tests:

============================= test session starts ==============================
platform darwin -- Python 3.11.5, pytest-8.3.3, pluggy-1.5.0
rootdir: /Users/reibs/Projects/tokencost
configfile: pyproject.toml
plugins: baserun-0.9.16, anyio-4.6.2.post1, requests-mock-1.12.1
collected 98 items

tests/test_costs.py ................F................F.................. [ 53%]
..................F.........................                             [ 97%]
tests/test_llama_index_callbacks.py ..                                   [100%]

=================================== FAILURES ===================================
___________________ test_count_message_tokens[claude-2.1-4] ____________________

model = 'claude-2.1', expected_output = 4

    @pytest.mark.parametrize(
        "model,expected_output",
        [
            ("gpt-3.5-turbo", 15),
            ("gpt-3.5-turbo-0301", 17),
            ("gpt-3.5-turbo-0613", 15),
            ("gpt-3.5-turbo-16k", 15),
            ("gpt-3.5-turbo-16k-0613", 15),
            ("gpt-3.5-turbo-1106", 15),
            ("gpt-3.5-turbo-instruct", 15),
            ("gpt-4", 15),
            ("gpt-4-0314", 15),
            ("gpt-4-0613", 15),
            ("gpt-4-32k", 15),
            ("gpt-4-32k-0314", 15),
            ("gpt-4-1106-preview", 15),
            ("gpt-4-vision-preview", 15),
            ("gpt-4o", 15),
            ("azure/gpt-4o", 15),
            ("claude-2.1", 4),
        ],
    )
    def test_count_message_tokens(model, expected_output):
        print(model)
>       assert count_message_tokens(MESSAGES, model) == expected_output

tests/test_costs.py:54:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
tokencost/costs.py:68: in count_message_tokens
    raise e
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

messages = [{'content': 'Hello', 'role': 'user'}, {'content': 'Hi there!', 'role': 'assistant'}]
model = 'claude-2.1'

    def count_message_tokens(messages: List[Dict[str, str]], model: str) -> int:
        """
        Return the total number of tokens in a prompt's messages.
        Args:
            messages (List[Dict[str, str]]): Message format for prompt requests. e.g.:
                [{ "role": "user", "content": "Hello world"},
                 { "role": "assistant", "content": "How may I assist you today?"}]
            model (str): Name of LLM to choose encoding for.
        Returns:
            Total number of tokens in message.
        """
        model = model.lower()
        model = strip_ft_model_name(model)

        # Anthropic token counting requires a valid API key
        if "claude-" in model:
            logger.warning(
                "Warning: Anthropic token counting API is currently in beta. Please expect differences in costs!"
            )
            if "claude-3-sonnet" in model:
                logger.warning(
                    f"Token counting (beta) is not supported for {model}. Returning num tokens using count from the string."
                )
                # For anthropic<0.39.0 this method is no more supported
                prompt = "".join(message["content"] for message in messages)
                return count_string_tokens(prompt, model)

            ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")

            try:
                client = anthropic.Client(api_key=ANTHROPIC_API_KEY)
>               num_tokens = client.beta.messages.count_tokens(
                    model=model,
                    messages=messages,
                ).input_tokens
E               AttributeError: 'Beta' object has no attribute 'messages'

tokencost/costs.py:60: AttributeError
----------------------------- Captured stdout call -----------------------------
claude-2.1
------------------------------ Captured log call -------------------------------
WARNING  tokencost.costs:costs.py:45 Warning: Anthropic token counting API is currently in beta. Please expect differences in costs!
______________ test_count_message_tokens_with_name[claude-2.1-4] _______________

model = 'claude-2.1', expected_output = 4

    @pytest.mark.parametrize(
        "model,expected_output",
        [
            ("gpt-3.5-turbo", 17),
            ("gpt-3.5-turbo-0301", 17),
            ("gpt-3.5-turbo-0613", 17),
            ("gpt-3.5-turbo-1106", 17),
            ("gpt-3.5-turbo-instruct", 17),
            ("gpt-3.5-turbo-16k", 17),
            ("gpt-3.5-turbo-16k-0613", 17),
            ("gpt-4", 17),
            ("gpt-4-0314", 17),
            ("gpt-4-0613", 17),
            ("gpt-4-32k", 17),
            ("gpt-4-32k-0314", 17),
            ("gpt-4-1106-preview", 17),
            ("gpt-4-vision-preview", 17),
            ("gpt-4o", 17),
            ("azure/gpt-4o", 17),
            ("claude-2.1", 4),
        ],
    )
    def test_count_message_tokens_with_name(model, expected_output):
        """Notice: name 'John' appears"""

>       assert count_message_tokens(MESSAGES_WITH_NAME, model) == expected_output

tests/test_costs.py:83:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
tokencost/costs.py:68: in count_message_tokens
    raise e
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

messages = [{'content': 'Hello', 'name': 'John', 'role': 'user'}, {'content': 'Hi there!', 'role': 'assistant'}]
model = 'claude-2.1'

    def count_message_tokens(messages: List[Dict[str, str]], model: str) -> int:
        """
        Return the total number of tokens in a prompt's messages.
        Args:
            messages (List[Dict[str, str]]): Message format for prompt requests. e.g.:
                [{ "role": "user", "content": "Hello world"},
                 { "role": "assistant", "content": "How may I assist you today?"}]
            model (str): Name of LLM to choose encoding for.
        Returns:
            Total number of tokens in message.
        """
        model = model.lower()
        model = strip_ft_model_name(model)

        # Anthropic token counting requires a valid API key
        if "claude-" in model:
            logger.warning(
                "Warning: Anthropic token counting API is currently in beta. Please expect differences in costs!"
            )
            if "claude-3-sonnet" in model:
                logger.warning(
                    f"Token counting (beta) is not supported for {model}. Returning num tokens using count from the string."
                )
                # For anthropic<0.39.0 this method is no more supported
                prompt = "".join(message["content"] for message in messages)
                return count_string_tokens(prompt, model)

            ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")

            try:
                client = anthropic.Client(api_key=ANTHROPIC_API_KEY)
>               num_tokens = client.beta.messages.count_tokens(
                    model=model,
                    messages=messages,
                ).input_tokens
E               AttributeError: 'Beta' object has no attribute 'messages'

tokencost/costs.py:60: AttributeError
------------------------------ Captured log call -------------------------------
WARNING  tokencost.costs:costs.py:45 Warning: Anthropic token counting API is currently in beta. Please expect differences in costs!
______ test_calculate_prompt_cost[prompt16-claude-2.1-expected_output16] _______

prompt = [{'content': 'Hello', 'role': 'user'}, {'content': 'Hi there!', 'role': 'assistant'}]
model = 'claude-2.1', expected_output = Decimal('0.000032')

    @pytest.mark.parametrize(
        "prompt,model,expected_output",
        [
            (MESSAGES, "gpt-3.5-turbo", Decimal("0.0000225")),
            (MESSAGES, "gpt-3.5-turbo-0301", Decimal("0.0000255")),
            (MESSAGES, "gpt-3.5-turbo-0613", Decimal("0.0000225")),
            (MESSAGES, "gpt-3.5-turbo-16k", Decimal("0.000045")),
            (MESSAGES, "gpt-3.5-turbo-16k-0613", Decimal("0.000045")),
            (MESSAGES, "gpt-3.5-turbo-1106", Decimal("0.000015")),
            (MESSAGES, "gpt-3.5-turbo-instruct", Decimal("0.0000225")),
            (MESSAGES, "gpt-4", Decimal("0.00045")),
            (MESSAGES, "gpt-4-0314", Decimal("0.00045")),
            (MESSAGES, "gpt-4-32k", Decimal("0.00090")),
            (MESSAGES, "gpt-4-32k-0314", Decimal("0.00090")),
            (MESSAGES, "gpt-4-0613", Decimal("0.00045")),
            (MESSAGES, "gpt-4-1106-preview", Decimal("0.00015")),
            (MESSAGES, "gpt-4-vision-preview", Decimal("0.00015")),
            (MESSAGES, "gpt-4o", Decimal("0.000075")),
            (MESSAGES, "azure/gpt-4o", Decimal("0.000075")),
            (MESSAGES, "claude-2.1", Decimal("0.000032")),
            (STRING, "text-embedding-ada-002", Decimal("0.0000004")),
        ],
    )
    def test_calculate_prompt_cost(prompt, model, expected_output):
        """Test that the cost calculation is correct."""

>       cost = calculate_prompt_cost(prompt, model)

tests/test_costs.py:165:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
tokencost/costs.py:226: in calculate_prompt_cost
    else count_message_tokens(prompt, model)
tokencost/costs.py:68: in count_message_tokens
    raise e
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

messages = [{'content': 'Hello', 'role': 'user'}, {'content': 'Hi there!', 'role': 'assistant'}]
model = 'claude-2.1'

    def count_message_tokens(messages: List[Dict[str, str]], model: str) -> int:
        """
        Return the total number of tokens in a prompt's messages.
        Args:
            messages (List[Dict[str, str]]): Message format for prompt requests. e.g.:
                [{ "role": "user", "content": "Hello world"},
                 { "role": "assistant", "content": "How may I assist you today?"}]
            model (str): Name of LLM to choose encoding for.
        Returns:
            Total number of tokens in message.
        """
        model = model.lower()
        model = strip_ft_model_name(model)

        # Anthropic token counting requires a valid API key
        if "claude-" in model:
            logger.warning(
                "Warning: Anthropic token counting API is currently in beta. Please expect differences in costs!"
            )
            if "claude-3-sonnet" in model:
                logger.warning(
                    f"Token counting (beta) is not supported for {model}. Returning num tokens using count from the string."
                )
                # For anthropic<0.39.0 this method is no more supported
                prompt = "".join(message["content"] for message in messages)
                return count_string_tokens(prompt, model)

            ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")

            try:
                client = anthropic.Client(api_key=ANTHROPIC_API_KEY)
>               num_tokens = client.beta.messages.count_tokens(
                    model=model,
                    messages=messages,
                ).input_tokens
E               AttributeError: 'Beta' object has no attribute 'messages'

tokencost/costs.py:60: AttributeError
------------------------------ Captured log call -------------------------------
WARNING  tokencost.costs:costs.py:45 Warning: Anthropic token counting API is currently in beta. Please expect differences in costs!
=========================== short test summary info ============================
FAILED tests/test_costs.py::test_count_message_tokens[claude-2.1-4] - AttributeError: 'Beta' object has no attribute 'messages'
FAILED tests/test_costs.py::test_count_message_tokens_with_name[claude-2.1-4] - AttributeError: 'Beta' object has no attribute 'messages'
FAILED tests/test_costs.py::test_calculate_prompt_cost[prompt16-claude-2.1-expected_output16] - AttributeError: 'Beta' object has no attribute 'messages'
========================= 3 failed, 95 passed in 7.39s =========================
(base) ➜  tokencost git:(feat/anthropic-token-count)```