mortal/strix
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Simplify cost ledger to one bucket (#531)

Browse Source
This commit is contained in:
Ahmed Allam
2026-06-08 15:56:28 -07:00
committed by GitHub
parent 1c9ab993bb
commit 6c99829325
5 changed files with 69 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pyproject.toml
-1
View File
@@ -220,7 +220,6 @@ ignore = [
# 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/config/models.py
+1 -1
View File
@@ -110,7 +110,7 @@ def _configure_litellm_compatibility() -> None:
def _register_litellm_cost_callback() -> None:
import litellm
from strix.report.cost_capture import litellm_cost_callback
from strix.report.state import litellm_cost_callback
for bucket_name in ("success_callback", "_async_success_callback"):
bucket = getattr(litellm, bucket_name, None)
strix/report/cost_capture.py
-53
View File
@@ -1,53 +0,0 @@
"""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/state.py
+45
View File
@@ -230,6 +230,9 @@ class ReportState:
):
self.save_run_data()
def record_observed_llm_cost(self, cost: float) -> None:
self._llm_usage.record_observed_cost(cost)
def get_total_llm_usage(self) -> dict[str, Any]:
return dict(self.run_record.get("llm_usage") or self._build_llm_usage_record())
@@ -343,3 +346,45 @@ class ReportState:
def _hydrate_llm_usage(self, raw_usage: Any) -> None:
self._llm_usage.hydrate(raw_usage)
self._sync_llm_usage_record()
def litellm_cost_callback(
kwargs: Any,
completion_response: Any,
_start_time: Any = None,
_end_time: Any = None,
) -> None:
"""LiteLLM ``success_callback`` adapter; forwards observed cost to the active scan."""
cost: float | None = None
raw = kwargs.get("response_cost") if isinstance(kwargs, dict) else None
if isinstance(raw, int | float) and raw > 0:
cost = float(raw)
if cost is None:
hidden = getattr(completion_response, "_hidden_params", None) or {}
candidate = hidden.get("response_cost") if isinstance(hidden, dict) else None
if isinstance(candidate, int | float) and candidate > 0:
cost = float(candidate)
else:
headers = hidden.get("additional_headers") or {} if isinstance(hidden, dict) else {}
raw = (
headers.get("llm_provider-x-litellm-response-cost")
if isinstance(headers, dict)
else None
)
try:
value = float(raw) if raw is not None else None
except (TypeError, ValueError):
value = None
if value is not None and value > 0:
cost = value
if cost is None or cost <= 0:
return
report_state = get_global_report_state()
if report_state is None:
return
try:
report_state.record_observed_llm_cost(cost)
except Exception:
logger.exception("Failed to record observed LiteLLM cost")
strix/report/usage.py
+23 -43
View File
@@ -18,9 +18,7 @@ class LLMUsageLedger:
self._total_usage = Usage()
self._agent_usage: dict[str, Usage] = {}
self._agent_metadata: dict[str, dict[str, str]] = {}
self._estimated_cost = 0.0
self._agent_estimated_cost: dict[str, float] = {}
self._observed_cost = 0.0
self._total_cost = 0.0
def record(
self,
@@ -44,36 +42,29 @@ class LLMUsageLedger:
metadata["model"] = model
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
)
estimated = _estimate_litellm_cost(usage, model)
if estimated:
self._total_cost += estimated
return True
def record_observed_cost(self, cost: float) -> None:
if isinstance(cost, int | float) and cost > 0:
self._observed_cost += float(cost)
self._total_cost += float(cost)
def to_record(self) -> dict[str, Any]:
record = serialize_usage(self._total_usage)
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["cost"] = _round_cost(self._total_cost)
record["agents"] = []
total_tokens = max(0, int(self._total_usage.total_tokens or 0))
agent_tokens = {aid: _resolve_total_tokens(u) for aid, u in self._agent_usage.items()}
total_tokens = sum(agent_tokens.values())
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_cost = (
self._total_cost * (agent_tokens[agent_id] / 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(
@@ -81,11 +72,7 @@ class LLMUsageLedger:
"agent_id": agent_id,
"agent_name": metadata.get("agent_name") or agent_id,
"model": metadata.get("model"),
"cost": _round_cost(agent_total),
"cost_source": _cost_source_label(
self._agent_estimated_cost.get(agent_id, 0.0),
observed_share,
),
"cost": _round_cost(agent_cost),
}
)
record["agents"].append(agent_record)
@@ -96,9 +83,7 @@ class LLMUsageLedger:
self._total_usage = Usage()
self._agent_usage.clear()
self._agent_metadata.clear()
self._estimated_cost = 0.0
self._agent_estimated_cost.clear()
self._observed_cost = 0.0
self._total_cost = 0.0
if not isinstance(raw_usage, dict):
return
@@ -109,14 +94,9 @@ class LLMUsageLedger:
logger.exception("Failed to hydrate aggregate llm_usage from run.json")
self._total_usage = Usage()
# 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
self._total_cost = _float_or_zero(raw_usage.get("cost"))
for raw_agent in agents:
for raw_agent in raw_usage.get("agents") or []:
if not isinstance(raw_agent, dict):
continue
agent_id = str(raw_agent.get("agent_id") or "").strip()
@@ -136,7 +116,15 @@ class LLMUsageLedger:
if isinstance(model, str) and model:
metadata["model"] = model
self._agent_metadata[agent_id] = metadata
self._agent_estimated_cost[agent_id] = _float_or_zero(raw_agent.get("cost"))
def _resolve_total_tokens(usage: Usage) -> int:
total = max(0, int(usage.total_tokens or 0))
if total > 0:
return total
prompt = _int_or_zero(getattr(usage, "input_tokens", 0))
completion = _int_or_zero(getattr(usage, "output_tokens", 0))
return prompt + completion
def _is_litellm_routed(model: str | None) -> bool:
@@ -148,14 +136,6 @@ def _is_litellm_routed(model: str | None) -> bool:
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:
return bool(
usage.requests