feat: add cost tracking to observability layer (APPENG-5370)

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",
 ]

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,
         }

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)

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:

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