Function test_token_usage_tool

tests/main/test_token_usage.py:99–124

Check token usage accumulation with tool/function-call

(fn, stream)

@pytest.mark.parametrize("fn", [True, False])
@pytest.mark.parametrize("stream", [True, False])
def test_token_usage_tool(fn, stream):
    """Check token usage accumulation with tool/function-call"""
    set_global(Settings(cache=False, stream=stream))
    cfg = _TestChatAgentConfig(
        llm=config,
        use_functions_api=fn,
        use_tools=not fn,
        system_message="Use the `capital` tool to tell me the capital of a country",
    )
    agent = ChatAgent(cfg)
    agent.llm.reset_usage_cost()
    agent.enable_message(CapitalTool, use=True, handle=True)

    question = "What is the capital of China?"
    response1 = agent.llm_response(question)
    result = agent.agent_response(response1)
    agent.llm_response(result)
    response3 = agent.llm_response(question)

    assert (
        response3.metadata.usage.prompt_tokens
        >= response1.metadata.usage.prompt_tokens
        + response1.metadata.usage.completion_tokens
        + agent.num_tokens(question)
        + agent.num_tokens(result.content)
    )


@pytest.mark.asyncio