"""ASV benchmarks for per-session and per-org cost budget checks. Measures the performance of: - SessionCostBudget creation and cost recording - OrgCostAccumulator creation and cost recording - CostBudgetService.check_budget_hierarchy with varying configurations - AutonomyGuardrailService.check_budget_hierarchy integration """ from __future__ import annotations import importlib import sys from pathlib import Path _SRC = str(Path(__file__).resolve().parents[1] / "src") if _SRC not in sys.path: sys.path.insert(0, _SRC) import cleveragents # noqa: E402 importlib.reload(cleveragents) from cleveragents.application.services.autonomy_guardrail_service import ( # noqa: E402 AutonomyGuardrailService, ) from cleveragents.application.services.cost_budget_service import ( # noqa: E402 CostBudgetService, ) from cleveragents.domain.models.core.autonomy_guardrails import ( # noqa: E402 AutonomyGuardrails, ) from cleveragents.domain.models.core.cost_budget import ( # noqa: E402 OrgCostAccumulator, SessionCostBudget, ThreadSafeOrgCostAccumulator, ) class SessionCostBudgetSuite: """Benchmark SessionCostBudget operations.""" def time_creation_defaults(self) -> None: """Benchmark creating a SessionCostBudget with defaults.""" SessionCostBudget() def time_creation_with_cap(self) -> None: """Benchmark creating a SessionCostBudget with a cap.""" SessionCostBudget(max_cost_usd=100.0) def time_record_cost(self) -> None: """Benchmark recording cost to a session budget.""" b = SessionCostBudget(max_cost_usd=1000.0) for _ in range(100): b.record_cost(1.0) def time_utilization(self) -> None: """Benchmark utilization calculation.""" b = SessionCostBudget(max_cost_usd=100.0, total_cost=50.0) for _ in range(1000): b.utilization() def time_would_exceed(self) -> None: """Benchmark would_exceed check.""" b = SessionCostBudget(max_cost_usd=100.0, total_cost=50.0) for _ in range(1000): b.would_exceed(10.0) class OrgCostAccumulatorSuite: """Benchmark OrgCostAccumulator operations.""" def time_creation(self) -> None: """Benchmark creating an OrgCostAccumulator.""" OrgCostAccumulator(org_id="bench-org", max_cost_usd=1000.0) def time_record_cost(self) -> None: """Benchmark recording cost to an org accumulator.""" acc = OrgCostAccumulator(org_id="bench-org", max_cost_usd=10000.0) for _ in range(100): acc.record_cost(1.0) def time_thread_safe_record(self) -> None: """Benchmark thread-safe cost recording.""" inner = OrgCostAccumulator(org_id="bench-org", max_cost_usd=10000.0) ts = ThreadSafeOrgCostAccumulator(inner) for _ in range(100): ts.record_cost(1.0) class CostBudgetServiceSuite: """Benchmark CostBudgetService hierarchy checks.""" def setup(self) -> None: """Configure service with session and org budgets.""" self._svc = CostBudgetService(warning_threshold=0.8) self._svc.configure_session_budget( "bench-sess", max_cost_usd=1000.0, org_id="bench-org" ) self._svc.configure_org_budget("bench-org", max_cost_usd=5000.0) def time_hierarchy_check_allowed(self) -> None: """Benchmark hierarchy check that is allowed.""" self._svc.check_budget_hierarchy( session_id="bench-sess", plan_cost=1.0, ) def time_hierarchy_check_with_plan_budget(self) -> None: """Benchmark hierarchy check with plan budget.""" self._svc.check_budget_hierarchy( session_id="bench-sess", plan_cost=1.0, plan_budget=100.0, plan_spent=10.0, ) def time_record_plan_cost(self) -> None: """Benchmark recording plan cost through the service.""" svc = CostBudgetService() svc.configure_session_budget("rs", max_cost_usd=100000.0) for _ in range(100): svc.record_plan_cost("rs", 1.0) class GuardrailBudgetIntegrationSuite: """Benchmark AutonomyGuardrailService budget hierarchy integration.""" def setup(self) -> None: """Configure guardrail service with budget hierarchy.""" self._cbs = CostBudgetService(warning_threshold=0.8) self._cbs.configure_session_budget( "g-sess", max_cost_usd=1000.0, org_id="g-org" ) self._cbs.configure_org_budget("g-org", max_cost_usd=5000.0) self._gs = AutonomyGuardrailService(cost_budget_service=self._cbs) self._gs.configure_guardrails("g-plan", AutonomyGuardrails(tool_budget=500.0)) self._gs.associate_plan_with_session("g-plan", "g-sess") def time_check_budget_hierarchy(self) -> None: """Benchmark guardrail service budget hierarchy check.""" self._gs.check_budget_hierarchy("g-plan", 1.0) def time_record_plan_cost_to_session(self) -> None: """Benchmark recording plan cost through guardrail service.""" cbs = CostBudgetService() cbs.configure_session_budget("rps", max_cost_usd=100000.0) gs = AutonomyGuardrailService(cost_budget_service=cbs) gs.configure_guardrails("rp", AutonomyGuardrails(tool_budget=100000.0)) gs.associate_plan_with_session("rp", "rps") for _ in range(100): gs.record_plan_cost_to_session("rp", 1.0)