Use observed LiteLLM cost for LiteLLM-routed calls (#529)

Register a litellm.success_callback that captures kwargs['response_cost']
into a new observed-cost bucket on LLMUsageLedger. record() skips the
tokens-times-registry estimate for LiteLLM-routed models so we do not
double-count with the callback; OpenAI direct routes keep estimating
since LiteLLM is not invoked for them. Per-agent attribution for
LiteLLM-routed calls is apportioned by token share at to_record() time.

Author: 0xallam

pyproject.toml

@@ -219,6 +219,9 @@ ignore = [
 # ReportState carries scan artifact/report fields and
 # a runtime ``Callable`` annotation on ``vulnerability_found_callback``.
 "strix/report/state.py" = ["TC003", "PLR0912", "PLR0915", "E501", "PERF401"]
+"strix/report/usage.py" = ["PLC0415"]
+"strix/report/cost_capture.py" = ["PLC0415"]
+"strix/config/models.py" = ["PLC0415"]
 # Interface utility branches per scope-mode / target-type combination;
 # splitting would obscure the decision tree without simplifying it.
 "strix/interface/utils.py" = ["PLR0912", "BLE001", "PLC0415"]

strix/config/models.py

@@ -104,6 +104,22 @@ def _configure_litellm_compatibility() -> None:
     litellm.disable_streaming_logging = True
     litellm.suppress_debug_info = True
 
+    _register_litellm_cost_callback()
+
+
+def _register_litellm_cost_callback() -> None:
+    import litellm
+
+    from strix.report.cost_capture import litellm_cost_callback
+
+    for bucket_name in ("success_callback", "_async_success_callback"):
+        bucket = getattr(litellm, bucket_name, None)
+        if not isinstance(bucket, list):
+            continue
+        if litellm_cost_callback in bucket:
+            continue
+        bucket.append(litellm_cost_callback)
+
 
 def _configure_litellm_default(name: str, value: str) -> None:
     """Set LiteLLM's module-level defaults without adding a provider wrapper."""

strix/report/cost_capture.py

@@ -0,0 +1,53 @@
+"""LiteLLM success-callback that feeds observed cost into the report ledger."""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+
+logger = logging.getLogger(__name__)
+
+
+def litellm_cost_callback(
+    kwargs: dict[str, Any],
+    completion_response: Any,
+    _start_time: Any = None,
+    _end_time: Any = None,
+) -> None:
+    cost = _extract_cost(kwargs, completion_response)
+    if cost is None or cost <= 0:
+        return
+
+    from strix.report.state import get_global_report_state
+
+    report_state = get_global_report_state()
+    if report_state is None:
+        return
+
+    try:
+        report_state._llm_usage.record_observed_cost(cost)
+    except Exception:
+        logger.exception("Failed to record observed LiteLLM cost")
+
+
+def _extract_cost(kwargs: dict[str, Any], completion_response: Any) -> float | None:
+    cost = kwargs.get("response_cost") if isinstance(kwargs, dict) else None
+    if isinstance(cost, int | float) and cost > 0:
+        return float(cost)
+
+    hidden = getattr(completion_response, "_hidden_params", None)
+    if isinstance(hidden, dict):
+        candidate = hidden.get("response_cost")
+        if isinstance(candidate, int | float) and candidate > 0:
+            return float(candidate)
+        headers = hidden.get("additional_headers") or {}
+        if isinstance(headers, dict):
+            from_header = headers.get("llm_provider-x-litellm-response-cost")
+            try:
+                value = float(from_header) if from_header is not None else None
+            except (TypeError, ValueError):
+                value = None
+            if value is not None and value > 0:
+                return value
+    return None

strix/report/usage.py

@@ -18,8 +18,9 @@ def __init__(self) -> None:
         self._total_usage = Usage()
         self._agent_usage: dict[str, Usage] = {}
         self._agent_metadata: dict[str, dict[str, str]] = {}
-        self._total_cost = 0.0
-        self._agent_cost: dict[str, float] = {}
+        self._estimated_cost = 0.0
+        self._agent_estimated_cost: dict[str, float] = {}
+        self._observed_cost = 0.0
 
     def record(
         self,
@@ -33,8 +34,6 @@ def record(
             return False
 
         normalized_agent_id = str(agent_id or "unknown")
-        estimated_cost = _estimate_litellm_cost(usage, model)
-
         self._total_usage.add(usage)
         self._agent_usage.setdefault(normalized_agent_id, Usage()).add(usage)
 
@@ -44,31 +43,49 @@ def record(
         if model:
             metadata["model"] = model
 
-        if estimated_cost is not None:
-            self._total_cost += estimated_cost
-            self._agent_cost[normalized_agent_id] = (
-                self._agent_cost.get(normalized_agent_id, 0.0) + estimated_cost
-            )
+        if not _is_litellm_routed(model):
+            estimated_cost = _estimate_litellm_cost(usage, model)
+            if estimated_cost is not None:
+                self._estimated_cost += estimated_cost
+                self._agent_estimated_cost[normalized_agent_id] = (
+                    self._agent_estimated_cost.get(normalized_agent_id, 0.0) + estimated_cost
+                )
 
         return True
 
+    def record_observed_cost(self, cost: float) -> None:
+        if isinstance(cost, int | float) and cost > 0:
+            self._observed_cost += float(cost)
+
     def to_record(self) -> dict[str, Any]:
         record = serialize_usage(self._total_usage)
-        record["cost"] = _round_cost(self._total_cost)
-        record["cost_source"] = "litellm_estimate"
+        grand_total = self._estimated_cost + self._observed_cost
+        record["cost"] = _round_cost(grand_total)
+        record["cost_source"] = _cost_source_label(self._estimated_cost, self._observed_cost)
         record["agents"] = []
 
+        total_tokens = max(0, int(self._total_usage.total_tokens or 0))
+
         for agent_id in sorted(self._agent_usage):
             usage = self._agent_usage[agent_id]
             metadata = self._agent_metadata.get(agent_id, {})
+            agent_tokens = max(0, int(usage.total_tokens or 0))
+            observed_share = (
+                self._observed_cost * (agent_tokens / total_tokens) if total_tokens else 0.0
+            )
+            agent_total = self._agent_estimated_cost.get(agent_id, 0.0) + observed_share
+
             agent_record = serialize_usage(usage)
             agent_record.update(
                 {
                     "agent_id": agent_id,
                     "agent_name": metadata.get("agent_name") or agent_id,
                     "model": metadata.get("model"),
-                    "cost": _round_cost(self._agent_cost.get(agent_id, 0.0)),
-                    "cost_source": "litellm_estimate",
+                    "cost": _round_cost(agent_total),
+                    "cost_source": _cost_source_label(
+                        self._agent_estimated_cost.get(agent_id, 0.0),
+                        observed_share,
+                    ),
                 }
             )
             record["agents"].append(agent_record)
@@ -79,8 +96,9 @@ def hydrate(self, raw_usage: Any) -> None:
         self._total_usage = Usage()
         self._agent_usage.clear()
         self._agent_metadata.clear()
-        self._total_cost = 0.0
-        self._agent_cost.clear()
+        self._estimated_cost = 0.0
+        self._agent_estimated_cost.clear()
+        self._observed_cost = 0.0
 
         if not isinstance(raw_usage, dict):
             return
@@ -91,7 +109,9 @@ def hydrate(self, raw_usage: Any) -> None:
             logger.exception("Failed to hydrate aggregate llm_usage from run.json")
             self._total_usage = Usage()
 
-        self._total_cost = _float_or_zero(raw_usage.get("cost"))
+        # Resumed runs have already-aggregated cost; treat as estimated. New calls
+        # in this resume add observed cost on top.
+        self._estimated_cost = _float_or_zero(raw_usage.get("cost"))
         agents = raw_usage.get("agents") or []
         if not isinstance(agents, list):
             return
@@ -116,7 +136,24 @@ def hydrate(self, raw_usage: Any) -> None:
             if isinstance(model, str) and model:
                 metadata["model"] = model
             self._agent_metadata[agent_id] = metadata
-            self._agent_cost[agent_id] = _float_or_zero(raw_agent.get("cost"))
+            self._agent_estimated_cost[agent_id] = _float_or_zero(raw_agent.get("cost"))
+
+
+def _is_litellm_routed(model: str | None) -> bool:
+    if not model:
+        return False
+    name = model.strip().lower()
+    if "/" not in name:
+        return False
+    return not name.startswith("openai/")
+
+
+def _cost_source_label(estimated: float, observed: float) -> str:
+    if observed > 0 and estimated > 0:
+        return "mixed"
+    if observed > 0:
+        return "litellm_observed"
+    return "litellm_estimate"
 
 
 def _usage_has_activity(usage: Usage) -> bool:
@@ -171,18 +208,25 @@ def _estimate_litellm_entry_cost(entry: Any, model: str) -> float | None:
     if completion_details:
         usage_payload["completion_tokens_details"] = completion_details
 
-    try:
-        from litellm import completion_cost
-
-        cost = completion_cost(
-            completion_response={
-                "model": model.split("/", 1)[-1],
-                "usage": usage_payload,
-            },
-            model=model,
-        )
-    except Exception:  # noqa: BLE001 - LiteLLM raises plain Exception for unknown model prices.
-        logger.debug("LiteLLM cost estimate unavailable for model %s", model, exc_info=True)
+    from litellm import completion_cost
+
+    candidates = [model]
+    if "/" in model:
+        candidates.append(model.split("/", 1)[-1])
+
+    cost: Any = None
+    for candidate in candidates:
+        try:
+            cost = completion_cost(
+                completion_response={"model": candidate, "usage": usage_payload},
+                model=model,
+            )
+            break
+        except Exception:  # nosec B112  # noqa: BLE001, S112
+            continue
+
+    if cost is None:
+        logger.debug("LiteLLM cost estimate unavailable for model %s", model)
         return None
 
     return cost if isinstance(cost, int | float) and cost >= 0 else None