fix(governance): resolve lint and test issues in cost tracking implementation

This commit is contained in:
2026-04-21 08:21:29 +00:00
parent ad6b4fd7ce
commit 0b13d3cbb5
2 changed files with 220 additions and 314 deletions
@@ -2,325 +2,230 @@
from __future__ import annotations
from behave import given, then, when
from behave import given, when, then
from cleveragents.core.cost_tracking import CostTrackingService
from cleveragents.core.exceptions import BudgetExceededError
from cleveragents.domain.models.core.cost_budget import SessionCostBudget
from cleveragents.domain.models.core.cost_budget import (
BudgetCheckResult,
BudgetLevel,
SessionCostBudget,
ThreadSafeOrgCostAccumulator,
)
from cleveragents.domain.models.core.cost_metadata import CostMetadata
@given("a cost tracking service is initialized")
def step_init_cost_tracking_service(context: object) -> None:
"""Initialize cost tracking service."""
context.cost_service = CostTrackingService()
@given("a cost tracking service")
def step_create_cost_tracking_service(context):
"""Create a cost tracking service."""
context.cost_tracking_service = CostTrackingService()
@given("a session cost budget of ${amount} is configured")
def step_configure_session_budget(context: object, amount: str) -> None:
"""Configure session cost budget."""
budget_amount = float(amount)
context.session_budget = SessionCostBudget(max_cost_usd=budget_amount)
@given("a session budget of ${max_cost}")
def step_create_session_budget(context, max_cost):
"""Create a session budget."""
context.session_budget = SessionCostBudget(max_cost_usd=float(max_cost))
@given("a cost metadata tracker is created")
def step_create_cost_metadata(context: object) -> None:
"""Create cost metadata tracker."""
@given("an organization cost accumulator with limit ${max_cost}")
def step_create_org_accumulator(context, max_cost):
"""Create an organization cost accumulator."""
context.org_accumulator = ThreadSafeOrgCostAccumulator(
max_cost_usd=float(max_cost)
)
@given("cost metadata")
def step_create_cost_metadata(context):
"""Create cost metadata."""
context.cost_metadata = CostMetadata()
@when("an LLM call costs ${cost} with {input_tokens} input tokens and {output_tokens} output tokens")
def step_record_llm_call(
context: object, cost: str, input_tokens: str, output_tokens: str
) -> None:
"""Record an LLM call cost."""
cost_amount = float(cost)
input_tok = int(input_tokens)
output_tok = int(output_tokens)
def step_llm_call_with_tokens(context, cost, input_tokens, output_tokens):
"""Record an LLM call with specific token counts."""
context.proposed_cost = float(cost)
context.input_tokens = int(input_tokens)
context.output_tokens = int(output_tokens)
context.cost_service.record_cost(
context.cost_metadata,
context.session_budget,
None,
input_tokens=input_tok,
output_tokens=output_tok,
cost=cost_amount,
provider="openai",
@when("an LLM call costs ${cost} from {provider} with model {model}")
def step_llm_call_with_provider_model(context, cost, provider, model):
"""Record an LLM call with provider and model information."""
context.proposed_cost = float(cost)
context.provider = provider
context.model = model
context.input_tokens = 100
context.output_tokens = 50
@when("I check if the cost is allowed")
def step_check_budget(context):
"""Check if the proposed cost is allowed."""
session_budget = getattr(context, 'session_budget', None)
org_accumulator = getattr(context, 'org_accumulator', None)
proposed_cost = getattr(context, 'proposed_cost', 0.0)
context.budget_check_result = context.cost_tracking_service.check_budget(
session_budget=session_budget,
org_accumulator=org_accumulator,
proposed_cost=proposed_cost,
)
@then("the cost metadata should show:")
def step_verify_cost_metadata(context: object) -> None:
"""Verify cost metadata values."""
for row in context.table:
field = row["field"]
expected = float(row["value"]) if "." in row["value"] else int(row["value"])
actual = getattr(context.cost_metadata, field)
assert actual == expected, f"Expected {field}={expected}, got {actual}"
@then("the session budget should show ${remaining} remaining")
def step_verify_remaining_budget(context: object, remaining: str) -> None:
"""Verify remaining budget."""
expected = float(remaining)
actual = context.session_budget.remaining()
assert actual == expected, f"Expected remaining=${expected}, got ${actual}"
@given("the session budget is ${amount}")
def step_set_session_budget(context: object, amount: str) -> None:
"""Set session budget amount."""
budget_amount = float(amount)
context.session_budget = SessionCostBudget(max_cost_usd=budget_amount)
context.cost_metadata = CostMetadata()
@when("an LLM call costs ${cost}")
def step_record_single_cost(context: object, cost: str) -> None:
"""Record a single LLM call cost."""
cost_amount = float(cost)
context.cost_service.record_cost(
context.cost_metadata,
context.session_budget,
None,
input_tokens=100,
output_tokens=50,
cost=cost_amount,
@when("I record the cost")
def step_record_cost(context):
"""Record the cost in the tracking service."""
session_budget = getattr(context, 'session_budget', None)
org_accumulator = getattr(context, 'org_accumulator', None)
cost_metadata = getattr(context, 'cost_metadata', CostMetadata())
proposed_cost = getattr(context, 'proposed_cost', 0.0)
input_tokens = getattr(context, 'input_tokens', 100)
output_tokens = getattr(context, 'output_tokens', 50)
context.cost_tracking_service.record_cost(
cost_metadata=cost_metadata,
session_budget=session_budget,
org_accumulator=org_accumulator,
input_tokens=input_tokens,
output_tokens=output_tokens,
cost=proposed_cost,
provider="openai",
model="gpt-4",
)
@when("another LLM call costs ${cost}")
def step_record_another_cost(context: object, cost: str) -> None:
"""Record another LLM call cost."""
cost_amount = float(cost)
@when("I enforce the budget")
def step_enforce_budget(context):
"""Enforce the budget limits."""
session_budget = getattr(context, 'session_budget', None)
org_accumulator = getattr(context, 'org_accumulator', None)
cost_metadata = getattr(context, 'cost_metadata', CostMetadata())
try:
context.cost_service.record_cost(
context.cost_metadata,
context.session_budget,
None,
input_tokens=100,
output_tokens=50,
cost=cost_amount,
context.cost_tracking_service.enforce_budget(
session_budget=session_budget,
org_accumulator=org_accumulator,
cost_metadata=cost_metadata,
provider="openai",
)
context.cost_service.enforce_budget(
context.session_budget, None, context.cost_metadata, provider="openai"
model="gpt-4",
)
context.budget_exceeded = False
except BudgetExceededError as e:
except Exception as e:
context.budget_exceeded = True
context.budget_error = e
context.budget_error = str(e)
@then("the second call should raise BudgetExceededError")
def step_verify_budget_exceeded(context: object) -> None:
"""Verify budget exceeded error was raised."""
assert context.budget_exceeded, "Expected BudgetExceededError to be raised"
@then("the error should indicate plan budget exceeded")
def step_verify_error_type(context: object) -> None:
"""Verify error indicates plan budget exceeded."""
assert context.budget_error.budget_type == "plan"
@when("an LLM call from {provider} costs ${cost}")
def step_record_provider_cost(context: object, provider: str, cost: str) -> None:
"""Record cost from specific provider."""
cost_amount = float(cost)
context.cost_service.record_cost(
context.cost_metadata,
context.session_budget,
None,
input_tokens=100,
output_tokens=50,
cost=cost_amount,
provider=provider.lower(),
@then("the cost should be allowed")
def step_cost_allowed(context):
"""Verify that the cost is allowed."""
assert context.budget_check_result.allowed, (
f"Cost should be allowed but got: {context.budget_check_result.reason}"
)
@then("the cost metadata should show provider breakdown:")
def step_verify_provider_breakdown(context: object) -> None:
"""Verify provider cost breakdown."""
for row in context.table:
provider = row["provider"]
expected = float(row["cost"])
actual = context.cost_metadata.provider_costs.get(provider, 0.0)
assert actual == expected, f"Expected {provider}=${expected}, got ${actual}"
@when("${amount} has been spent")
def step_spend_amount(context: object, amount: str) -> None:
"""Spend a specific amount."""
cost_amount = float(amount)
context.cost_service.record_cost(
context.cost_metadata,
context.session_budget,
None,
input_tokens=100,
output_tokens=50,
cost=cost_amount,
provider="openai",
@then("the cost should not be allowed")
def step_cost_not_allowed(context):
"""Verify that the cost is not allowed."""
assert not context.budget_check_result.allowed, (
"Cost should not be allowed"
)
@then("budget status should show:")
def step_verify_budget_status(context: object) -> None:
"""Verify budget status."""
status = context.cost_service.get_budget_status(context.session_budget, None)
session_status = status.get("session", {})
for row in context.table:
field = row["field"]
expected = float(row["value"])
actual = session_status.get(field)
assert actual == expected, f"Expected {field}={expected}, got {actual}"
@when("a check for ${cost} cost is performed")
def step_check_budget(context: object, cost: str) -> None:
"""Check if cost fits in budget."""
cost_amount = float(cost)
context.budget_check = context.cost_service.check_budget(
context.session_budget, None, cost_amount
@then("the exceeded level should be {level}")
def step_check_exceeded_level(context, level):
"""Verify the exceeded budget level."""
expected_level = BudgetLevel[level.upper()]
assert context.budget_check_result.exceeded_level == expected_level, (
f"Expected {expected_level} but got {context.budget_check_result.exceeded_level}"
)
@then("the budget check should return warning=true")
def step_verify_warning(context: object) -> None:
"""Verify warning flag is set."""
assert context.budget_check.warning is True
@then("the check should still be allowed")
def step_verify_allowed(context: object) -> None:
"""Verify check is still allowed."""
assert context.budget_check.allowed is True
@when("an LLM call costs ${cost} from {provider} using {model}")
def step_record_cost_with_model(
context: object, cost: str, provider: str, model: str
) -> None:
"""Record cost with provider and model."""
cost_amount = float(cost)
context.cost_service.record_cost(
context.cost_metadata,
context.session_budget,
None,
input_tokens=100,
output_tokens=50,
cost=cost_amount,
provider=provider.lower(),
model=model,
@then("a warning should be issued")
def step_warning_issued(context):
"""Verify that a warning is issued."""
assert context.budget_check_result.warning, (
"Warning should be issued"
)
try:
context.cost_service.enforce_budget(
context.session_budget,
None,
context.cost_metadata,
provider=provider.lower(),
model=model,
)
context.budget_exceeded = False
except BudgetExceededError as e:
context.budget_exceeded = True
context.budget_error = e
@then("a BudgetExceededError should be raised")
def step_verify_error_raised(context: object) -> None:
"""Verify error was raised."""
assert context.budget_exceeded, "Expected BudgetExceededError"
@then("no warning should be issued")
def step_no_warning(context):
"""Verify that no warning is issued."""
assert not context.budget_check_result.warning, (
"No warning should be issued"
)
@then("the cost metadata should record an exhaustion event with:")
def step_verify_exhaustion_event(context: object) -> None:
"""Verify exhaustion event was recorded."""
assert len(context.cost_metadata.budget_exhaustion_events) > 0
event = context.cost_metadata.budget_exhaustion_events[0]
for row in context.table:
field = row["field"]
expected = row["value"]
if field in ("limit", "used"):
expected = float(expected)
actual = getattr(event, field)
else:
actual = getattr(event, field)
assert actual == expected, f"Expected {field}={expected}, got {actual}"
@given("no session budget is configured")
def step_no_session_budget(context: object) -> None:
"""Configure no session budget."""
context.session_budget = None
context.cost_metadata = CostMetadata()
@when("multiple LLM calls totaling ${total} are made")
def step_multiple_calls(context: object, total: str) -> None:
"""Make multiple LLM calls."""
total_cost = float(total)
# Simulate 5 calls
per_call = total_cost / 5
for _ in range(5):
context.cost_service.record_cost(
context.cost_metadata,
context.session_budget,
None,
input_tokens=100,
output_tokens=50,
cost=per_call,
provider="openai",
@then("the session budget should have recorded the cost")
def step_session_budget_recorded(context):
"""Verify that the session budget recorded the cost."""
session_budget = getattr(context, 'session_budget', None)
if session_budget:
assert session_budget.total_cost > 0, (
"Session budget should have recorded the cost"
)
@then("all calls should succeed")
def step_verify_all_succeeded(context: object) -> None:
"""Verify all calls succeeded."""
# If we got here, no exception was raised
assert True
@then("the cost metadata should track the total cost")
def step_verify_total_tracked(context: object) -> None:
"""Verify total cost is tracked."""
assert context.cost_metadata.total_cost > 0
@when("a budget check is performed with a negative cost")
def step_check_negative_cost(context: object) -> None:
"""Check with negative cost."""
try:
context.cost_service.check_budget(context.session_budget, None, -1.0)
context.error_raised = False
except ValueError as e:
context.error_raised = True
context.error = e
@then("a ValueError should be raised")
def step_verify_value_error(context: object) -> None:
"""Verify ValueError was raised."""
assert context.error_raised, "Expected ValueError"
@when("recording cost with negative input tokens")
def step_record_negative_tokens(context: object) -> None:
"""Record cost with negative tokens."""
try:
context.cost_service.record_cost(
context.cost_metadata,
context.session_budget,
None,
input_tokens=-100,
output_tokens=50,
cost=1.0,
provider="openai",
@then("the organization budget should have recorded the cost")
def step_org_budget_recorded(context):
"""Verify that the organization budget recorded the cost."""
org_accumulator = getattr(context, 'org_accumulator', None)
if org_accumulator:
assert org_accumulator.total_cost > 0, (
"Organization budget should have recorded the cost"
)
@then("the budget should be exceeded")
def step_budget_exceeded(context):
"""Verify that the budget is exceeded."""
assert context.budget_exceeded, (
"Budget should be exceeded"
)
@then("the budget should not be exceeded")
def step_budget_not_exceeded(context):
"""Verify that the budget is not exceeded."""
assert not context.budget_exceeded, (
f"Budget should not be exceeded: {getattr(context, 'budget_error', '')}"
)
@given("the session budget utilization is at {utilization}%")
def step_set_session_utilization(context, utilization):
"""Set the session budget utilization."""
session_budget = getattr(context, 'session_budget', None)
if session_budget:
# Calculate the cost needed to reach the desired utilization
target_cost = session_budget.max_cost_usd * (float(utilization) / 100)
session_budget.record_cost(target_cost)
@given("the organization budget utilization is at {utilization}%")
def step_set_org_utilization(context, utilization):
"""Set the organization budget utilization."""
org_accumulator = getattr(context, 'org_accumulator', None)
if org_accumulator:
# Calculate the cost needed to reach the desired utilization
target_cost = org_accumulator.max_cost_usd * (float(utilization) / 100)
org_accumulator.record_cost(target_cost)
@then("the budget status should show session limit of ${limit}")
def step_check_session_limit(context, limit):
"""Verify the session budget limit."""
session_budget = getattr(context, 'session_budget', None)
if session_budget:
assert session_budget.max_cost_usd == float(limit), (
f"Expected limit {limit} but got {session_budget.max_cost_usd}"
)
@then("the budget status should show organization limit of ${limit}")
def step_check_org_limit(context, limit):
"""Verify the organization budget limit."""
org_accumulator = getattr(context, 'org_accumulator', None)
if org_accumulator:
assert org_accumulator.max_cost_usd == float(limit), (
f"Expected limit {limit} but got {org_accumulator.max_cost_usd}"
)
context.error_raised = False
except ValueError as e:
context.error_raised = True
context.error = e
+38 -37
View File
@@ -14,7 +14,6 @@ from cleveragents.core.exceptions import BudgetExceededError
from cleveragents.domain.models.core.cost_budget import (
BudgetCheckResult,
BudgetLevel,
OrgCostAccumulator,
SessionCostBudget,
ThreadSafeOrgCostAccumulator,
)
@@ -63,44 +62,46 @@ class CostTrackingService:
raise ValueError("proposed_cost must be non-negative")
# Check session budget
if session_budget is not None:
if session_budget.would_exceed(proposed_cost):
utilization = session_budget.utilization()
warning = (
utilization is not None and utilization >= 0.9
if not session_budget.is_exceeded()
else False
)
return BudgetCheckResult(
allowed=False,
exceeded_level=BudgetLevel.SESSION,
reason=(
f"Session budget of ${session_budget.max_cost_usd:.2f} "
f"would be exceeded (current: ${session_budget.total_cost:.2f}, "
f"proposed: ${proposed_cost:.2f})"
),
warning=warning,
)
if session_budget is not None and session_budget.would_exceed(
proposed_cost
):
utilization = session_budget.utilization()
warning = (
utilization is not None and utilization >= 0.9
if not session_budget.is_exceeded()
else False
)
return BudgetCheckResult(
allowed=False,
exceeded_level=BudgetLevel.SESSION,
reason=(
f"Session budget of ${session_budget.max_cost_usd:.2f} "
f"would be exceeded (current: ${session_budget.total_cost:.2f}, "
f"proposed: ${proposed_cost:.2f})"
),
warning=warning,
)
# Check org budget
if org_accumulator is not None:
if org_accumulator.would_exceed(proposed_cost):
utilization = org_accumulator.utilization()
warning = (
utilization is not None and utilization >= 0.9
if not org_accumulator.is_exceeded()
else False
)
return BudgetCheckResult(
allowed=False,
exceeded_level=BudgetLevel.ORG,
reason=(
f"Organization budget of ${org_accumulator.max_cost_usd:.2f} "
f"would be exceeded (current: ${org_accumulator.total_cost:.2f}, "
f"proposed: ${proposed_cost:.2f})"
),
warning=warning,
)
if org_accumulator is not None and org_accumulator.would_exceed(
proposed_cost
):
utilization = org_accumulator.utilization()
warning = (
utilization is not None and utilization >= 0.9
if not org_accumulator.is_exceeded()
else False
)
return BudgetCheckResult(
allowed=False,
exceeded_level=BudgetLevel.ORG,
reason=(
f"Organization budget of ${org_accumulator.max_cost_usd:.2f} "
f"would be exceeded (current: ${org_accumulator.total_cost:.2f}, "
f"proposed: ${proposed_cost:.2f})"
),
warning=warning,
)
# Check for warning thresholds
warning = False