freemo
5a3678cd6c
CI / typecheck (pull_request) Successful in 59s
CI / security (pull_request) Successful in 54s
CI / quality (pull_request) Successful in 40s
CI / lint (pull_request) Successful in 3m21s
CI / build (pull_request) Successful in 17s
CI / helm (pull_request) Successful in 21s
CI / unit_tests (pull_request) Successful in 6m50s
CI / e2e_tests (pull_request) Successful in 19m16s
CI / docker (pull_request) Successful in 1m21s
CI / integration_tests (pull_request) Successful in 21m37s
CI / coverage (pull_request) Successful in 14m51s
CI / status-check (pull_request) Successful in 1s
CI / benchmark-publish (pull_request) Has been skipped
CI / benchmark-regression (pull_request) Successful in 57m40s
The record_usage method performed an unprotected read-modify-write on _daily_costs, a classic TOCTOU race condition. In multi-threaded execution (ThreadPoolExecutor for parallel plan steps), two threads could interleave their .get() read and = write, causing one thread's cost increment to be silently overwritten by the other. Changes: - Add _daily_costs_lock: threading.Lock to CostTracker.__init__ - Wrap the read-modify-write in record_usage with self._daily_costs_lock - Protect reads in check_daily_budget and get_daily_spend with the same lock - Add Behave @concurrency scenario: 20 threads concurrently call record_usage; final daily spend must equal the sum of all individual costs - Add Robot Framework integration smoke test via helper_cost_controls.py cost-tracker-concurrent command All accesses to _daily_costs are now lock-protected. nox -e typecheck passes with zero errors. Existing cost_controls tests continue to pass. ISSUES CLOSED: #1919
1121 lines
37 KiB
Python
1121 lines
37 KiB
Python
"""Step definitions for cost controls feature tests."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import threading
|
|
from typing import Any
|
|
|
|
from behave import given, then, when # type: ignore[import-untyped]
|
|
from behave.runner import Context # type: ignore[import-untyped]
|
|
|
|
# ---- CostEntry steps ----
|
|
|
|
|
|
@given("I import the CostEntry class")
|
|
def step_import_cost_entry(context: Context) -> None:
|
|
from cleveragents.providers.cost_table import CostEntry
|
|
|
|
context.CostEntry = CostEntry
|
|
|
|
|
|
@when("I create a CostEntry with negative input cost")
|
|
def step_create_cost_entry_negative_input(context: Context) -> None:
|
|
context.cost_entry_error = None
|
|
try:
|
|
context.CostEntry(input_cost_per_token=-0.001, output_cost_per_token=0.001)
|
|
except ValueError as exc:
|
|
context.cost_entry_error = exc
|
|
|
|
|
|
@when("I create a CostEntry with negative output cost")
|
|
def step_create_cost_entry_negative_output(context: Context) -> None:
|
|
context.cost_entry_error = None
|
|
try:
|
|
context.CostEntry(input_cost_per_token=0.001, output_cost_per_token=-0.001)
|
|
except ValueError as exc:
|
|
context.cost_entry_error = exc
|
|
|
|
|
|
@then("a ValueError should be raised for CostEntry")
|
|
def step_cost_entry_valueerror(context: Context) -> None:
|
|
assert context.cost_entry_error is not None, "Expected ValueError"
|
|
assert isinstance(context.cost_entry_error, ValueError)
|
|
|
|
|
|
@when(
|
|
"I create a CostEntry with input cost {input_cost:f} and output cost {output_cost:f}"
|
|
)
|
|
def step_create_cost_entry_with_costs(
|
|
context: Context, input_cost: float, output_cost: float
|
|
) -> None:
|
|
context.cost_entry = context.CostEntry(
|
|
input_cost_per_token=input_cost,
|
|
output_cost_per_token=output_cost,
|
|
)
|
|
|
|
|
|
@when(
|
|
"I estimate cost for {input_tokens:d} input tokens and {output_tokens:d} output tokens"
|
|
)
|
|
def step_estimate_cost(context: Context, input_tokens: int, output_tokens: int) -> None:
|
|
context.estimated_cost = context.cost_entry.estimate_cost(
|
|
input_tokens, output_tokens
|
|
)
|
|
|
|
|
|
@then("the estimated cost should be {expected:f}")
|
|
def step_check_estimated_cost(context: Context, expected: float) -> None:
|
|
assert abs(context.estimated_cost - expected) < 1e-9, (
|
|
f"Expected {expected}, got {context.estimated_cost}"
|
|
)
|
|
|
|
|
|
@then("estimating cost with negative input tokens raises ValueError")
|
|
def step_estimate_negative_input(context: Context) -> None:
|
|
try:
|
|
context.cost_entry.estimate_cost(-1, 0)
|
|
raise AssertionError("Expected ValueError")
|
|
except ValueError:
|
|
pass
|
|
|
|
|
|
@then("estimating cost with negative output tokens raises ValueError")
|
|
def step_estimate_negative_output(context: Context) -> None:
|
|
try:
|
|
context.cost_entry.estimate_cost(0, -1)
|
|
raise AssertionError("Expected ValueError")
|
|
except ValueError:
|
|
pass
|
|
|
|
|
|
# ---- ProviderCostTable steps ----
|
|
|
|
|
|
@given("I create a ProviderCostTable with defaults")
|
|
def step_create_default_cost_table(context: Context) -> None:
|
|
from cleveragents.providers.cost_table import ProviderCostTable
|
|
|
|
context.cost_table = ProviderCostTable()
|
|
|
|
|
|
@when('I look up cost for provider "{provider}" model "{model}"')
|
|
def step_lookup_cost(context: Context, provider: str, model: str) -> None:
|
|
context.cost_entry_result = context.cost_table.get_cost_entry(provider, model)
|
|
|
|
|
|
@then("the cost entry should be the default cost")
|
|
def step_check_default_cost(context: Context) -> None:
|
|
from cleveragents.providers.cost_table import ProviderCostTable
|
|
|
|
assert context.cost_entry_result == ProviderCostTable.DEFAULT_COST
|
|
|
|
|
|
@then("the cost entry input cost should be {expected:f}")
|
|
def step_check_input_cost(context: Context, expected: float) -> None:
|
|
assert abs(context.cost_entry_result.input_cost_per_token - expected) < 1e-12, (
|
|
f"Expected {expected}, got {context.cost_entry_result.input_cost_per_token}"
|
|
)
|
|
|
|
|
|
@then("the cost entry output cost should be {expected:f}")
|
|
def step_check_output_cost(context: Context, expected: float) -> None:
|
|
assert abs(context.cost_entry_result.output_cost_per_token - expected) < 1e-12, (
|
|
f"Expected {expected}, got {context.cost_entry_result.output_cost_per_token}"
|
|
)
|
|
|
|
|
|
@given('I create a ProviderCostTable with custom entry for "{provider}" "{model}"')
|
|
def step_create_custom_cost_table(context: Context, provider: str, model: str) -> None:
|
|
from cleveragents.providers.cost_table import CostEntry, ProviderCostTable
|
|
|
|
custom: dict[str, dict[str, CostEntry]] = {
|
|
provider: {
|
|
model: CostEntry(input_cost_per_token=0.001, output_cost_per_token=0.002)
|
|
}
|
|
}
|
|
context.cost_table = ProviderCostTable(custom_entries=custom)
|
|
|
|
|
|
@then("looking up cost for empty provider raises ValueError")
|
|
def step_lookup_empty_provider(context: Context) -> None:
|
|
try:
|
|
context.cost_table.get_cost_entry("", "model")
|
|
raise AssertionError("Expected ValueError")
|
|
except ValueError:
|
|
pass
|
|
|
|
|
|
@then("looking up cost for empty model raises ValueError")
|
|
def step_lookup_empty_model(context: Context) -> None:
|
|
try:
|
|
context.cost_table.get_cost_entry("openai", "")
|
|
raise AssertionError("Expected ValueError")
|
|
except ValueError:
|
|
pass
|
|
|
|
|
|
@when("I list providers from the cost table")
|
|
def step_list_providers(context: Context) -> None:
|
|
context.provider_list = context.cost_table.list_providers()
|
|
|
|
|
|
@then('the provider list should include "{provider}"')
|
|
def step_check_provider_in_list(context: Context, provider: str) -> None:
|
|
assert provider in context.provider_list, (
|
|
f"{provider} not in {context.provider_list}"
|
|
)
|
|
|
|
|
|
@when('I list models for provider "{provider}"')
|
|
def step_list_models(context: Context, provider: str) -> None:
|
|
context.model_list = context.cost_table.list_models(provider)
|
|
|
|
|
|
@then('the model list should include "{model}"')
|
|
def step_check_model_in_list(context: Context, model: str) -> None:
|
|
assert model in context.model_list, f"{model} not in {context.model_list}"
|
|
|
|
|
|
@when("I create a ProviderCostTable with non-dict custom entries")
|
|
def step_create_cost_table_non_dict(context: Context) -> None:
|
|
from cleveragents.providers.cost_table import ProviderCostTable
|
|
|
|
context.cost_table_error = None
|
|
try:
|
|
ProviderCostTable(custom_entries="not-a-dict") # type: ignore[arg-type]
|
|
except TypeError as exc:
|
|
context.cost_table_error = exc
|
|
|
|
|
|
@then("a TypeError should be raised for ProviderCostTable")
|
|
def step_cost_table_typeerror(context: Context) -> None:
|
|
assert context.cost_table_error is not None, "Expected TypeError"
|
|
assert isinstance(context.cost_table_error, TypeError)
|
|
|
|
|
|
# ---- CostMetadata steps ----
|
|
|
|
|
|
@given("I create a fresh CostMetadata")
|
|
def step_create_cost_metadata(context: Context) -> None:
|
|
from cleveragents.domain.models.core.cost_metadata import CostMetadata
|
|
|
|
context.cost_metadata = CostMetadata()
|
|
|
|
|
|
@when(
|
|
"I record usage of {input_t:d} input tokens and {output_t:d} output tokens "
|
|
'at cost {cost:f} for provider "{provider}"'
|
|
)
|
|
def step_record_usage(
|
|
context: Context, input_t: int, output_t: int, cost: float, provider: str
|
|
) -> None:
|
|
context.cost_metadata.record_usage(
|
|
input_tokens=input_t,
|
|
output_tokens=output_t,
|
|
cost=cost,
|
|
provider=provider,
|
|
)
|
|
|
|
|
|
@then("total tokens should be {expected:d}")
|
|
def step_check_total_tokens(context: Context, expected: int) -> None:
|
|
assert context.cost_metadata.total_tokens == expected, (
|
|
f"Expected {expected}, got {context.cost_metadata.total_tokens}"
|
|
)
|
|
|
|
|
|
@then("total cost should be {expected:f}")
|
|
def step_check_total_cost(context: Context, expected: float) -> None:
|
|
assert abs(context.cost_metadata.total_cost - expected) < 1e-9, (
|
|
f"Expected {expected}, got {context.cost_metadata.total_cost}"
|
|
)
|
|
|
|
|
|
@then('provider costs for "{provider}" should be {expected:f}')
|
|
def step_check_provider_costs(context: Context, provider: str, expected: float) -> None:
|
|
actual = context.cost_metadata.provider_costs.get(provider, 0.0)
|
|
assert abs(actual - expected) < 1e-9, f"Expected {expected}, got {actual}"
|
|
|
|
|
|
@then("recording usage with negative input tokens raises ValueError")
|
|
def step_record_negative_input(context: Context) -> None:
|
|
try:
|
|
context.cost_metadata.record_usage(
|
|
input_tokens=-1, output_tokens=0, cost=0.0, provider="test"
|
|
)
|
|
raise AssertionError("Expected ValueError")
|
|
except ValueError:
|
|
pass
|
|
|
|
|
|
@then("recording usage with negative output tokens raises ValueError")
|
|
def step_record_negative_output_tokens(context: Context) -> None:
|
|
try:
|
|
context.cost_metadata.record_usage(
|
|
input_tokens=0, output_tokens=-1, cost=0.0, provider="test"
|
|
)
|
|
raise AssertionError("Expected ValueError")
|
|
except ValueError:
|
|
pass
|
|
|
|
|
|
@then("recording usage with negative cost raises ValueError")
|
|
def step_record_negative_cost(context: Context) -> None:
|
|
try:
|
|
context.cost_metadata.record_usage(
|
|
input_tokens=0, output_tokens=0, cost=-1.0, provider="test"
|
|
)
|
|
raise AssertionError("Expected ValueError")
|
|
except ValueError:
|
|
pass
|
|
|
|
|
|
@then("recording usage with empty provider raises ValueError")
|
|
def step_record_empty_provider(context: Context) -> None:
|
|
try:
|
|
context.cost_metadata.record_usage(
|
|
input_tokens=0, output_tokens=0, cost=0.0, provider=""
|
|
)
|
|
raise AssertionError("Expected ValueError")
|
|
except ValueError:
|
|
pass
|
|
|
|
|
|
@when("I get the display dict from cost metadata")
|
|
def step_get_display_dict(context: Context) -> None:
|
|
context.display_dict = context.cost_metadata.as_display_dict()
|
|
|
|
|
|
@then('the display dict should have key "{key}"')
|
|
def step_check_display_dict_key(context: Context, key: str) -> None:
|
|
assert key in context.display_dict, f"Key {key} not in {context.display_dict}"
|
|
|
|
|
|
@when("I set budget remaining to {value:f}")
|
|
def step_set_budget_remaining(context: Context, value: float) -> None:
|
|
context.cost_metadata.budget_remaining = value
|
|
|
|
|
|
# ---- BudgetExhaustionEvent steps ----
|
|
|
|
|
|
@given("I import the BudgetExhaustionEvent class")
|
|
def step_import_exhaustion_event(context: Context) -> None:
|
|
from cleveragents.domain.models.core.cost_metadata import BudgetExhaustionEvent
|
|
|
|
context.BudgetExhaustionEvent = BudgetExhaustionEvent
|
|
|
|
|
|
@when("I create a BudgetExhaustionEvent with invalid budget type")
|
|
def step_create_invalid_event(context: Context) -> None:
|
|
context.exhaustion_event_error = None
|
|
try:
|
|
context.BudgetExhaustionEvent(budget_type="invalid", limit=10.0, used=5.0)
|
|
except Exception as exc:
|
|
context.exhaustion_event_error = exc
|
|
|
|
|
|
@then("a validation error should be raised for BudgetExhaustionEvent")
|
|
def step_check_exhaustion_validation_error(context: Context) -> None:
|
|
assert context.exhaustion_event_error is not None, "Expected validation error"
|
|
|
|
|
|
@when('I create a BudgetExhaustionEvent with budget type "{btype}"')
|
|
def step_create_valid_event(context: Context, btype: str) -> None:
|
|
context.exhaustion_event = context.BudgetExhaustionEvent(
|
|
budget_type=btype, limit=10.0, used=5.0
|
|
)
|
|
|
|
|
|
@then("the event should be created successfully")
|
|
def step_check_event_created(context: Context) -> None:
|
|
assert context.exhaustion_event is not None
|
|
|
|
|
|
# ---- CostTracker steps ----
|
|
|
|
|
|
@given("I create a CostTracker with no limits")
|
|
def step_create_tracker_no_limits(context: Context) -> None:
|
|
from cleveragents.providers.cost_tracker import CostTracker
|
|
|
|
context.cost_tracker = CostTracker()
|
|
|
|
|
|
@given("I create a CostTracker with plan budget {budget:f}")
|
|
def step_create_tracker_plan_budget(context: Context, budget: float) -> None:
|
|
from cleveragents.providers.cost_tracker import CostTracker
|
|
|
|
context.cost_tracker = CostTracker(budget_per_plan=budget)
|
|
|
|
|
|
@given("I create a CostTracker with daily budget {budget:f}")
|
|
def step_create_tracker_daily_budget(context: Context, budget: float) -> None:
|
|
from cleveragents.providers.cost_tracker import CostTracker
|
|
|
|
context.cost_tracker = CostTracker(budget_per_day=budget)
|
|
context.fallback_cost_tracker = context.cost_tracker
|
|
|
|
|
|
@given("I create a CostTracker with plan budget {plan:f} and daily budget {daily:f}")
|
|
def step_create_tracker_both_budgets(
|
|
context: Context, plan: float, daily: float
|
|
) -> None:
|
|
from cleveragents.providers.cost_tracker import CostTracker
|
|
|
|
context.cost_tracker = CostTracker(budget_per_plan=plan, budget_per_day=daily)
|
|
|
|
|
|
@when(
|
|
"I record tracked usage of {inp:d} input and {out:d} output "
|
|
'for "{provider}" "{model}"'
|
|
)
|
|
def step_record_tracked_usage(
|
|
context: Context, inp: int, out: int, provider: str, model: str
|
|
) -> None:
|
|
if not hasattr(context, "cost_metadata"):
|
|
from cleveragents.domain.models.core.cost_metadata import CostMetadata
|
|
|
|
context.cost_metadata = CostMetadata()
|
|
|
|
context.budget_result = context.cost_tracker.record_usage(
|
|
context.cost_metadata,
|
|
provider=provider,
|
|
model=model,
|
|
input_tokens=inp,
|
|
output_tokens=out,
|
|
)
|
|
|
|
|
|
@then("the budget check should be under budget")
|
|
def step_check_under_budget(context: Context) -> None:
|
|
from cleveragents.providers.cost_tracker import BudgetStatus
|
|
|
|
assert context.budget_result.status == BudgetStatus.UNDER_BUDGET, (
|
|
f"Expected UNDER_BUDGET, got {context.budget_result.status}"
|
|
)
|
|
|
|
|
|
@then("the budget check should be warning or exceeded")
|
|
def step_check_warning_or_exceeded(context: Context) -> None:
|
|
from cleveragents.providers.cost_tracker import BudgetStatus
|
|
|
|
assert context.budget_result.status in (
|
|
BudgetStatus.WARNING,
|
|
BudgetStatus.EXCEEDED,
|
|
), f"Expected WARNING or EXCEEDED, got {context.budget_result.status}"
|
|
|
|
|
|
@then("the budget check should be exceeded")
|
|
def step_check_exceeded(context: Context) -> None:
|
|
from cleveragents.providers.cost_tracker import BudgetStatus
|
|
|
|
assert context.budget_result.status == BudgetStatus.EXCEEDED, (
|
|
f"Expected EXCEEDED, got {context.budget_result.status}"
|
|
)
|
|
|
|
|
|
@then("the cost metadata total tokens should be {expected:d}")
|
|
def step_check_metadata_total_tokens(context: Context, expected: int) -> None:
|
|
assert context.cost_metadata.total_tokens == expected, (
|
|
f"Expected {expected}, got {context.cost_metadata.total_tokens}"
|
|
)
|
|
|
|
|
|
@then("the cost metadata should have budget exhaustion events")
|
|
def step_check_exhaustion_events(context: Context) -> None:
|
|
assert len(context.cost_metadata.budget_exhaustion_events) > 0, (
|
|
"Expected budget exhaustion events"
|
|
)
|
|
|
|
|
|
@when("I create a CostTracker with negative plan budget")
|
|
def step_create_tracker_negative_plan(context: Context) -> None:
|
|
from cleveragents.providers.cost_tracker import CostTracker
|
|
|
|
context.tracker_error = None
|
|
try:
|
|
CostTracker(budget_per_plan=-1.0)
|
|
except ValueError as exc:
|
|
context.tracker_error = exc
|
|
|
|
|
|
@when("I create a CostTracker with negative daily budget")
|
|
def step_create_tracker_negative_daily(context: Context) -> None:
|
|
from cleveragents.providers.cost_tracker import CostTracker
|
|
|
|
context.tracker_error = None
|
|
try:
|
|
CostTracker(budget_per_day=-1.0)
|
|
except ValueError as exc:
|
|
context.tracker_error = exc
|
|
|
|
|
|
@then("a ValueError should be raised for CostTracker")
|
|
def step_check_tracker_valueerror(context: Context) -> None:
|
|
assert context.tracker_error is not None, "Expected ValueError"
|
|
assert isinstance(context.tracker_error, ValueError)
|
|
|
|
|
|
@when(
|
|
'I estimate cost for "{provider}" "{model}" with {inp:d} input and {out:d} output'
|
|
)
|
|
def step_estimate_tracker_cost(
|
|
context: Context, provider: str, model: str, inp: int, out: int
|
|
) -> None:
|
|
context.tracker_estimated_cost = context.cost_tracker.estimate_cost(
|
|
provider, model, inp, out
|
|
)
|
|
|
|
|
|
@then("the estimated tracker cost should be positive")
|
|
def step_check_tracker_positive_cost(context: Context) -> None:
|
|
assert context.tracker_estimated_cost > 0.0, (
|
|
f"Expected positive cost, got {context.tracker_estimated_cost}"
|
|
)
|
|
|
|
|
|
@then("estimating cost with empty provider raises ValueError on tracker")
|
|
def step_estimate_empty_provider_tracker(context: Context) -> None:
|
|
try:
|
|
context.cost_tracker.estimate_cost("", "gpt-4o", 100, 50)
|
|
raise AssertionError("Expected ValueError")
|
|
except ValueError:
|
|
pass
|
|
|
|
|
|
@then("estimating cost with empty model raises ValueError on tracker")
|
|
def step_estimate_empty_model_tracker(context: Context) -> None:
|
|
try:
|
|
context.cost_tracker.estimate_cost("openai", "", 100, 50)
|
|
raise AssertionError("Expected ValueError")
|
|
except ValueError:
|
|
pass
|
|
|
|
|
|
@then("recording tracked usage with negative input raises ValueError")
|
|
def step_record_tracked_negative_input(context: Context) -> None:
|
|
try:
|
|
context.cost_tracker.record_usage(
|
|
context.cost_metadata,
|
|
provider="openai",
|
|
model="gpt-4o",
|
|
input_tokens=-1,
|
|
output_tokens=0,
|
|
)
|
|
raise AssertionError("Expected ValueError")
|
|
except ValueError:
|
|
pass
|
|
|
|
|
|
@when("I check plan budget")
|
|
def step_check_plan_budget(context: Context) -> None:
|
|
context.budget_result = context.cost_tracker.check_plan_budget(
|
|
context.cost_metadata
|
|
)
|
|
|
|
|
|
@when("I check daily budget")
|
|
def step_check_daily_budget(context: Context) -> None:
|
|
context.budget_result = context.cost_tracker.check_daily_budget()
|
|
|
|
|
|
@then("the plan budget check should be under budget")
|
|
def step_plan_under_budget(context: Context) -> None:
|
|
from cleveragents.providers.cost_tracker import BudgetStatus
|
|
|
|
assert context.budget_result.status == BudgetStatus.UNDER_BUDGET
|
|
|
|
|
|
@then("the plan budget check should be exceeded")
|
|
def step_plan_exceeded(context: Context) -> None:
|
|
from cleveragents.providers.cost_tracker import BudgetStatus
|
|
|
|
assert context.budget_result.status == BudgetStatus.EXCEEDED
|
|
|
|
|
|
@then("the daily budget check should be under budget")
|
|
def step_daily_under_budget(context: Context) -> None:
|
|
from cleveragents.providers.cost_tracker import BudgetStatus
|
|
|
|
assert context.budget_result.status == BudgetStatus.UNDER_BUDGET
|
|
|
|
|
|
@then("the daily budget check should be exceeded")
|
|
def step_daily_exceeded(context: Context) -> None:
|
|
from cleveragents.providers.cost_tracker import BudgetStatus
|
|
|
|
assert context.budget_result.status == BudgetStatus.EXCEEDED
|
|
|
|
|
|
@then("the daily spend should be {expected:f}")
|
|
def step_check_daily_spend(context: Context, expected: float) -> None:
|
|
actual = context.cost_tracker.get_daily_spend()
|
|
assert abs(actual - expected) < 1e-9, f"Expected {expected}, got {actual}"
|
|
|
|
|
|
@then("the tracker plan budget should be {expected:f}")
|
|
def step_check_tracker_plan_budget(context: Context, expected: float) -> None:
|
|
assert context.cost_tracker.budget_per_plan == expected
|
|
|
|
|
|
@then("the tracker daily budget should be {expected:f}")
|
|
def step_check_tracker_daily_budget(context: Context, expected: float) -> None:
|
|
assert context.cost_tracker.budget_per_day == expected
|
|
|
|
|
|
@then("getting cost entry with empty provider raises ValueError")
|
|
def step_get_cost_entry_empty_provider(context: Context) -> None:
|
|
try:
|
|
context.cost_tracker.get_cost_entry("", "gpt-4o")
|
|
raise AssertionError("Expected ValueError")
|
|
except ValueError:
|
|
pass
|
|
|
|
|
|
@then("getting cost entry with empty model raises ValueError")
|
|
def step_get_cost_entry_empty_model(context: Context) -> None:
|
|
try:
|
|
context.cost_tracker.get_cost_entry("openai", "")
|
|
raise AssertionError("Expected ValueError")
|
|
except ValueError:
|
|
pass
|
|
|
|
|
|
@then("recording tracked usage with empty provider raises ValueError")
|
|
def step_record_tracked_empty_provider(context: Context) -> None:
|
|
try:
|
|
context.cost_tracker.record_usage(
|
|
context.cost_metadata,
|
|
provider="",
|
|
model="gpt-4o",
|
|
input_tokens=1,
|
|
output_tokens=1,
|
|
)
|
|
raise AssertionError("Expected ValueError")
|
|
except ValueError:
|
|
pass
|
|
|
|
|
|
@then("recording tracked usage with empty model raises ValueError")
|
|
def step_record_tracked_empty_model(context: Context) -> None:
|
|
try:
|
|
context.cost_tracker.record_usage(
|
|
context.cost_metadata,
|
|
provider="openai",
|
|
model="",
|
|
input_tokens=1,
|
|
output_tokens=1,
|
|
)
|
|
raise AssertionError("Expected ValueError")
|
|
except ValueError:
|
|
pass
|
|
|
|
|
|
# ---- CostTracker concurrency steps ----
|
|
|
|
|
|
@when(
|
|
"{n:d} threads concurrently record usage of {inp:d} input and {out:d} output "
|
|
'for "{provider}" "{model}"'
|
|
)
|
|
def step_concurrent_record_usage(
|
|
context: Context, n: int, inp: int, out: int, provider: str, model: str
|
|
) -> None:
|
|
from cleveragents.domain.models.core.cost_metadata import CostMetadata
|
|
|
|
# Each thread gets its own CostMetadata to avoid data races on that object.
|
|
# We only care about the shared _daily_costs in CostTracker.
|
|
errors: list[Exception] = []
|
|
barrier = threading.Barrier(n)
|
|
|
|
def worker() -> None:
|
|
meta = CostMetadata()
|
|
try:
|
|
barrier.wait() # All threads start simultaneously
|
|
context.cost_tracker.record_usage(
|
|
meta,
|
|
provider=provider,
|
|
model=model,
|
|
input_tokens=inp,
|
|
output_tokens=out,
|
|
)
|
|
except Exception as exc:
|
|
errors.append(exc)
|
|
|
|
threads = [threading.Thread(target=worker) for _ in range(n)]
|
|
for t in threads:
|
|
t.start()
|
|
for t in threads:
|
|
t.join()
|
|
|
|
assert not errors, f"Threads raised exceptions: {errors}"
|
|
|
|
# Store expected total for the assertion step
|
|
single_cost = context.cost_tracker.estimate_cost(provider, model, inp, out)
|
|
context.expected_concurrent_total = single_cost * n
|
|
|
|
|
|
@then("the daily spend should equal the sum of all concurrent costs")
|
|
def step_check_concurrent_daily_spend(context: Context) -> None:
|
|
actual = context.cost_tracker.get_daily_spend()
|
|
expected = context.expected_concurrent_total
|
|
assert abs(actual - expected) < 1e-9, (
|
|
f"Expected daily spend {expected:.6f}, got {actual:.6f} — "
|
|
"possible race condition: cost increments were lost"
|
|
)
|
|
|
|
|
|
# ---- FallbackSelector steps ----
|
|
|
|
|
|
@given("I create a ProviderRegistry with no API keys for fallback")
|
|
def step_create_registry_no_keys(context: Context) -> None:
|
|
import os
|
|
|
|
from cleveragents.config.settings import Settings
|
|
from cleveragents.providers.registry import ProviderRegistry
|
|
|
|
# Temporarily clear all provider API keys to simulate no configured providers
|
|
key_env_vars = [
|
|
"OPENAI_API_KEY",
|
|
"ANTHROPIC_API_KEY",
|
|
"GOOGLE_API_KEY",
|
|
"AZURE_OPENAI_API_KEY",
|
|
"OPENROUTER_API_KEY",
|
|
"GEMINI_API_KEY",
|
|
"GOOGLE_GENAI_API_KEY",
|
|
"GROQ_API_KEY",
|
|
"TOGETHER_API_KEY",
|
|
"COHERE_API_KEY",
|
|
]
|
|
saved: dict[str, str | None] = {}
|
|
for var in key_env_vars:
|
|
saved[var] = os.environ.pop(var, None)
|
|
try:
|
|
settings = Settings()
|
|
context.fallback_registry = ProviderRegistry(settings)
|
|
finally:
|
|
for var, val in saved.items():
|
|
if val is not None:
|
|
os.environ[var] = val
|
|
|
|
|
|
@given("I create a ProviderRegistry with mock provider for fallback")
|
|
def step_create_registry_mock(context: Context) -> None:
|
|
import os
|
|
|
|
from cleveragents.config.settings import Settings
|
|
from cleveragents.providers.registry import ProviderRegistry
|
|
|
|
# Clear all real API keys, the mock provider is always discovered
|
|
key_env_vars = [
|
|
"OPENAI_API_KEY",
|
|
"ANTHROPIC_API_KEY",
|
|
"GOOGLE_API_KEY",
|
|
"AZURE_OPENAI_API_KEY",
|
|
"OPENROUTER_API_KEY",
|
|
"GEMINI_API_KEY",
|
|
"GOOGLE_GENAI_API_KEY",
|
|
"GROQ_API_KEY",
|
|
"TOGETHER_API_KEY",
|
|
"COHERE_API_KEY",
|
|
]
|
|
saved: dict[str, str | None] = {}
|
|
for var in key_env_vars:
|
|
saved[var] = os.environ.pop(var, None)
|
|
try:
|
|
settings = Settings()
|
|
context.fallback_registry = ProviderRegistry(settings)
|
|
finally:
|
|
for var, val in saved.items():
|
|
if val is not None:
|
|
os.environ[var] = val
|
|
|
|
|
|
@given("I create a ProviderRegistry with openai key for fallback")
|
|
def step_create_registry_openai(context: Context) -> None:
|
|
from cleveragents.config.settings import Settings
|
|
from cleveragents.providers.registry import ProviderRegistry
|
|
|
|
old_key = os.environ.get("OPENAI_API_KEY")
|
|
os.environ["OPENAI_API_KEY"] = "test-key-for-fallback"
|
|
try:
|
|
settings = Settings()
|
|
context.fallback_registry = ProviderRegistry(settings)
|
|
finally:
|
|
if old_key is None:
|
|
os.environ.pop("OPENAI_API_KEY", None)
|
|
else:
|
|
os.environ["OPENAI_API_KEY"] = old_key
|
|
|
|
|
|
@when("I create a FallbackSelector with the registry")
|
|
def step_create_fallback_selector(context: Context) -> None:
|
|
from cleveragents.providers.fallback_selector import FallbackSelector
|
|
|
|
tracker = getattr(context, "fallback_cost_tracker", None)
|
|
context.fallback_selector = FallbackSelector(
|
|
registry=context.fallback_registry,
|
|
cost_tracker=tracker,
|
|
)
|
|
|
|
|
|
@when("I select a fallback provider")
|
|
def step_select_fallback(context: Context) -> None:
|
|
context.fallback_result = context.fallback_selector.select()
|
|
|
|
|
|
@when("I select a fallback provider requiring tool calls")
|
|
def step_select_fallback_tool_calls(context: Context) -> None:
|
|
context.fallback_result = context.fallback_selector.select(require_tool_calls=True)
|
|
|
|
|
|
@then("the fallback result provider should be None")
|
|
def step_check_fallback_none(context: Context) -> None:
|
|
assert context.fallback_result.provider_type is None, (
|
|
f"Expected None, got {context.fallback_result.provider_type}"
|
|
)
|
|
|
|
|
|
@then("the fallback result should have skipped entries")
|
|
def step_check_fallback_skipped(context: Context) -> None:
|
|
assert len(context.fallback_result.skipped) > 0
|
|
|
|
|
|
@then("the fallback result should skip mock for missing tool calls")
|
|
def step_check_mock_skipped_tool_calls(context: Context) -> None:
|
|
# All configured providers were skipped (none configured or missing caps)
|
|
assert len(context.fallback_result.skipped) > 0
|
|
|
|
|
|
@then('the fallback result provider type should be "{expected}"')
|
|
def step_check_fallback_provider_type(context: Context, expected: str) -> None:
|
|
assert context.fallback_result.provider_type is not None
|
|
assert context.fallback_result.provider_type.value == expected, (
|
|
f"Expected {expected}, got {context.fallback_result.provider_type.value}"
|
|
)
|
|
|
|
|
|
@when('I create a FallbackSelector with custom order "{providers}"')
|
|
def step_create_fallback_custom_order(context: Context, providers: str) -> None:
|
|
from cleveragents.providers.fallback_selector import FallbackSelector
|
|
|
|
provider_list = [p.strip() for p in providers.split(",")]
|
|
context.fallback_selector = FallbackSelector(
|
|
registry=context.fallback_registry,
|
|
fallback_providers=provider_list,
|
|
)
|
|
|
|
|
|
@when("I create a FallbackSelector with None registry")
|
|
def step_create_fallback_none_registry(context: Context) -> None:
|
|
from cleveragents.providers.fallback_selector import FallbackSelector
|
|
|
|
context.fallback_error = None
|
|
try:
|
|
FallbackSelector(registry=None) # type: ignore[arg-type]
|
|
except TypeError as exc:
|
|
context.fallback_error = exc
|
|
|
|
|
|
@then("a TypeError should be raised for FallbackSelector")
|
|
def step_check_fallback_typeerror(context: Context) -> None:
|
|
assert context.fallback_error is not None, "Expected TypeError"
|
|
assert isinstance(context.fallback_error, TypeError)
|
|
|
|
|
|
@when("I create a FallbackSelector with empty string in fallback list")
|
|
def step_create_fallback_empty_string(context: Context) -> None:
|
|
from cleveragents.providers.fallback_selector import FallbackSelector
|
|
|
|
context.fallback_error = None
|
|
try:
|
|
FallbackSelector(
|
|
registry=context.fallback_registry,
|
|
fallback_providers=["openai", ""],
|
|
)
|
|
except ValueError as exc:
|
|
context.fallback_error = exc
|
|
|
|
|
|
@then("a ValueError should be raised for FallbackSelector")
|
|
def step_check_fallback_valueerror(context: Context) -> None:
|
|
assert context.fallback_error is not None, "Expected ValueError"
|
|
assert isinstance(context.fallback_error, ValueError)
|
|
|
|
|
|
@when("I create a FallbackSelector with cost tracker")
|
|
def step_create_fallback_with_tracker(context: Context) -> None:
|
|
from cleveragents.providers.fallback_selector import FallbackSelector
|
|
|
|
context.fallback_selector = FallbackSelector(
|
|
registry=context.fallback_registry,
|
|
cost_tracker=context.fallback_cost_tracker,
|
|
)
|
|
|
|
|
|
# ---- Settings config keys ----
|
|
|
|
|
|
@when("I load cost control settings with defaults")
|
|
def step_load_default_settings(context: Context) -> None:
|
|
from cleveragents.config.settings import Settings
|
|
|
|
context.cost_settings = Settings()
|
|
|
|
|
|
@then("budget_per_plan should be None")
|
|
def step_check_budget_per_plan_none(context: Context) -> None:
|
|
assert context.cost_settings.budget_per_plan is None
|
|
|
|
|
|
@then("budget_per_day should be None")
|
|
def step_check_budget_per_day_none(context: Context) -> None:
|
|
assert context.cost_settings.budget_per_day is None
|
|
|
|
|
|
@then("fallback_providers should be an empty list")
|
|
def step_check_fallback_empty(context: Context) -> None:
|
|
assert context.cost_settings.fallback_providers == []
|
|
|
|
|
|
@when("I load cost control settings with budget_per_plan {value:f}")
|
|
def step_load_settings_plan_budget(context: Context, value: float) -> None:
|
|
from cleveragents.config.settings import Settings
|
|
|
|
old = os.environ.get("CLEVERAGENTS_BUDGET_PER_PLAN")
|
|
os.environ["CLEVERAGENTS_BUDGET_PER_PLAN"] = str(value)
|
|
try:
|
|
context.cost_settings = Settings()
|
|
finally:
|
|
if old is None:
|
|
os.environ.pop("CLEVERAGENTS_BUDGET_PER_PLAN", None)
|
|
else:
|
|
os.environ["CLEVERAGENTS_BUDGET_PER_PLAN"] = old
|
|
|
|
|
|
@then("budget_per_plan should be {expected:f}")
|
|
def step_check_budget_per_plan_value(context: Context, expected: float) -> None:
|
|
assert context.cost_settings.budget_per_plan == expected, (
|
|
f"Expected {expected}, got {context.cost_settings.budget_per_plan}"
|
|
)
|
|
|
|
|
|
@when("I load cost control settings with budget_per_day {value:f}")
|
|
def step_load_settings_daily_budget(context: Context, value: float) -> None:
|
|
from cleveragents.config.settings import Settings
|
|
|
|
old = os.environ.get("CLEVERAGENTS_BUDGET_PER_DAY")
|
|
os.environ["CLEVERAGENTS_BUDGET_PER_DAY"] = str(value)
|
|
try:
|
|
context.cost_settings = Settings()
|
|
finally:
|
|
if old is None:
|
|
os.environ.pop("CLEVERAGENTS_BUDGET_PER_DAY", None)
|
|
else:
|
|
os.environ["CLEVERAGENTS_BUDGET_PER_DAY"] = old
|
|
|
|
|
|
@then("budget_per_day should be {expected:f}")
|
|
def step_check_budget_per_day_value(context: Context, expected: float) -> None:
|
|
assert context.cost_settings.budget_per_day == expected, (
|
|
f"Expected {expected}, got {context.cost_settings.budget_per_day}"
|
|
)
|
|
|
|
|
|
# ---- Plan model steps ----
|
|
|
|
|
|
def _make_minimal_plan(
|
|
cost_meta: Any = None,
|
|
) -> Any:
|
|
"""Create a minimal Plan v3 model for testing."""
|
|
from cleveragents.domain.models.core.plan import (
|
|
NamespacedName,
|
|
Plan,
|
|
PlanIdentity,
|
|
)
|
|
|
|
return Plan(
|
|
identity=PlanIdentity(plan_id="01ARZJG0J0000000000000000A"),
|
|
namespaced_name=NamespacedName(namespace="local", name="test-plan"),
|
|
description="Test plan for cost metadata",
|
|
action_name="local/test-action",
|
|
cost_metadata=cost_meta,
|
|
)
|
|
|
|
|
|
@given("I create a minimal Plan with cost metadata")
|
|
def step_create_plan_with_cost(context: Context) -> None:
|
|
from cleveragents.domain.models.core.cost_metadata import CostMetadata
|
|
|
|
meta = CostMetadata(total_tokens=100, total_cost=0.01)
|
|
context.plan_with_cost = _make_minimal_plan(cost_meta=meta)
|
|
|
|
|
|
@then("the plan should have cost_metadata set")
|
|
def step_check_plan_cost_metadata(context: Context) -> None:
|
|
assert context.plan_with_cost.cost_metadata is not None
|
|
|
|
|
|
@when("I get the CLI dict from the plan")
|
|
def step_get_cli_dict(context: Context) -> None:
|
|
plan = getattr(context, "plan_with_cost", None) or getattr(
|
|
context, "plan_without_cost", None
|
|
)
|
|
context.cli_dict = plan.as_cli_dict()
|
|
|
|
|
|
@then('the CLI dict should have a "{key}" key')
|
|
def step_check_cli_dict_has_key(context: Context, key: str) -> None:
|
|
assert key in context.cli_dict, f"Key {key} not in {list(context.cli_dict.keys())}"
|
|
|
|
|
|
@given("I create a minimal Plan without cost metadata")
|
|
def step_create_plan_without_cost(context: Context) -> None:
|
|
context.plan_without_cost = _make_minimal_plan(cost_meta=None)
|
|
|
|
|
|
@then('the CLI dict should not have a "{key}" key')
|
|
def step_check_cli_dict_no_key(context: Context, key: str) -> None:
|
|
assert key not in context.cli_dict, f"Key {key} unexpectedly in CLI dict"
|
|
|
|
|
|
@when('I add a budget exhaustion event of type "{btype}"')
|
|
def step_add_exhaustion_event(context: Context, btype: str) -> None:
|
|
from cleveragents.domain.models.core.cost_metadata import BudgetExhaustionEvent
|
|
|
|
event = BudgetExhaustionEvent(
|
|
budget_type=btype,
|
|
limit=10.0,
|
|
used=11.0,
|
|
provider="openai",
|
|
model="gpt-4o",
|
|
)
|
|
context.cost_metadata.budget_exhaustion_events.append(event)
|
|
|
|
|
|
# ---- Additional coverage steps ----
|
|
|
|
|
|
@given(
|
|
"I create a ProviderCostTable with custom entry for new provider "
|
|
'"{provider}" model "{model}"'
|
|
)
|
|
def step_create_cost_table_new_provider(
|
|
context: Context, provider: str, model: str
|
|
) -> None:
|
|
from cleveragents.providers.cost_table import CostEntry, ProviderCostTable
|
|
|
|
custom: dict[str, dict[str, CostEntry]] = {
|
|
provider: {
|
|
model: CostEntry(input_cost_per_token=0.005, output_cost_per_token=0.01)
|
|
}
|
|
}
|
|
context.cost_table = ProviderCostTable(custom_entries=custom)
|
|
|
|
|
|
@then("listing models for empty provider raises ValueError")
|
|
def step_list_models_empty_provider(context: Context) -> None:
|
|
try:
|
|
context.cost_table.list_models("")
|
|
raise AssertionError("Expected ValueError")
|
|
except ValueError:
|
|
pass
|
|
|
|
|
|
@then("estimating cost with negative input tokens raises ValueError on tracker")
|
|
def step_estimate_negative_input_tracker(context: Context) -> None:
|
|
try:
|
|
context.cost_tracker.estimate_cost("openai", "gpt-4o", -1, 0)
|
|
raise AssertionError("Expected ValueError")
|
|
except ValueError:
|
|
pass
|
|
|
|
|
|
@then("estimating cost with negative output tokens raises ValueError on tracker")
|
|
def step_estimate_negative_output_tracker(context: Context) -> None:
|
|
try:
|
|
context.cost_tracker.estimate_cost("openai", "gpt-4o", 0, -1)
|
|
raise AssertionError("Expected ValueError")
|
|
except ValueError:
|
|
pass
|
|
|
|
|
|
@then("recording tracked usage with negative output raises ValueError")
|
|
def step_record_tracked_negative_output(context: Context) -> None:
|
|
try:
|
|
context.cost_tracker.record_usage(
|
|
context.cost_metadata,
|
|
provider="openai",
|
|
model="gpt-4o",
|
|
input_tokens=0,
|
|
output_tokens=-1,
|
|
)
|
|
raise AssertionError("Expected ValueError")
|
|
except ValueError:
|
|
pass
|
|
|
|
|
|
@then("the budget check should be warning")
|
|
def step_check_warning(context: Context) -> None:
|
|
from cleveragents.providers.cost_tracker import BudgetStatus
|
|
|
|
assert context.budget_result.status == BudgetStatus.WARNING, (
|
|
f"Expected WARNING, got {context.budget_result.status}"
|
|
)
|
|
|
|
|
|
@when('I get cost entry for "{provider}" "{model}"')
|
|
def step_get_cost_entry(context: Context, provider: str, model: str) -> None:
|
|
context.returned_cost_entry = context.cost_tracker.get_cost_entry(provider, model)
|
|
|
|
|
|
@then("the returned cost entry should have positive input cost")
|
|
def step_check_returned_cost_entry(context: Context) -> None:
|
|
assert context.returned_cost_entry.input_cost_per_token > 0
|
|
|
|
|
|
@then('the fallback skipped should mention "{reason}"')
|
|
def step_check_fallback_skipped_reason(context: Context, reason: str) -> None:
|
|
reasons = [r for _, r in context.fallback_result.skipped]
|
|
assert any(reason in r for r in reasons), (
|
|
f"Expected '{reason}' in skipped reasons: {reasons}"
|
|
)
|
|
|
|
|
|
@when("I select a fallback provider requiring streaming")
|
|
def step_select_fallback_streaming(context: Context) -> None:
|
|
context.fallback_result = context.fallback_selector.select(require_streaming=True)
|
|
|
|
|
|
@when("I select a fallback provider requiring vision")
|
|
def step_select_fallback_vision(context: Context) -> None:
|
|
context.fallback_result = context.fallback_selector.select(require_vision=True)
|
|
|
|
|
|
@when("I select a fallback provider requiring json mode")
|
|
def step_select_fallback_json_mode(context: Context) -> None:
|
|
context.fallback_result = context.fallback_selector.select(require_json_mode=True)
|
|
|
|
|
|
@then("the fallback result should skip for missing streaming")
|
|
def step_check_skip_streaming(context: Context) -> None:
|
|
reasons = [r for _, r in context.fallback_result.skipped]
|
|
assert len(reasons) > 0, "Expected skipped entries"
|
|
|
|
|
|
@then("the fallback result should skip for missing vision")
|
|
def step_check_skip_vision(context: Context) -> None:
|
|
reasons = [r for _, r in context.fallback_result.skipped]
|
|
assert len(reasons) > 0, "Expected skipped entries"
|
|
|
|
|
|
@then("the fallback result should skip for missing json mode")
|
|
def step_check_skip_json_mode(context: Context) -> None:
|
|
reasons = [r for _, r in context.fallback_result.skipped]
|
|
assert len(reasons) > 0, "Expected skipped entries"
|