test_count_tokens_approximately_usage_metadata_scaling() — langchain Function Reference

Source Code

libs/core/tests/unit_tests/messages/test_utils.py lines 1598–1634

def test_count_tokens_approximately_usage_metadata_scaling() -> None:
    messages = [
        HumanMessage("text"),
        AIMessage(
            "text",
            response_metadata={"model_provider": "openai"},
            usage_metadata={"input_tokens": 0, "output_tokens": 0, "total_tokens": 100},
        ),
        HumanMessage("text"),
        AIMessage(
            "text",
            response_metadata={"model_provider": "openai"},
            usage_metadata={"input_tokens": 0, "output_tokens": 0, "total_tokens": 200},
        ),
    ]

    unscaled = count_tokens_approximately(messages)
    scaled = count_tokens_approximately(messages, use_usage_metadata_scaling=True)

    ratio = scaled / unscaled
    assert 1 <= round(ratio, 1) <= 1.2  # we ceil scale token counts, so can be > 1.2

    messages.extend([ToolMessage("text", tool_call_id="abc123")] * 3)

    unscaled_extended = count_tokens_approximately(messages)
    scaled_extended = count_tokens_approximately(
        messages, use_usage_metadata_scaling=True
    )

    # scaling should still be based on the most recent AIMessage with total_tokens=200
    assert unscaled_extended > unscaled
    assert scaled_extended > scaled

    # And the scaled total should be the unscaled total multiplied by the same ratio.
    # ratio = 200 / unscaled (as of last AI message)
    expected_scaled_extended = math.ceil(unscaled_extended * ratio)
    assert scaled_extended <= expected_scaled_extended <= scaled_extended + 1