diff --git a/CHANGELOG.md b/CHANGELOG.md index fdf6fab13..f2efc617e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,12 @@ ## Unreleased +- Added token/cost tracking, budget enforcement (per-plan and per-day), provider fallback + selection with capability filtering, and cost metadata for plan execution. New config keys + `budget_per_plan`, `budget_per_day`, and `fallback_providers` control spending limits and + provider ordering. Budget warnings are emitted at 90% usage, and requests are blocked at 100%. + Per-provider cost table includes default token cost estimates for offline reporting. Budget + exhaustion events are persisted in plan metadata for auditability. (#324) - Added plan-level and project-level advisory locking with configurable timeouts, re-entrant acquisition, conflict detection, lock renewal, graceful shutdown release, startup cleanup of expired locks, and diagnostics check for stale lock reporting. (#327) diff --git a/benchmarks/cost_controls_bench.py b/benchmarks/cost_controls_bench.py new file mode 100644 index 000000000..982ef2f3d --- /dev/null +++ b/benchmarks/cost_controls_bench.py @@ -0,0 +1,144 @@ +"""ASV benchmarks for cost controls overhead. + +Measures performance of: +- CostEntry creation and estimation +- ProviderCostTable lookups +- CostTracker budget checks and recording +- FallbackSelector provider selection +""" + +from __future__ import annotations + +import importlib +import sys +from pathlib import Path + +# Ensure the local source tree is importable +_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.config.settings import Settings # noqa: E402 +from cleveragents.domain.models.core.cost_metadata import CostMetadata # noqa: E402 +from cleveragents.providers.cost_table import ( # noqa: E402 + CostEntry, + ProviderCostTable, +) +from cleveragents.providers.cost_tracker import CostTracker # noqa: E402 +from cleveragents.providers.fallback_selector import FallbackSelector # noqa: E402 +from cleveragents.providers.registry import ProviderRegistry # noqa: E402 + + +class TimeCostEntry: + """Benchmark CostEntry creation and estimation.""" + + def time_create_cost_entry(self) -> None: + CostEntry(input_cost_per_token=0.00001, output_cost_per_token=0.00003) + + def time_estimate_cost(self) -> None: + entry = CostEntry(input_cost_per_token=0.00001, output_cost_per_token=0.00003) + entry.estimate_cost(1000, 500) + + +class TimeCostTable: + """Benchmark ProviderCostTable lookups.""" + + def setup(self) -> None: + self.table = ProviderCostTable() + + def time_lookup_known_model(self) -> None: + self.table.get_cost_entry("openai", "gpt-4o") + + def time_lookup_unknown_model(self) -> None: + self.table.get_cost_entry("openai", "unknown-model") + + def time_list_providers(self) -> None: + self.table.list_providers() + + def time_list_models(self) -> None: + self.table.list_models("openai") + + +class TimeCostTracker: + """Benchmark CostTracker budget checks and recording.""" + + def setup(self) -> None: + self.tracker = CostTracker(budget_per_plan=10.0, budget_per_day=50.0) + self.tracker_unlimited = CostTracker() + + def time_record_usage_with_budget(self) -> None: + metadata = CostMetadata() + self.tracker.record_usage( + metadata, + provider="openai", + model="gpt-4o", + input_tokens=1000, + output_tokens=500, + ) + + def time_record_usage_unlimited(self) -> None: + metadata = CostMetadata() + self.tracker_unlimited.record_usage( + metadata, + provider="openai", + model="gpt-4o", + input_tokens=1000, + output_tokens=500, + ) + + def time_check_plan_budget(self) -> None: + metadata = CostMetadata() + self.tracker.check_plan_budget(metadata) + + def time_check_daily_budget(self) -> None: + self.tracker.check_daily_budget() + + def time_estimate_cost(self) -> None: + self.tracker.estimate_cost("openai", "gpt-4o", 1000, 500) + + +class TimeFallbackSelector: + """Benchmark FallbackSelector provider selection.""" + + def setup(self) -> None: + settings = Settings() + self.registry = ProviderRegistry(settings) + self.selector = FallbackSelector(registry=self.registry) + + def time_select_no_requirements(self) -> None: + self.selector.select() + + def time_select_with_tool_calls(self) -> None: + self.selector.select(require_tool_calls=True) + + def time_select_with_all_requirements(self) -> None: + self.selector.select( + require_tool_calls=True, + require_streaming=True, + require_vision=True, + require_json_mode=True, + ) + + +class TimeCostMetadata: + """Benchmark CostMetadata operations.""" + + def time_create_metadata(self) -> None: + CostMetadata() + + def time_record_usage(self) -> None: + meta = CostMetadata() + meta.record_usage( + input_tokens=1000, + output_tokens=500, + cost=0.025, + provider="openai", + ) + + def time_display_dict(self) -> None: + meta = CostMetadata(total_tokens=1500, total_cost=0.025) + meta.as_display_dict() diff --git a/docs/reference/cost_controls.md b/docs/reference/cost_controls.md new file mode 100644 index 000000000..7a769af4a --- /dev/null +++ b/docs/reference/cost_controls.md @@ -0,0 +1,111 @@ +# Cost Controls + +CleverAgents provides cost tracking, budget enforcement, and provider +fallback controls for AI provider usage. + +## Configuration Keys + +| Key | Env Variable | Type | Default | Description | +|-----|-------------|------|---------|-------------| +| `budget_per_plan` | `CLEVERAGENTS_BUDGET_PER_PLAN` | `float \| None` | `None` | Max USD spend per plan execution. `None` = unlimited. | +| `budget_per_day` | `CLEVERAGENTS_BUDGET_PER_DAY` | `float \| None` | `None` | Max USD spend per calendar day. `None` = unlimited. | +| `fallback_providers` | `CLEVERAGENTS_FALLBACK_PROVIDERS` | `list[str]` | `[]` | Ordered provider fallback list. Empty uses default order. | + +## Budget Thresholds + +| Threshold | Value | Behavior | +|-----------|-------|----------| +| Warning | 90% of limit | Log warning, continue execution | +| Block | 100% of limit | Block further provider calls, persist exhaustion event | + +## Default Fallback Order + +When `fallback_providers` is empty, the default order is: + +1. OpenAI +2. Anthropic +3. Google +4. Azure +5. OpenRouter +6. Groq +7. Together +8. Cohere + +Providers are skipped if they: + +- Are not configured (no API key) +- Lack required capabilities (tool calling, streaming, vision, JSON mode) +- Would exceed the daily budget + +## Cost Tracking Fields + +Plan execution metadata includes: + +| Field | Type | Description | +|-------|------|-------------| +| `total_tokens` | `int` | Total tokens consumed (input + output) | +| `input_tokens` | `int` | Input/prompt tokens | +| `output_tokens` | `int` | Output/completion tokens | +| `total_cost` | `float` | Estimated cost in USD | +| `budget_remaining` | `float \| None` | Remaining plan budget (None if unlimited) | +| `provider_costs` | `dict[str, float]` | Per-provider cost breakdown | +| `budget_exhaustion_events` | `list` | Audit log of budget limit events | + +## Viewing Costs + +Use `agents plan status` to see cost data for a plan: + +``` +agents plan status +``` + +Cost data appears under the `cost` key in the output. + +## Per-Provider Cost Table + +Default token cost estimates are provided for offline reporting. +Models not in the table use a conservative default estimate. + +### Supported Providers + +- **OpenAI**: gpt-4o, gpt-4o-mini, gpt-4-turbo +- **Anthropic**: claude-sonnet-4-20250514, claude-3-5-haiku-20241022, claude-opus-4-20250514 +- **Google**: gemini-2.0-flash, gemini-1.5-pro +- **Groq**: llama-3.1-70b-versatile +- **Together**: meta-llama/Llama-3.1-70B-Instruct-Turbo +- **Cohere**: command-r-plus +- **Mock**: mock-gpt (zero cost) + +## Budget Exhaustion Events + +When a budget limit is reached, an exhaustion event is persisted to the +plan's cost metadata with: + +- Timestamp +- Budget type (`plan` or `daily`) +- Limit amount +- Used amount +- Provider and model that triggered it + +These events are retained for auditability and appear in `plan status` +output. + +## Examples + +### Set a per-plan budget + +```bash +export CLEVERAGENTS_BUDGET_PER_PLAN=5.00 +``` + +### Set a daily budget + +```bash +export CLEVERAGENTS_BUDGET_PER_DAY=50.00 +``` + +### Configure fallback providers + +```bash +export CLEVERAGENTS_FALLBACK_PROVIDERS='["anthropic", "openai", "groq"]' +``` diff --git a/features/cost_controls.feature b/features/cost_controls.feature new file mode 100644 index 000000000..3e82e5698 --- /dev/null +++ b/features/cost_controls.feature @@ -0,0 +1,463 @@ +Feature: Cost controls and provider fallback + As a developer using CleverAgents + I want cost tracking, budget enforcement, and provider fallback + So that I can control AI spending and ensure service continuity + + # ---- CostEntry ---- + + @unit @cost + Scenario: CostEntry validates non-negative input cost + Given I import the CostEntry class + When I create a CostEntry with negative input cost + Then a ValueError should be raised for CostEntry + + @unit @cost + Scenario: CostEntry validates non-negative output cost + Given I import the CostEntry class + When I create a CostEntry with negative output cost + Then a ValueError should be raised for CostEntry + + @unit @cost + Scenario: CostEntry estimates cost correctly + Given I import the CostEntry class + When I create a CostEntry with input cost 0.00001 and output cost 0.00003 + And I estimate cost for 1000 input tokens and 500 output tokens + Then the estimated cost should be 0.025 + + @unit @cost + Scenario: CostEntry rejects negative token counts + Given I import the CostEntry class + When I create a CostEntry with input cost 0.00001 and output cost 0.00003 + Then estimating cost with negative input tokens raises ValueError + And estimating cost with negative output tokens raises ValueError + + # ---- ProviderCostTable ---- + + @unit @cost + Scenario: ProviderCostTable returns default cost for unknown model + Given I create a ProviderCostTable with defaults + When I look up cost for provider "openai" model "unknown-model-xyz" + Then the cost entry should be the default cost + + @unit @cost + Scenario: ProviderCostTable returns known cost for gpt-4o + Given I create a ProviderCostTable with defaults + When I look up cost for provider "openai" model "gpt-4o" + Then the cost entry input cost should be 0.0000025 + And the cost entry output cost should be 0.00001 + + @unit @cost + Scenario: ProviderCostTable with custom entries overrides defaults + Given I create a ProviderCostTable with custom entry for "openai" "gpt-4o" + When I look up cost for provider "openai" model "gpt-4o" + Then the cost entry input cost should be 0.001 + And the cost entry output cost should be 0.002 + + @unit @cost + Scenario: ProviderCostTable validates empty provider + Given I create a ProviderCostTable with defaults + Then looking up cost for empty provider raises ValueError + + @unit @cost + Scenario: ProviderCostTable validates empty model + Given I create a ProviderCostTable with defaults + Then looking up cost for empty model raises ValueError + + @unit @cost + Scenario: ProviderCostTable lists providers and models + Given I create a ProviderCostTable with defaults + When I list providers from the cost table + Then the provider list should include "openai" + And the provider list should include "anthropic" + When I list models for provider "openai" + Then the model list should include "gpt-4o" + + @unit @cost + Scenario: ProviderCostTable rejects non-dict custom entries + When I create a ProviderCostTable with non-dict custom entries + Then a TypeError should be raised for ProviderCostTable + + # ---- CostMetadata ---- + + @unit @cost + Scenario: CostMetadata records usage correctly + Given I create a fresh CostMetadata + When I record usage of 100 input tokens and 50 output tokens at cost 0.01 for provider "openai" + Then total tokens should be 150 + And total cost should be 0.01 + And provider costs for "openai" should be 0.01 + + @unit @cost + Scenario: CostMetadata accumulates multiple usages + Given I create a fresh CostMetadata + When I record usage of 100 input tokens and 50 output tokens at cost 0.01 for provider "openai" + And I record usage of 200 input tokens and 100 output tokens at cost 0.02 for provider "anthropic" + Then total tokens should be 450 + And total cost should be 0.03 + And provider costs for "openai" should be 0.01 + And provider costs for "anthropic" should be 0.02 + + @unit @cost + Scenario: CostMetadata rejects negative input tokens + Given I create a fresh CostMetadata + Then recording usage with negative input tokens raises ValueError + + @unit @cost + Scenario: CostMetadata rejects negative cost + Given I create a fresh CostMetadata + Then recording usage with negative cost raises ValueError + + @unit @cost + Scenario: CostMetadata rejects empty provider + Given I create a fresh CostMetadata + Then recording usage with empty provider raises ValueError + + @unit @cost + Scenario: CostMetadata display dict includes all fields + Given I create a fresh CostMetadata + When I record usage of 100 input tokens and 50 output tokens at cost 0.01 for provider "openai" + And I get the display dict from cost metadata + Then the display dict should have key "total_tokens" + And the display dict should have key "total_cost_usd" + + @unit @cost + Scenario: CostMetadata display dict includes budget remaining when set + Given I create a fresh CostMetadata + When I set budget remaining to 5.0 + And I get the display dict from cost metadata + Then the display dict should have key "budget_remaining_usd" + + # ---- BudgetExhaustionEvent ---- + + @unit @cost + Scenario: BudgetExhaustionEvent validates budget type + Given I import the BudgetExhaustionEvent class + When I create a BudgetExhaustionEvent with invalid budget type + Then a validation error should be raised for BudgetExhaustionEvent + + @unit @cost + Scenario: BudgetExhaustionEvent accepts plan budget type + Given I import the BudgetExhaustionEvent class + When I create a BudgetExhaustionEvent with budget type "plan" + Then the event should be created successfully + + @unit @cost + Scenario: BudgetExhaustionEvent accepts daily budget type + Given I import the BudgetExhaustionEvent class + When I create a BudgetExhaustionEvent with budget type "daily" + Then the event should be created successfully + + # ---- CostTracker ---- + + @unit @cost + Scenario: CostTracker with no budget limits + Given I create a CostTracker with no limits + When I record tracked usage of 1000 input and 500 output for "openai" "gpt-4o" + Then the budget check should be under budget + And the cost metadata total tokens should be 1500 + + @unit @cost + Scenario: CostTracker warns at 90% of plan budget + Given I create a CostTracker with plan budget 0.001 + And I create a fresh CostMetadata + When I record tracked usage of 500 input and 500 output for "openai" "gpt-4o" + Then the budget check should be warning or exceeded + + @unit @cost + Scenario: CostTracker blocks at 100% of plan budget + Given I create a CostTracker with plan budget 0.0000001 + And I create a fresh CostMetadata + When I record tracked usage of 10000 input and 5000 output for "openai" "gpt-4o" + Then the budget check should be exceeded + + @unit @cost + Scenario: CostTracker blocks at 100% of daily budget + Given I create a CostTracker with daily budget 0.0000001 + And I create a fresh CostMetadata + When I record tracked usage of 10000 input and 5000 output for "openai" "gpt-4o" + Then the budget check should be exceeded + + @unit @cost + Scenario: CostTracker persists exhaustion events + Given I create a CostTracker with plan budget 0.0000001 + And I create a fresh CostMetadata + When I record tracked usage of 10000 input and 5000 output for "openai" "gpt-4o" + Then the cost metadata should have budget exhaustion events + + @unit @cost + Scenario: CostTracker validates negative plan budget + When I create a CostTracker with negative plan budget + Then a ValueError should be raised for CostTracker + + @unit @cost + Scenario: CostTracker validates negative daily budget + When I create a CostTracker with negative daily budget + Then a ValueError should be raised for CostTracker + + @unit @cost + Scenario: CostTracker estimate_cost delegates to cost table + Given I create a CostTracker with no limits + When I estimate cost for "openai" "gpt-4o" with 1000 input and 500 output + Then the estimated tracker cost should be positive + + @unit @cost + Scenario: CostTracker validates empty provider in estimate + Given I create a CostTracker with no limits + Then estimating cost with empty provider raises ValueError on tracker + + @unit @cost + Scenario: CostTracker validates empty model in estimate + Given I create a CostTracker with no limits + Then estimating cost with empty model raises ValueError on tracker + + @unit @cost + Scenario: CostTracker validates negative tokens in record + Given I create a CostTracker with no limits + And I create a fresh CostMetadata + Then recording tracked usage with negative input raises ValueError + + @unit @cost + Scenario: CostTracker check_plan_budget with no limit returns under budget + Given I create a CostTracker with no limits + And I create a fresh CostMetadata + When I check plan budget + Then the plan budget check should be under budget + + @unit @cost + Scenario: CostTracker check_daily_budget with no limit returns under budget + Given I create a CostTracker with no limits + When I check daily budget + Then the daily budget check should be under budget + + @unit @cost + Scenario: CostTracker get_daily_spend returns zero initially + Given I create a CostTracker with no limits + Then the daily spend should be 0.0 + + @unit @cost + Scenario: CostTracker budget properties are accessible + Given I create a CostTracker with plan budget 10.0 and daily budget 50.0 + Then the tracker plan budget should be 10.0 + And the tracker daily budget should be 50.0 + + @unit @cost + Scenario: CostTracker get_cost_entry validates empty provider + Given I create a CostTracker with no limits + Then getting cost entry with empty provider raises ValueError + + @unit @cost + Scenario: CostTracker get_cost_entry validates empty model + Given I create a CostTracker with no limits + Then getting cost entry with empty model raises ValueError + + @unit @cost + Scenario: CostTracker validates empty provider in record_usage + Given I create a CostTracker with no limits + And I create a fresh CostMetadata + Then recording tracked usage with empty provider raises ValueError + + @unit @cost + Scenario: CostTracker validates empty model in record_usage + Given I create a CostTracker with no limits + And I create a fresh CostMetadata + Then recording tracked usage with empty model raises ValueError + + @unit @cost + Scenario: CostTracker with zero plan budget returns exceeded + Given I create a CostTracker with plan budget 0.0 + And I create a fresh CostMetadata + When I check plan budget + Then the plan budget check should be exceeded + + @unit @cost + Scenario: CostTracker with zero daily budget returns exceeded + Given I create a CostTracker with daily budget 0.0 + When I check daily budget + Then the daily budget check should be exceeded + + # ---- FallbackSelector ---- + + @unit @cost + Scenario: FallbackSelector with no configured providers returns None + Given I create a ProviderRegistry with no API keys for fallback + When I create a FallbackSelector with the registry + And I select a fallback provider + Then the fallback result provider should be None + And the fallback result should have skipped entries + + @unit @cost + Scenario: FallbackSelector skips providers missing tool calls + Given I create a ProviderRegistry with mock provider for fallback + When I create a FallbackSelector with the registry + And I select a fallback provider requiring tool calls + Then the fallback result should skip mock for missing tool calls + + @unit @cost + Scenario: FallbackSelector selects first capable provider + Given I create a ProviderRegistry with openai key for fallback + When I create a FallbackSelector with the registry + And I select a fallback provider requiring tool calls + Then the fallback result provider type should be "openai" + + @unit @cost + Scenario: FallbackSelector uses custom fallback order + Given I create a ProviderRegistry with openai key for fallback + When I create a FallbackSelector with custom order "openai" + And I select a fallback provider + Then the fallback result provider type should be "openai" + + @unit @cost + Scenario: FallbackSelector rejects None registry + When I create a FallbackSelector with None registry + Then a TypeError should be raised for FallbackSelector + + @unit @cost + Scenario: FallbackSelector rejects empty string in fallback list + Given I create a ProviderRegistry with no API keys for fallback + When I create a FallbackSelector with empty string in fallback list + Then a ValueError should be raised for FallbackSelector + + @unit @cost + Scenario: FallbackSelector skips provider when daily budget exceeded + Given I create a ProviderRegistry with openai key for fallback + And I create a CostTracker with daily budget 0.0 + When I create a FallbackSelector with cost tracker + And I select a fallback provider + Then the fallback result provider should be None + + # ---- Settings config keys ---- + + @unit @cost + Scenario: Settings budget_per_plan defaults to None + When I load cost control settings with defaults + Then budget_per_plan should be None + + @unit @cost + Scenario: Settings budget_per_day defaults to None + When I load cost control settings with defaults + Then budget_per_day should be None + + @unit @cost + Scenario: Settings fallback_providers defaults to empty list + When I load cost control settings with defaults + Then fallback_providers should be an empty list + + @unit @cost + Scenario: Settings budget_per_plan accepts positive value + When I load cost control settings with budget_per_plan 10.0 + Then budget_per_plan should be 10.0 + + @unit @cost + Scenario: Settings budget_per_day accepts positive value + When I load cost control settings with budget_per_day 50.0 + Then budget_per_day should be 50.0 + + # ---- Plan model cost_metadata ---- + + @unit @cost + Scenario: Plan model includes cost_metadata field + Given I create a minimal Plan with cost metadata + Then the plan should have cost_metadata set + + @unit @cost + Scenario: Plan as_cli_dict includes cost when present + Given I create a minimal Plan with cost metadata + When I get the CLI dict from the plan + Then the CLI dict should have a "cost" key + + @unit @cost + Scenario: Plan as_cli_dict omits cost when None + Given I create a minimal Plan without cost metadata + When I get the CLI dict from the plan + Then the CLI dict should not have a "cost" key + + @unit @cost + Scenario: CostMetadata display dict with exhaustion events + Given I create a fresh CostMetadata + When I add a budget exhaustion event of type "plan" + And I get the display dict from cost metadata + Then the display dict should have key "budget_exhaustion_events" + + @unit @cost + Scenario: CostMetadata records negative output tokens rejection + Given I create a fresh CostMetadata + Then recording usage with negative output tokens raises ValueError + + # ---- Additional coverage scenarios ---- + + @unit @cost + Scenario: ProviderCostTable custom entries for brand new provider + Given I create a ProviderCostTable with custom entry for new provider "myprovider" model "mymodel" + When I look up cost for provider "myprovider" model "mymodel" + Then the cost entry input cost should be 0.005 + And the cost entry output cost should be 0.01 + + @unit @cost + Scenario: ProviderCostTable list_models with empty provider raises ValueError + Given I create a ProviderCostTable with defaults + Then listing models for empty provider raises ValueError + + @unit @cost + Scenario: CostTracker estimate_cost rejects negative input tokens + Given I create a CostTracker with no limits + Then estimating cost with negative input tokens raises ValueError on tracker + + @unit @cost + Scenario: CostTracker estimate_cost rejects negative output tokens + Given I create a CostTracker with no limits + Then estimating cost with negative output tokens raises ValueError on tracker + + @unit @cost + Scenario: CostTracker record_usage rejects negative output tokens + Given I create a CostTracker with no limits + And I create a fresh CostMetadata + Then recording tracked usage with negative output raises ValueError + + @unit @cost + Scenario: CostTracker plan budget warning at 90 percent + Given I create a CostTracker with plan budget 0.10 + And I create a fresh CostMetadata + When I record tracked usage of 3600 input and 8500 output for "openai" "gpt-4o" + Then the budget check should be warning + + @unit @cost + Scenario: CostTracker daily budget warning at 90 percent + Given I create a CostTracker with daily budget 0.10 + And I create a fresh CostMetadata + When I record tracked usage of 3600 input and 8500 output for "openai" "gpt-4o" + Then the budget check should be warning + + @unit @cost + Scenario: CostTracker get_cost_entry returns valid entry + Given I create a CostTracker with no limits + When I get cost entry for "openai" "gpt-4o" + Then the returned cost entry should have positive input cost + + @unit @cost + Scenario: FallbackSelector skips unknown provider types + Given I create a ProviderRegistry with no API keys for fallback + When I create a FallbackSelector with custom order "nonexistent_provider" + And I select a fallback provider + Then the fallback result provider should be None + And the fallback skipped should mention "unknown provider type" + + @unit @cost + Scenario: FallbackSelector skips providers missing streaming + Given I create a ProviderRegistry with mock provider for fallback + When I create a FallbackSelector with the registry + And I select a fallback provider requiring streaming + Then the fallback result should skip for missing streaming + + @unit @cost + Scenario: FallbackSelector skips providers missing vision + Given I create a ProviderRegistry with mock provider for fallback + When I create a FallbackSelector with the registry + And I select a fallback provider requiring vision + Then the fallback result should skip for missing vision + + @unit @cost + Scenario: FallbackSelector skips providers missing json mode + Given I create a ProviderRegistry with mock provider for fallback + When I create a FallbackSelector with the registry + And I select a fallback provider requiring json mode + Then the fallback result should skip for missing json mode diff --git a/features/steps/cost_controls_steps.py b/features/steps/cost_controls_steps.py new file mode 100644 index 000000000..1cc9c430c --- /dev/null +++ b/features/steps/cost_controls_steps.py @@ -0,0 +1,1065 @@ +"""Step definitions for cost controls feature tests.""" + +from __future__ import annotations + +import os +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 + + +# ---- 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" diff --git a/robot/cost_controls.robot b/robot/cost_controls.robot new file mode 100644 index 000000000..792e866b9 --- /dev/null +++ b/robot/cost_controls.robot @@ -0,0 +1,71 @@ +*** Settings *** +Resource ${CURDIR}/common.resource +Library OperatingSystem +Library Process + +*** Variables *** +${PYTHON} python +${SRC_DIR} ${CURDIR}/.. + +*** Test Cases *** +Cost Entry Creation And Estimation + [Documentation] Verify CostEntry creation and cost estimation + [Tags] cost_controls provider model + ${result}= Run Process ${PYTHON} robot/helper_cost_controls.py cost-entry + ... cwd=${SRC_DIR} + Log Process Failure ${result} + Should Be Equal As Integers ${result.rc} 0 + Should Contain ${result.stdout} cost-entry-ok + +Provider Cost Table Lookup + [Documentation] Verify ProviderCostTable model lookup + [Tags] cost_controls provider model + ${result}= Run Process ${PYTHON} robot/helper_cost_controls.py cost-table + ... cwd=${SRC_DIR} + Log Process Failure ${result} + Should Be Equal As Integers ${result.rc} 0 + Should Contain ${result.stdout} cost-table-ok + +Cost Tracker Budget Enforcement + [Documentation] Verify CostTracker budget enforcement + [Tags] cost_controls provider budget + ${result}= Run Process ${PYTHON} robot/helper_cost_controls.py cost-tracker + ... cwd=${SRC_DIR} + Log Process Failure ${result} + Should Be Equal As Integers ${result.rc} 0 + Should Contain ${result.stdout} cost-tracker-budget-ok + +Fallback Selector No Providers + [Documentation] Verify FallbackSelector with no configured providers + [Tags] cost_controls provider fallback + ${result}= Run Process ${PYTHON} robot/helper_cost_controls.py fallback-selector + ... cwd=${SRC_DIR} + Log Process Failure ${result} + Should Be Equal As Integers ${result.rc} 0 + Should Contain ${result.stdout} fallback-selector-ok + +Settings Cost Control Defaults + [Documentation] Verify Settings cost control config defaults + [Tags] cost_controls config settings + ${result}= Run Process ${PYTHON} robot/helper_cost_controls.py settings-defaults + ... cwd=${SRC_DIR} + Log Process Failure ${result} + Should Be Equal As Integers ${result.rc} 0 + Should Contain ${result.stdout} settings-defaults-ok + +Cost Metadata Display Dict + [Documentation] Verify CostMetadata display dictionary + [Tags] cost_controls provider model + ${result}= Run Process ${PYTHON} robot/helper_cost_controls.py cost-metadata + ... cwd=${SRC_DIR} + Log Process Failure ${result} + Should Be Equal As Integers ${result.rc} 0 + Should Contain ${result.stdout} cost-metadata-display-ok + +*** Keywords *** +Log Process Failure + [Arguments] ${result} + Run Keyword If ${result.rc} == 0 Return From Keyword + Log To Console Process failed with rc=${result.rc} + Log To Console STDOUT:${\n}${result.stdout} + Log To Console STDERR:${\n}${result.stderr} diff --git a/robot/helper_cost_controls.py b/robot/helper_cost_controls.py new file mode 100644 index 000000000..179871c78 --- /dev/null +++ b/robot/helper_cost_controls.py @@ -0,0 +1,146 @@ +#!/usr/bin/env python3 +"""Helper script for cost controls Robot Framework integration tests.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +# Ensure local source is importable +_SRC = str(Path(__file__).resolve().parents[1] / "src") +if _SRC not in sys.path: + sys.path.insert(0, _SRC) + + +def test_cost_entry_creation() -> None: + """Verify CostEntry creation and estimation.""" + from cleveragents.providers.cost_table import CostEntry + + entry = CostEntry(input_cost_per_token=0.00001, output_cost_per_token=0.00003) + cost = entry.estimate_cost(1000, 500) + assert cost > 0.0, f"Expected positive cost, got {cost}" + print("cost-entry-ok") + + +def test_cost_table_lookup() -> None: + """Verify ProviderCostTable lookup.""" + from cleveragents.providers.cost_table import ProviderCostTable + + table = ProviderCostTable() + entry = table.get_cost_entry("openai", "gpt-4o") + assert entry.input_cost_per_token > 0 + providers = table.list_providers() + assert "openai" in providers + print("cost-table-ok") + + +def test_cost_tracker_budget() -> None: + """Verify CostTracker budget enforcement.""" + from cleveragents.domain.models.core.cost_metadata import CostMetadata + from cleveragents.providers.cost_tracker import BudgetStatus, CostTracker + + tracker = CostTracker(budget_per_plan=0.0001) + metadata = CostMetadata() + result = tracker.record_usage( + metadata, + provider="openai", + model="gpt-4o", + input_tokens=10000, + output_tokens=5000, + ) + # With a tiny budget, usage should exceed it + assert result.status in (BudgetStatus.WARNING, BudgetStatus.EXCEEDED), ( + f"Expected warning/exceeded, got {result.status}" + ) + print("cost-tracker-budget-ok") + + +def test_fallback_selector() -> None: + """Verify FallbackSelector with no configured providers.""" + import os + + from cleveragents.config.settings import Settings + from cleveragents.providers.fallback_selector import FallbackSelector + from cleveragents.providers.registry import ProviderRegistry + + # Clear all API keys so no providers are configured + api_key_vars = [ + "OPENAI_API_KEY", + "ANTHROPIC_API_KEY", + "GOOGLE_API_KEY", + "AZURE_OPENAI_API_KEY", + "AZURE_API_KEY", + "OPENROUTER_API_KEY", + "GEMINI_API_KEY", + "GOOGLE_GEMINI_API_KEY", + "GROQ_API_KEY", + "TOGETHER_API_KEY", + "COHERE_API_KEY", + "PERPLEXITY_API_KEY", + ] + saved: dict[str, str] = {} + for var in api_key_vars: + val = os.environ.pop(var, None) + if val is not None: + saved[var] = val + try: + settings = Settings() + registry = ProviderRegistry(settings) + selector = FallbackSelector(registry=registry) + result = selector.select() + # With no API keys, no provider should be selected + assert result.provider_type is None + assert len(result.skipped) > 0 + print("fallback-selector-ok") + finally: + # Restore any keys we removed + os.environ.update(saved) + + +def test_settings_defaults() -> None: + """Verify Settings cost control defaults.""" + from cleveragents.config.settings import Settings + + settings = Settings() + assert settings.budget_per_plan is None + assert settings.budget_per_day is None + assert settings.fallback_providers == [] + print("settings-defaults-ok") + + +def test_cost_metadata_display() -> None: + """Verify CostMetadata display dict.""" + from cleveragents.domain.models.core.cost_metadata import CostMetadata + + meta = CostMetadata() + meta.record_usage( + input_tokens=100, + output_tokens=50, + cost=0.01, + provider="openai", + ) + display = meta.as_display_dict() + assert "total_tokens" in display + assert "total_cost_usd" in display + print("cost-metadata-display-ok") + + +if __name__ == "__main__": + if len(sys.argv) < 2: + print("Usage: helper_cost_controls.py ", file=sys.stderr) + sys.exit(1) + + test_name = sys.argv[1] + tests = { + "cost-entry": test_cost_entry_creation, + "cost-table": test_cost_table_lookup, + "cost-tracker": test_cost_tracker_budget, + "fallback-selector": test_fallback_selector, + "settings-defaults": test_settings_defaults, + "cost-metadata": test_cost_metadata_display, + } + test_fn = tests.get(test_name) + if test_fn is None: + print(f"Unknown test: {test_name}", file=sys.stderr) + sys.exit(1) + test_fn() diff --git a/src/cleveragents/config/settings.py b/src/cleveragents/config/settings.py index 4cf0c0a1d..d9e314a39 100644 --- a/src/cleveragents/config/settings.py +++ b/src/cleveragents/config/settings.py @@ -187,6 +187,34 @@ class Settings(BaseSettings): ), ) + # Cost controls (M4) + budget_per_plan: float | None = Field( + default=None, + ge=0.0, + validation_alias=AliasChoices("CLEVERAGENTS_BUDGET_PER_PLAN"), + description=( + "Maximum USD spend per plan execution. None (default) means unlimited." + ), + ) + budget_per_day: float | None = Field( + default=None, + ge=0.0, + validation_alias=AliasChoices("CLEVERAGENTS_BUDGET_PER_DAY"), + description=( + "Maximum USD spend per calendar day across all plans. " + "None (default) means unlimited." + ), + ) + fallback_providers: list[str] = Field( + default_factory=list, + validation_alias=AliasChoices("CLEVERAGENTS_FALLBACK_PROVIDERS"), + description=( + "Ordered list of provider names to try when the primary " + "provider is unavailable or over budget. " + "Empty list uses the default fallback order." + ), + ) + # Persistence database_url: str = Field( default="sqlite:///cleveragents.db", diff --git a/src/cleveragents/domain/models/core/cost_metadata.py b/src/cleveragents/domain/models/core/cost_metadata.py new file mode 100644 index 000000000..5409c7d26 --- /dev/null +++ b/src/cleveragents/domain/models/core/cost_metadata.py @@ -0,0 +1,162 @@ +"""Cost metadata model for plan execution tracking. + +Tracks token usage, monetary costs, and budget exhaustion events +attached to plan execution lifecycle. +""" + +from __future__ import annotations + +from datetime import datetime +from typing import Any + +from pydantic import BaseModel, ConfigDict, Field, field_validator + + +class BudgetExhaustionEvent(BaseModel): + """Record of a budget limit being reached during plan execution. + + Persisted in plan metadata for auditability (issue #324). + """ + + timestamp: datetime = Field( + default_factory=datetime.now, + description="When the budget threshold was reached", + ) + budget_type: str = Field( + ..., + min_length=1, + description="Type of budget exhausted: 'plan' or 'daily'", + ) + limit: float = Field( + ..., + ge=0.0, + description="The budget limit that was reached (USD)", + ) + used: float = Field( + ..., + ge=0.0, + description="Amount used at time of exhaustion (USD)", + ) + provider: str = Field( + default="", + description="Provider that triggered the exhaustion", + ) + model: str = Field( + default="", + description="Model that triggered the exhaustion", + ) + + @field_validator("budget_type") + @classmethod + def validate_budget_type(cls: type[BudgetExhaustionEvent], v: str) -> str: + allowed = {"plan", "daily"} + if v not in allowed: + raise ValueError(f"budget_type must be one of {allowed}, got {v!r}") + return v + + model_config = ConfigDict(validate_assignment=True) + + +class CostMetadata(BaseModel): + """Cost and token tracking metadata for a plan execution. + + Attached to plan execution to provide cost visibility and + auditability. Surfaced in ``plan status`` CLI output. + """ + + total_tokens: int = Field( + default=0, + ge=0, + description="Total tokens consumed (input + output)", + ) + input_tokens: int = Field( + default=0, + ge=0, + description="Input/prompt tokens consumed", + ) + output_tokens: int = Field( + default=0, + ge=0, + description="Output/completion tokens consumed", + ) + total_cost: float = Field( + default=0.0, + ge=0.0, + description="Total estimated cost in USD", + ) + budget_remaining: float | None = Field( + default=None, + description="Remaining budget in USD (None if unlimited)", + ) + provider_costs: dict[str, float] = Field( + default_factory=dict, + description="Cost breakdown by provider name", + ) + budget_exhaustion_events: list[BudgetExhaustionEvent] = Field( + default_factory=list, + description="Audit log of budget limit events", + ) + + def record_usage( + self, + *, + input_tokens: int, + output_tokens: int, + cost: float, + provider: str, + ) -> None: + """Record token usage and cost from a provider call. + + Args: + input_tokens: Number of input tokens consumed. + output_tokens: Number of output tokens consumed. + cost: Estimated cost in USD. + provider: Provider name for cost breakdown. + + Raises: + ValueError: If any numeric argument is negative. + """ + if input_tokens < 0: + raise ValueError(f"input_tokens must be >= 0, got {input_tokens}") + if output_tokens < 0: + raise ValueError(f"output_tokens must be >= 0, got {output_tokens}") + if cost < 0.0: + raise ValueError(f"cost must be >= 0.0, got {cost}") + if not provider: + raise ValueError("provider must be a non-empty string") + + self.input_tokens += input_tokens + self.output_tokens += output_tokens + self.total_tokens += input_tokens + output_tokens + self.total_cost += cost + self.provider_costs[provider] = self.provider_costs.get(provider, 0.0) + cost + + def as_display_dict(self) -> dict[str, Any]: + """Return a dictionary suitable for CLI display.""" + result: dict[str, Any] = { + "total_tokens": self.total_tokens, + "input_tokens": self.input_tokens, + "output_tokens": self.output_tokens, + "total_cost_usd": round(self.total_cost, 6), + } + if self.budget_remaining is not None: + result["budget_remaining_usd"] = round(self.budget_remaining, 6) + if self.provider_costs: + result["provider_costs"] = { + k: round(v, 6) for k, v in self.provider_costs.items() + } + if self.budget_exhaustion_events: + result["budget_exhaustion_events"] = [ + { + "timestamp": evt.timestamp.isoformat(), + "budget_type": evt.budget_type, + "limit": evt.limit, + "used": evt.used, + "provider": evt.provider, + "model": evt.model, + } + for evt in self.budget_exhaustion_events + ] + return result + + model_config = ConfigDict(validate_assignment=True) diff --git a/src/cleveragents/domain/models/core/plan.py b/src/cleveragents/domain/models/core/plan.py index 5b54f2b51..757704c1b 100644 --- a/src/cleveragents/domain/models/core/plan.py +++ b/src/cleveragents/domain/models/core/plan.py @@ -63,6 +63,8 @@ from typing import Any, ClassVar from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator +from cleveragents.domain.models.core.cost_metadata import CostMetadata + # ULID is 26 characters, all uppercase alphanumeric (Crockford's base32) ULID_PATTERN = r"^[0-9A-HJKMNP-TV-Z]{26}$" @@ -640,6 +642,12 @@ class Plan(BaseModel): description="ULID of the root decision node in the decision tree", ) + # Cost tracking (M4) + cost_metadata: CostMetadata | None = Field( + default=None, + description="Token/cost tracking and budget exhaustion events", + ) + # Metadata created_by: str | None = Field(None, description="User/session that created plan") tags: list[str] = Field(default_factory=list, description="Tags for organization") @@ -861,6 +869,8 @@ class Plan(BaseModel): result["parent_plan_id"] = self.identity.parent_plan_id if self.has_subplans: result["subplan_count"] = len(self.subplan_statuses) + if self.cost_metadata is not None: + result["cost"] = self.cost_metadata.as_display_dict() return result model_config = ConfigDict( diff --git a/src/cleveragents/providers/__init__.py b/src/cleveragents/providers/__init__.py index 1cd4e7499..2ab0c0b14 100644 --- a/src/cleveragents/providers/__init__.py +++ b/src/cleveragents/providers/__init__.py @@ -1,3 +1,6 @@ +from .cost_table import CostEntry, ProviderCostTable +from .cost_tracker import BudgetCheckResult, BudgetStatus, CostTracker +from .fallback_selector import FallbackResult, FallbackSelector from .registry import ( ProviderRegistry, get_provider_registry, @@ -5,6 +8,13 @@ from .registry import ( ) __all__ = [ + "BudgetCheckResult", + "BudgetStatus", + "CostEntry", + "CostTracker", + "FallbackResult", + "FallbackSelector", + "ProviderCostTable", "ProviderRegistry", "get_provider_registry", "reset_provider_registry", diff --git a/src/cleveragents/providers/cost_table.py b/src/cleveragents/providers/cost_table.py new file mode 100644 index 000000000..bb23b482f --- /dev/null +++ b/src/cleveragents/providers/cost_table.py @@ -0,0 +1,216 @@ +"""Per-provider cost table with default token cost estimates. + +Provides offline cost estimation for token usage across supported +providers and models. Costs are in USD per token. + +Default estimates are approximations for common models and should be +updated as provider pricing changes. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import ClassVar + + +@dataclass(frozen=True, slots=True) +class CostEntry: + """Cost per token for a specific provider/model combination. + + Attributes: + input_cost_per_token: USD cost per input/prompt token. + output_cost_per_token: USD cost per output/completion token. + """ + + input_cost_per_token: float + output_cost_per_token: float + + def __post_init__(self) -> None: + if self.input_cost_per_token < 0.0: + raise ValueError( + f"input_cost_per_token must be >= 0, got {self.input_cost_per_token}" + ) + if self.output_cost_per_token < 0.0: + raise ValueError( + f"output_cost_per_token must be >= 0, got {self.output_cost_per_token}" + ) + + def estimate_cost(self, input_tokens: int, output_tokens: int) -> float: + """Estimate the cost for a given token count. + + Args: + input_tokens: Number of input tokens. + output_tokens: Number of output tokens. + + Returns: + Estimated cost in USD. + + Raises: + ValueError: If token counts are negative. + """ + if input_tokens < 0: + raise ValueError(f"input_tokens must be >= 0, got {input_tokens}") + if output_tokens < 0: + raise ValueError(f"output_tokens must be >= 0, got {output_tokens}") + return ( + self.input_cost_per_token * input_tokens + + self.output_cost_per_token * output_tokens + ) + + +class ProviderCostTable: + """Lookup table for provider/model cost estimates. + + Maintains a registry of known cost entries and provides a fallback + default for unknown models. Thread-safe for read operations. + + Example:: + + table = ProviderCostTable() + entry = table.get_cost_entry("openai", "gpt-4o") + cost = entry.estimate_cost(input_tokens=1000, output_tokens=500) + """ + + # Default fallback cost for unknown models (conservative estimate) + DEFAULT_COST: ClassVar[CostEntry] = CostEntry( + input_cost_per_token=0.000005, + output_cost_per_token=0.000015, + ) + + # Default cost entries for known provider/model combinations. + # Costs are approximate and based on published pricing as of 2025. + _DEFAULT_ENTRIES: ClassVar[dict[str, dict[str, CostEntry]]] = { + "openai": { + "gpt-4o": CostEntry( + input_cost_per_token=0.0000025, + output_cost_per_token=0.00001, + ), + "gpt-4o-mini": CostEntry( + input_cost_per_token=0.00000015, + output_cost_per_token=0.0000006, + ), + "gpt-4-turbo": CostEntry( + input_cost_per_token=0.00001, + output_cost_per_token=0.00003, + ), + }, + "anthropic": { + "claude-sonnet-4-20250514": CostEntry( + input_cost_per_token=0.000003, + output_cost_per_token=0.000015, + ), + "claude-3-5-haiku-20241022": CostEntry( + input_cost_per_token=0.0000008, + output_cost_per_token=0.000004, + ), + "claude-opus-4-20250514": CostEntry( + input_cost_per_token=0.000015, + output_cost_per_token=0.000075, + ), + }, + "google": { + "gemini-2.0-flash": CostEntry( + input_cost_per_token=0.0000001, + output_cost_per_token=0.0000004, + ), + "gemini-1.5-pro": CostEntry( + input_cost_per_token=0.00000125, + output_cost_per_token=0.000005, + ), + }, + "groq": { + "llama-3.1-70b-versatile": CostEntry( + input_cost_per_token=0.00000059, + output_cost_per_token=0.00000079, + ), + }, + "together": { + "meta-llama/Llama-3.1-70B-Instruct-Turbo": CostEntry( + input_cost_per_token=0.00000088, + output_cost_per_token=0.00000088, + ), + }, + "cohere": { + "command-r-plus": CostEntry( + input_cost_per_token=0.0000025, + output_cost_per_token=0.00001, + ), + }, + "mock": { + "mock-gpt": CostEntry( + input_cost_per_token=0.0, + output_cost_per_token=0.0, + ), + }, + } + + def __init__( + self, + custom_entries: dict[str, dict[str, CostEntry]] | None = None, + ) -> None: + """Initialize the cost table. + + Args: + custom_entries: Optional custom cost entries to merge with + defaults. Custom entries take precedence. + + Raises: + TypeError: If custom_entries is not a dict. + """ + if custom_entries is not None and not isinstance(custom_entries, dict): + raise TypeError( + f"custom_entries must be a dict, got {type(custom_entries)}" + ) + self._entries: dict[str, dict[str, CostEntry]] = {} + for provider, models in self._DEFAULT_ENTRIES.items(): + self._entries[provider] = dict(models) + if custom_entries: + for provider, models in custom_entries.items(): + if provider not in self._entries: + self._entries[provider] = {} + self._entries[provider].update(models) + + def get_cost_entry(self, provider: str, model: str) -> CostEntry: + """Look up the cost entry for a provider/model pair. + + Args: + provider: Provider name (e.g. 'openai', 'anthropic'). + model: Model identifier (e.g. 'gpt-4o'). + + Returns: + The matching CostEntry, or DEFAULT_COST if not found. + + Raises: + ValueError: If provider or model is empty. + """ + if not provider: + raise ValueError("provider must be a non-empty string") + if not model: + raise ValueError("model must be a non-empty string") + provider_lower = provider.strip().lower() + model_lower = model.strip().lower() + provider_models = self._entries.get(provider_lower, {}) + for key, entry in provider_models.items(): + if key.lower() == model_lower: + return entry + return self.DEFAULT_COST + + def list_providers(self) -> list[str]: + """Return sorted list of providers with cost data.""" + return sorted(self._entries.keys()) + + def list_models(self, provider: str) -> list[str]: + """Return sorted list of models for a provider. + + Args: + provider: Provider name. + + Returns: + List of model names, empty if provider unknown. + + Raises: + ValueError: If provider is empty. + """ + if not provider: + raise ValueError("provider must be a non-empty string") + return sorted(self._entries.get(provider.lower(), {}).keys()) diff --git a/src/cleveragents/providers/cost_tracker.py b/src/cleveragents/providers/cost_tracker.py new file mode 100644 index 000000000..2443e1046 --- /dev/null +++ b/src/cleveragents/providers/cost_tracker.py @@ -0,0 +1,367 @@ +"""Cost tracking and budget enforcement for provider usage. + +Tracks token consumption and monetary costs per plan execution, +enforces budget limits, and emits warnings/blocks when thresholds +are approached or exceeded. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from datetime import date, datetime +from enum import StrEnum +from typing import ClassVar + +from cleveragents.domain.models.core.cost_metadata import ( + BudgetExhaustionEvent, + CostMetadata, +) +from cleveragents.providers.cost_table import CostEntry, ProviderCostTable + +logger = logging.getLogger(__name__) + + +class BudgetStatus(StrEnum): + """Budget enforcement status levels.""" + + UNDER_BUDGET = "under_budget" + WARNING = "warning" + EXCEEDED = "exceeded" + + +# Severity ordering for BudgetStatus (higher = worse) +_BUDGET_SEVERITY: dict[BudgetStatus, int] = { + BudgetStatus.UNDER_BUDGET: 0, + BudgetStatus.WARNING: 1, + BudgetStatus.EXCEEDED: 2, +} + + +@dataclass(frozen=True, slots=True) +class BudgetCheckResult: + """Result of a budget check operation. + + Attributes: + status: Current budget status. + used: Amount used so far (USD). + limit: Budget limit (USD). None if unlimited. + remaining: Amount remaining (USD). None if unlimited. + budget_type: Which budget was checked ('plan' or 'daily'). + """ + + status: BudgetStatus + used: float + limit: float | None + remaining: float | None + budget_type: str + + +class CostTracker: + """Track costs and enforce budget limits for provider usage. + + Manages per-plan and per-day budget tracking with configurable + warning and blocking thresholds. Uses dependency injection for + the cost table. + + Attributes: + WARN_THRESHOLD: Fraction of budget at which warnings are emitted. + BLOCK_THRESHOLD: Fraction of budget at which requests are blocked. + """ + + WARN_THRESHOLD: ClassVar[float] = 0.9 + BLOCK_THRESHOLD: ClassVar[float] = 1.0 + + def __init__( + self, + *, + budget_per_plan: float | None = None, + budget_per_day: float | None = None, + cost_table: ProviderCostTable | None = None, + ) -> None: + """Initialize the cost tracker. + + Args: + budget_per_plan: Maximum USD spend per plan (None = unlimited). + budget_per_day: Maximum USD spend per day (None = unlimited). + cost_table: Provider cost lookup table (uses defaults if None). + + Raises: + ValueError: If budget values are negative. + """ + if budget_per_plan is not None and budget_per_plan < 0.0: + raise ValueError(f"budget_per_plan must be >= 0, got {budget_per_plan}") + if budget_per_day is not None and budget_per_day < 0.0: + raise ValueError(f"budget_per_day must be >= 0, got {budget_per_day}") + self._budget_per_plan = budget_per_plan + self._budget_per_day = budget_per_day + self._cost_table = cost_table or ProviderCostTable() + self._daily_costs: dict[str, float] = {} + + @property + def budget_per_plan(self) -> float | None: + """Return the per-plan budget limit.""" + return self._budget_per_plan + + @property + def budget_per_day(self) -> float | None: + """Return the per-day budget limit.""" + return self._budget_per_day + + def get_cost_entry(self, provider: str, model: str) -> CostEntry: + """Look up the cost entry for a provider/model. + + Args: + provider: Provider name. + model: Model identifier. + + Returns: + The cost entry for the provider/model pair. + + Raises: + ValueError: If provider or model is empty. + """ + if not provider: + raise ValueError("provider must be a non-empty string") + if not model: + raise ValueError("model must be a non-empty string") + return self._cost_table.get_cost_entry(provider, model) + + def estimate_cost( + self, + provider: str, + model: str, + input_tokens: int, + output_tokens: int, + ) -> float: + """Estimate cost for a token usage. + + Args: + provider: Provider name. + model: Model identifier. + input_tokens: Number of input tokens. + output_tokens: Number of output tokens. + + Returns: + Estimated cost in USD. + + Raises: + ValueError: If arguments are invalid. + """ + if not provider: + raise ValueError("provider must be a non-empty string") + if not model: + raise ValueError("model must be a non-empty string") + if input_tokens < 0: + raise ValueError(f"input_tokens must be >= 0, got {input_tokens}") + if output_tokens < 0: + raise ValueError(f"output_tokens must be >= 0, got {output_tokens}") + entry = self._cost_table.get_cost_entry(provider, model) + return entry.estimate_cost(input_tokens, output_tokens) + + def record_usage( + self, + cost_metadata: CostMetadata, + *, + provider: str, + model: str, + input_tokens: int, + output_tokens: int, + ) -> BudgetCheckResult: + """Record usage and check budget status. + + Updates cost_metadata in-place with token/cost data. + Also updates internal daily cost tracking. + Records budget exhaustion events when thresholds are crossed. + + Args: + cost_metadata: The plan's cost metadata to update. + provider: Provider name. + model: Model identifier. + input_tokens: Input tokens used. + output_tokens: Output tokens used. + + Returns: + Budget check result after recording usage. + + Raises: + ValueError: If arguments are invalid. + """ + if not provider: + raise ValueError("provider must be a non-empty string") + if not model: + raise ValueError("model must be a non-empty string") + if input_tokens < 0: + raise ValueError(f"input_tokens must be >= 0, got {input_tokens}") + if output_tokens < 0: + raise ValueError(f"output_tokens must be >= 0, got {output_tokens}") + + cost = self.estimate_cost(provider, model, input_tokens, output_tokens) + cost_metadata.record_usage( + input_tokens=input_tokens, + output_tokens=output_tokens, + cost=cost, + provider=provider, + ) + + today_key = date.today().isoformat() + self._daily_costs[today_key] = self._daily_costs.get(today_key, 0.0) + cost + + if self._budget_per_plan is not None: + cost_metadata.budget_remaining = max( + 0.0, self._budget_per_plan - cost_metadata.total_cost + ) + + result = self._check_budgets(cost_metadata, provider, model) + return result + + def check_plan_budget(self, cost_metadata: CostMetadata) -> BudgetCheckResult: + """Check the per-plan budget without recording usage. + + Args: + cost_metadata: The plan's cost metadata. + + Returns: + Budget check result for the plan budget. + """ + return self._evaluate_budget( + used=cost_metadata.total_cost, + limit=self._budget_per_plan, + budget_type="plan", + ) + + def check_daily_budget(self) -> BudgetCheckResult: + """Check the per-day budget without recording usage. + + Returns: + Budget check result for the daily budget. + """ + today_key = date.today().isoformat() + used = self._daily_costs.get(today_key, 0.0) + return self._evaluate_budget( + used=used, + limit=self._budget_per_day, + budget_type="daily", + ) + + def get_daily_spend(self) -> float: + """Return total spend for the current day.""" + today_key = date.today().isoformat() + return self._daily_costs.get(today_key, 0.0) + + def _check_budgets( + self, + cost_metadata: CostMetadata, + provider: str, + model: str, + ) -> BudgetCheckResult: + """Check both plan and daily budgets, return worst status.""" + plan_result = self.check_plan_budget(cost_metadata) + daily_result = self.check_daily_budget() + + worst = plan_result + if _BUDGET_SEVERITY.get(daily_result.status, 0) > _BUDGET_SEVERITY.get( + worst.status, 0 + ): + worst = daily_result + + if plan_result.status == BudgetStatus.WARNING: + logger.warning( + "Plan budget warning: %.2f%% used (%.4f / %.4f USD)", + (plan_result.used / plan_result.limit * 100) + if plan_result.limit + else 0, + plan_result.used, + plan_result.limit or 0, + ) + if daily_result.status == BudgetStatus.WARNING: + logger.warning( + "Daily budget warning: %.2f%% used (%.4f / %.4f USD)", + (daily_result.used / daily_result.limit * 100) + if daily_result.limit + else 0, + daily_result.used, + daily_result.limit or 0, + ) + + if plan_result.status == BudgetStatus.EXCEEDED: + self._record_exhaustion_event(cost_metadata, "plan", provider, model) + if daily_result.status == BudgetStatus.EXCEEDED: + self._record_exhaustion_event(cost_metadata, "daily", provider, model) + + return worst + + def _evaluate_budget( + self, + *, + used: float, + limit: float | None, + budget_type: str, + ) -> BudgetCheckResult: + """Evaluate budget status for a given usage and limit.""" + if limit is None: + return BudgetCheckResult( + status=BudgetStatus.UNDER_BUDGET, + used=used, + limit=None, + remaining=None, + budget_type=budget_type, + ) + + if limit <= 0.0: + return BudgetCheckResult( + status=BudgetStatus.EXCEEDED, + used=used, + limit=limit, + remaining=0.0, + budget_type=budget_type, + ) + + ratio = used / limit + remaining = max(0.0, limit - used) + + if ratio >= self.BLOCK_THRESHOLD: + status = BudgetStatus.EXCEEDED + elif ratio >= self.WARN_THRESHOLD: + status = BudgetStatus.WARNING + else: + status = BudgetStatus.UNDER_BUDGET + + return BudgetCheckResult( + status=status, + used=used, + limit=limit, + remaining=remaining, + budget_type=budget_type, + ) + + def _record_exhaustion_event( + self, + cost_metadata: CostMetadata, + budget_type: str, + provider: str, + model: str, + ) -> None: + """Persist a budget exhaustion event to cost metadata.""" + limit = self._budget_per_plan if budget_type == "plan" else self._budget_per_day + if limit is None: + return + + event = BudgetExhaustionEvent( + timestamp=datetime.now(), + budget_type=budget_type, + limit=limit, + used=cost_metadata.total_cost, + provider=provider, + model=model, + ) + cost_metadata.budget_exhaustion_events.append(event) + logger.warning( + "Budget exhausted: %s budget %.4f USD exceeded " + "(used: %.4f, provider: %s, model: %s)", + budget_type, + limit, + cost_metadata.total_cost, + provider, + model, + ) diff --git a/src/cleveragents/providers/fallback_selector.py b/src/cleveragents/providers/fallback_selector.py new file mode 100644 index 000000000..ea177f248 --- /dev/null +++ b/src/cleveragents/providers/fallback_selector.py @@ -0,0 +1,206 @@ +"""Provider fallback selection logic. + +Selects providers based on required capabilities and budget availability, +skipping providers that lack tool calling, streaming, or other features. +""" + +from __future__ import annotations + +import logging +from dataclasses import dataclass +from typing import ClassVar + +from cleveragents.providers.cost_tracker import BudgetStatus, CostTracker +from cleveragents.providers.registry import ( + ProviderCapabilities, + ProviderInfo, + ProviderRegistry, + ProviderType, +) + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True, slots=True) +class FallbackResult: + """Result of a fallback provider selection. + + Attributes: + provider_type: Selected provider type, or None if no suitable + provider found. + provider_info: Provider info for the selected provider. + skipped: List of (provider_type, reason) pairs for skipped providers. + """ + + provider_type: ProviderType | None + provider_info: ProviderInfo | None + skipped: list[tuple[str, str]] + + +class FallbackSelector: + """Select providers using fallback logic with capability filtering. + + Iterates through configured providers (in user-defined or default + priority order) and selects the first that meets all required + capabilities and is within budget. + + Example:: + + selector = FallbackSelector( + registry=registry, + cost_tracker=cost_tracker, + fallback_providers=["openai", "anthropic", "groq"], + ) + result = selector.select( + require_tool_calls=True, + require_streaming=True, + ) + if result.provider_type: + # Use result.provider_type + ... + """ + + DEFAULT_FALLBACK_ORDER: ClassVar[list[str]] = [ + "openai", + "anthropic", + "google", + "azure", + "openrouter", + "groq", + "together", + "cohere", + ] + + def __init__( + self, + *, + registry: ProviderRegistry, + cost_tracker: CostTracker | None = None, + fallback_providers: list[str] | None = None, + ) -> None: + """Initialize the fallback selector. + + Args: + registry: Provider registry for capability lookups. + cost_tracker: Optional cost tracker for budget checks. + fallback_providers: Ordered list of provider names to try. + Uses DEFAULT_FALLBACK_ORDER if not specified. + + Raises: + TypeError: If registry is None. + ValueError: If fallback_providers contains empty strings. + """ + if registry is None: + raise TypeError("registry must not be None") + if fallback_providers is not None: + for idx, name in enumerate(fallback_providers): + if not name or not name.strip(): + raise ValueError( + f"fallback_providers[{idx}] must be a non-empty string" + ) + self._registry = registry + self._cost_tracker = cost_tracker + self._fallback_order = ( + [p.strip().lower() for p in fallback_providers] + if fallback_providers + else list(self.DEFAULT_FALLBACK_ORDER) + ) + + def select( + self, + *, + require_tool_calls: bool = False, + require_streaming: bool = False, + require_vision: bool = False, + require_json_mode: bool = False, + ) -> FallbackResult: + """Select the best available provider matching requirements. + + Iterates through the fallback order, skipping providers that: + - Are not configured (no API key) + - Lack required capabilities + - Would exceed budget limits + + Args: + require_tool_calls: Require tool/function calling support. + require_streaming: Require streaming response support. + require_vision: Require image/vision input support. + require_json_mode: Require structured JSON output. + + Returns: + FallbackResult with the selected provider or None. + """ + skipped: list[tuple[str, str]] = [] + + for provider_name in self._fallback_order: + info = self._registry.get_provider_info(provider_name) + if info is None: + skipped.append((provider_name, "unknown provider type")) + continue + + if not info.is_configured: + skipped.append((provider_name, "not configured (no API key)")) + continue + + caps = info.capabilities + skip_reason = self._check_capabilities( + caps, + require_tool_calls=require_tool_calls, + require_streaming=require_streaming, + require_vision=require_vision, + require_json_mode=require_json_mode, + ) + if skip_reason: + skipped.append((provider_name, skip_reason)) + continue + + if self._cost_tracker is not None: + daily_check = self._cost_tracker.check_daily_budget() + if daily_check.status == BudgetStatus.EXCEEDED: + skipped.append((provider_name, "daily budget exceeded")) + continue + + logger.info( + "Fallback selected provider: %s (skipped %d)", + provider_name, + len(skipped), + ) + return FallbackResult( + provider_type=info.provider_type, + provider_info=info, + skipped=skipped, + ) + + logger.warning( + "No suitable fallback provider found (tried %d providers)", + len(self._fallback_order), + ) + return FallbackResult( + provider_type=None, + provider_info=None, + skipped=skipped, + ) + + def _check_capabilities( + self, + caps: ProviderCapabilities, + *, + require_tool_calls: bool, + require_streaming: bool, + require_vision: bool, + require_json_mode: bool, + ) -> str | None: + """Check if capabilities meet requirements. + + Returns: + Reason string if capabilities are insufficient, None if OK. + """ + if require_tool_calls and not caps.supports_tool_calls: + return "missing capability: tool_calls" + if require_streaming and not caps.supports_streaming: + return "missing capability: streaming" + if require_vision and not caps.supports_vision: + return "missing capability: vision" + if require_json_mode and not caps.supports_json_mode: + return "missing capability: json_mode" + return None diff --git a/vulture_whitelist.py b/vulture_whitelist.py index 20972872b..3af437e66 100644 --- a/vulture_whitelist.py +++ b/vulture_whitelist.py @@ -294,3 +294,32 @@ LockModel # noqa: B018, F821 LockConflictError # noqa: B018, F821 LockExpiredError # noqa: B018, F821 _check_stale_locks # noqa: B018, F821 + +# Cost controls (M4) — public API surface used by CLI, tests, and benchmarks +CostEntry # noqa: B018, F821 +ProviderCostTable # noqa: B018, F821 +CostTracker # noqa: B018, F821 +BudgetStatus # noqa: B018, F821 +BudgetCheckResult # noqa: B018, F821 +FallbackSelector # noqa: B018, F821 +FallbackResult # noqa: B018, F821 +CostMetadata # noqa: B018, F821 +BudgetExhaustionEvent # noqa: B018, F821 +budget_per_plan # noqa: B018, F821 +budget_per_day # noqa: B018, F821 +fallback_providers # noqa: B018, F821 +cost_metadata # noqa: B018, F821 +budget_remaining # noqa: B018, F821 +input_tokens # noqa: B018, F821 +output_tokens # noqa: B018, F821 +provider_costs # noqa: B018, F821 +budget_exhaustion_events # noqa: B018, F821 +as_display_dict # noqa: B018, F821 +record_usage # noqa: B018, F821 +estimate_cost # noqa: B018, F821 +get_cost_entry # noqa: B018, F821 +list_providers # noqa: B018, F821 +list_models # noqa: B018, F821 +check_plan_budget # noqa: B018, F821 +check_daily_budget # noqa: B018, F821 +get_daily_spend # noqa: B018, F821