test_usage_callback() — langchain Function Reference

Architecture documentation for the test_usage_callback() function in test_usage_callback.py from the langchain codebase.

Function python CoreAbstractions Serialization

Entity Profile

CoreAbstractions→ Serialization→ test_usage_callback() — langchain Function Reference

Dependency Diagram

graph TD
  490888df_da06_36ea_4704_453257730e3c["test_usage_callback()"]
  b9fdf1fe_41aa_aa38_11e1_ae02e5bde19d["test_usage_callback.py"]
  490888df_da06_36ea_4704_453257730e3c -->|defined in| b9fdf1fe_41aa_aa38_11e1_ae02e5bde19d
  style 490888df_da06_36ea_4704_453257730e3c fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

libs/core/tests/unit_tests/callbacks/test_usage_callback.py lines 60–97

def test_usage_callback() -> None:
    llm = FakeChatModelWithResponseMetadata(
        messages=iter(messages), model_name="test_model"
    )

    # Test context manager
    with get_usage_metadata_callback() as cb:
        _ = llm.invoke("Message 1")
        _ = llm.invoke("Message 2")
        total_1_2 = add_usage(usage1, usage2)
        assert cb.usage_metadata == {"test_model": total_1_2}
        _ = llm.invoke("Message 3")
        _ = llm.invoke("Message 4")
        total_3_4 = add_usage(usage3, usage4)
        assert cb.usage_metadata == {"test_model": add_usage(total_1_2, total_3_4)}

    # Test via config
    llm = FakeChatModelWithResponseMetadata(
        messages=iter(messages[:2]), model_name="test_model"
    )
    callback = UsageMetadataCallbackHandler()
    _ = llm.batch(["Message 1", "Message 2"], config={"callbacks": [callback]})
    assert callback.usage_metadata == {"test_model": total_1_2}

    # Test multiple models
    llm_1 = FakeChatModelWithResponseMetadata(
        messages=iter(messages[:2]), model_name="test_model_1"
    )
    llm_2 = FakeChatModelWithResponseMetadata(
        messages=iter(messages[2:4]), model_name="test_model_2"
    )
    callback = UsageMetadataCallbackHandler()
    _ = llm_1.batch(["Message 1", "Message 2"], config={"callbacks": [callback]})
    _ = llm_2.batch(["Message 3", "Message 4"], config={"callbacks": [callback]})
    assert callback.usage_metadata == {
        "test_model_1": total_1_2,
        "test_model_2": total_3_4,
    }