diff --git a/pyproject.toml b/pyproject.toml index d12c3b3..6065f6c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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. diff --git a/strix/config/models.py b/strix/config/models.py index 6c050fc..f582643 100644 --- a/strix/config/models.py +++ b/strix/config/models.py @@ -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) diff --git a/strix/report/cost_capture.py b/strix/report/cost_capture.py deleted file mode 100644 index 78c05bb..0000000 --- a/strix/report/cost_capture.py +++ /dev/null @@ -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 diff --git a/strix/report/state.py b/strix/report/state.py index a69ced7..1ee72c5 100644 --- a/strix/report/state.py +++ b/strix/report/state.py @@ -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") diff --git a/strix/report/usage.py b/strix/report/usage.py index eaf7aad..58977c6 100644 --- a/strix/report/usage.py +++ b/strix/report/usage.py @@ -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