diff --git a/abevalflow/observability/__init__.py b/abevalflow/observability/__init__.py index 22567959..6739513d 100644 --- a/abevalflow/observability/__init__.py +++ b/abevalflow/observability/__init__.py @@ -1,12 +1,14 @@ """Observability layer for Agentic Eval Flow pipeline metrics and tracing.""" from abevalflow.observability.context import MetricsContext, TimingRecord, TokenUsage +from abevalflow.observability.cost import calculate_cost from abevalflow.observability.otel import get_tracer, is_otel_enabled __all__ = [ "MetricsContext", "TimingRecord", "TokenUsage", + "calculate_cost", "get_tracer", "is_otel_enabled", ] diff --git a/abevalflow/observability/context.py b/abevalflow/observability/context.py index 00616867..bad78c53 100644 --- a/abevalflow/observability/context.py +++ b/abevalflow/observability/context.py @@ -120,6 +120,11 @@ def load_checkpoint(cls, workspace_path: Path) -> MetricsContext | None: def to_observability_dict(self) -> dict: """Convert to kwargs for ObservabilityMetricsRow.""" + from abevalflow.observability.cost import calculate_cost + + prompt = self.total_prompt_tokens + completion = self.total_completion_tokens + return { "submission_name": self.submission_name, "model_name": self.model_name, @@ -129,8 +134,9 @@ def to_observability_dict(self) -> dict: "evaluate_duration_ms": self.timing_ms("evaluate"), "analyze_duration_ms": self.timing_ms("analyze"), "store_duration_ms": self.timing_ms("store"), - "total_prompt_tokens": self.total_prompt_tokens or None, - "total_completion_tokens": self.total_completion_tokens or None, + "total_prompt_tokens": prompt or None, + "total_completion_tokens": completion or None, "total_tokens": self.total_tokens or None, + "estimated_cost_usd": calculate_cost(prompt, completion, self.model_name), "llm_calls_count": self.llm_calls_count or None, } diff --git a/abevalflow/observability/cost.py b/abevalflow/observability/cost.py new file mode 100644 index 00000000..b089834c --- /dev/null +++ b/abevalflow/observability/cost.py @@ -0,0 +1,67 @@ +"""Cost estimation from token counts and per-model pricing. + +Reads model_costs.yaml for per-model rates (USD per 1K tokens). +Lookup order: exact match → prefix match → _default entry → zero. +""" + +from __future__ import annotations + +import logging +from functools import lru_cache +from pathlib import Path + +import yaml + +logger = logging.getLogger(__name__) + +_COSTS_PATH = Path(__file__).resolve().parents[2] / "config" / "observability" / "model_costs.yaml" + + +@lru_cache(maxsize=1) +def _load_model_costs(path: Path = _COSTS_PATH) -> dict[str, dict[str, float]]: + if not path.is_file(): + logger.warning("Model costs file not found: %s", path) + return {} + with path.open() as f: + data = yaml.safe_load(f) or {} + return {k: v for k, v in data.items() if isinstance(v, dict)} + + +def _resolve_rates(model_name: str, costs: dict[str, dict[str, float]]) -> dict[str, float]: + if model_name in costs: + return costs[model_name] + + for key in costs: + if key == "_default": + continue + if model_name.startswith(key) or key.startswith(model_name): + return costs[key] + + return costs.get("_default", {}) + + +def calculate_cost( + prompt_tokens: int, + completion_tokens: int, + model_name: str | None, + costs_path: Path | None = None, +) -> float | None: + """Estimate USD cost from token counts and model name. + + Returns None if model_name is not provided or no pricing is available. + """ + if not model_name: + return None + + costs = _load_model_costs(costs_path or _COSTS_PATH) + if not costs: + return None + + rates = _resolve_rates(model_name, costs) + if not rates: + return None + + input_rate = rates.get("input_per_1k", 0) + output_rate = rates.get("output_per_1k", 0) + + return round((prompt_tokens / 1000) * input_rate + (completion_tokens / 1000) * output_rate, 6) diff --git a/abevalflow/observability/mlflow_observer.py b/abevalflow/observability/mlflow_observer.py index 9a1970fc..b8b46395 100644 --- a/abevalflow/observability/mlflow_observer.py +++ b/abevalflow/observability/mlflow_observer.py @@ -164,14 +164,22 @@ def _log_observability_metrics(report_dir: Path) -> None: total_calls += usage.get("call_count", 0) if total_prompt > 0 or total_completion > 0: - mlflow.log_metrics( - { - "total_prompt_tokens": total_prompt, - "total_completion_tokens": total_completion, - "total_tokens": total_prompt + total_completion, - "llm_calls_count": total_calls, - } - ) + metrics = { + "total_prompt_tokens": total_prompt, + "total_completion_tokens": total_completion, + "total_tokens": total_prompt + total_completion, + "llm_calls_count": total_calls, + } + + model_name = checkpoint.get("model_name") + if model_name: + from abevalflow.observability.cost import calculate_cost + + cost = calculate_cost(total_prompt, total_completion, model_name) + if cost is not None: + metrics["estimated_cost_usd"] = cost + + mlflow.log_metrics(metrics) model_name = checkpoint.get("model_name") if model_name: diff --git a/tests/test_cost.py b/tests/test_cost.py new file mode 100644 index 00000000..119d3121 --- /dev/null +++ b/tests/test_cost.py @@ -0,0 +1,117 @@ +"""Tests for abevalflow.observability.cost — cost estimation from token counts.""" + +from __future__ import annotations + +from pathlib import Path +from textwrap import dedent + +import pytest + +from abevalflow.observability.cost import _load_model_costs, _resolve_rates, calculate_cost + + +@pytest.fixture() +def costs_file(tmp_path: Path) -> Path: + p = tmp_path / "model_costs.yaml" + p.write_text( + dedent("""\ + claude-sonnet: + input_per_1k: 0.003 + output_per_1k: 0.015 + gpt-4o: + input_per_1k: 0.005 + output_per_1k: 0.015 + _default: + input_per_1k: 0.002 + output_per_1k: 0.010 + """) + ) + return p + + +class TestLoadModelCosts: + def test_load_valid_file(self, costs_file: Path) -> None: + _load_model_costs.cache_clear() + costs = _load_model_costs(costs_file) + assert "claude-sonnet" in costs + assert costs["claude-sonnet"]["input_per_1k"] == 0.003 + + def test_load_missing_file(self, tmp_path: Path) -> None: + _load_model_costs.cache_clear() + costs = _load_model_costs(tmp_path / "nonexistent.yaml") + assert costs == {} + + def test_load_empty_file(self, tmp_path: Path) -> None: + _load_model_costs.cache_clear() + p = tmp_path / "empty.yaml" + p.write_text("") + costs = _load_model_costs(p) + assert costs == {} + + +class TestResolveRates: + def test_exact_match(self) -> None: + costs = {"claude-sonnet": {"input_per_1k": 0.003, "output_per_1k": 0.015}} + assert _resolve_rates("claude-sonnet", costs) == costs["claude-sonnet"] + + def test_prefix_match(self) -> None: + costs = {"claude-sonnet": {"input_per_1k": 0.003, "output_per_1k": 0.015}} + assert _resolve_rates("claude-sonnet-4-20250514", costs) == costs["claude-sonnet"] + + def test_default_fallback(self) -> None: + costs = { + "gpt-4o": {"input_per_1k": 0.005, "output_per_1k": 0.015}, + "_default": {"input_per_1k": 0.002, "output_per_1k": 0.010}, + } + assert _resolve_rates("unknown-model", costs) == costs["_default"] + + def test_no_match_no_default(self) -> None: + costs = {"gpt-4o": {"input_per_1k": 0.005, "output_per_1k": 0.015}} + assert _resolve_rates("unknown-model", costs) == {} + + +class TestCalculateCost: + def test_basic_calculation(self, costs_file: Path) -> None: + _load_model_costs.cache_clear() + cost = calculate_cost(1000, 500, "claude-sonnet", costs_path=costs_file) + # (1000/1000)*0.003 + (500/1000)*0.015 = 0.003 + 0.0075 = 0.0105 + assert cost == 0.0105 + + def test_zero_tokens(self, costs_file: Path) -> None: + _load_model_costs.cache_clear() + cost = calculate_cost(0, 0, "claude-sonnet", costs_path=costs_file) + assert cost == 0.0 + + def test_no_model_name(self, costs_file: Path) -> None: + _load_model_costs.cache_clear() + cost = calculate_cost(1000, 500, None, costs_path=costs_file) + assert cost is None + + def test_unknown_model_uses_default(self, costs_file: Path) -> None: + _load_model_costs.cache_clear() + cost = calculate_cost(1000, 500, "some-unknown-model", costs_path=costs_file) + # (1000/1000)*0.002 + (500/1000)*0.010 = 0.002 + 0.005 = 0.007 + assert cost == 0.007 + + def test_missing_costs_file(self, tmp_path: Path) -> None: + _load_model_costs.cache_clear() + cost = calculate_cost(1000, 500, "claude-sonnet", costs_path=tmp_path / "nope.yaml") + assert cost is None + + def test_prefix_match_versioned_model(self, costs_file: Path) -> None: + _load_model_costs.cache_clear() + cost = calculate_cost(2000, 1000, "claude-sonnet-4-20250514", costs_path=costs_file) + # (2000/1000)*0.003 + (1000/1000)*0.015 = 0.006 + 0.015 = 0.021 + assert cost == 0.021 + + def test_precision(self, costs_file: Path) -> None: + _load_model_costs.cache_clear() + cost = calculate_cost(15000, 3500, "claude-sonnet", costs_path=costs_file) + # (15000/1000)*0.003 + (3500/1000)*0.015 = 0.045 + 0.0525 = 0.0975 + assert cost == 0.0975 + + def test_large_token_counts(self, costs_file: Path) -> None: + _load_model_costs.cache_clear() + cost = calculate_cost(1_000_000, 500_000, "claude-sonnet", costs_path=costs_file) + # (1M/1000)*0.003 + (500K/1000)*0.015 = 3.0 + 7.5 = 10.5 + assert cost == 10.5