Function test_task_cost

tests/main/test_task.py:35–63

Test that max_cost, max_tokens are respected by Task.run()

(test_settings: Settings)

def test_task_cost(test_settings: Settings):
    """Test that max_cost, max_tokens are respected by Task.run()"""

    set_global(test_settings)
    settings.cache = False
    agent = ChatAgent(ChatAgentConfig(name="Test"))
    agent.llm.reset_usage_cost()
    task = Task(
        agent,
        interactive=False,
        single_round=False,
        system_message="User will send you a number. Repeat the number.",
    )
    sub_agent = ChatAgent(ChatAgentConfig(name="Sub"))
    sub_agent.llm.reset_usage_cost()
    sub = Task(
        sub_agent,
        interactive=False,
        single_round=True,
        system_message="User will send you a number. Return its double",
    )
    task.add_sub_task(sub)
    response = task.run("4", turns=10, max_cost=0.0005, max_tokens=100)
    settings.cache = True
    assert response is not None
    assert response.metadata.status in [
        lr.StatusCode.MAX_COST,
        lr.StatusCode.MAX_TOKENS,
    ]


@pytest.mark.parametrize("restart", [True, False])