diff --git a/pyproject.toml b/pyproject.toml index ce97bbc..d12c3b3 100644 --- a/pyproject.toml +++ b/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"] diff --git a/strix/config/models.py b/strix/config/models.py index 4d2e072..6c050fc 100644 --- a/strix/config/models.py +++ b/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.""" diff --git a/strix/report/cost_capture.py b/strix/report/cost_capture.py new file mode 100644 index 0000000..78c05bb --- /dev/null +++ b/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 diff --git a/strix/report/usage.py b/strix/report/usage.py index 69f6f28..eaf7aad 100644 --- a/strix/report/usage.py +++ b/strix/report/usage.py @@ -18,8 +18,9 @@ class LLMUsageLedger: 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 @@ class LLMUsageLedger: 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 @@ class LLMUsageLedger: 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 @@ class LLMUsageLedger: 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 @@ class LLMUsageLedger: 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 @@ class LLMUsageLedger: 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 + 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) + 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