diff --git a/strix/core/hooks.py b/strix/core/hooks.py index 64562d74..133375ab 100644 --- a/strix/core/hooks.py +++ b/strix/core/hooks.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging +import math from typing import TYPE_CHECKING, Any from agents.lifecycle import RunHooks @@ -27,8 +28,6 @@ class ReportUsageHooks(RunHooks[dict[str, Any]]): """Persist SDK-native usage after every model response.""" def __init__(self, *, model: str, max_budget_usd: float | None = None) -> None: - import math - if max_budget_usd is not None and ( not math.isfinite(max_budget_usd) or max_budget_usd <= 0 ): diff --git a/strix/report/state.py b/strix/report/state.py index 178c4047..217b266d 100644 --- a/strix/report/state.py +++ b/strix/report/state.py @@ -534,16 +534,10 @@ def litellm_cost_callback( cost = value if cost is None: - usage: Any = getattr(completion_response, "usage", None) - if usage is None and isinstance(completion_response, dict): - usage = cast("dict[str, Any]", completion_response).get("usage") - usage_cost: Any - if isinstance(usage, dict): - usage_cost = cast("dict[str, Any]", usage).get("cost") - else: - usage_cost = getattr(usage, "cost", None) - if isinstance(usage_cost, int | float) and usage_cost > 0: - cost = float(usage_cost) + cost = _usage_reported_cost(completion_response) + + if cost is None: + cost = _estimate_response_cost(kwargs, completion_response) if cost is None or cost <= 0: return @@ -554,3 +548,101 @@ def litellm_cost_callback( report_state.record_observed_llm_cost(cost) except Exception: logger.exception("Failed to record observed LiteLLM cost") + + +def _usage_reported_cost(completion_response: Any) -> float | None: + """Provider-reported cost from the ``usage`` block (e.g. OpenRouter). + + Non-BYOK responses charge everything to ``usage.cost``. BYOK responses + charge only the OpenRouter fee to ``usage.cost`` (often 0) and report the + provider charge in ``usage.cost_details.upstream_inference_cost``, so the + true BYOK total is the sum of the two. + """ + usage: Any = getattr(completion_response, "usage", None) + if usage is None and isinstance(completion_response, dict): + usage = cast("dict[str, Any]", completion_response).get("usage") + if usage is None: + return None + + def _field(container: Any, name: str) -> Any: + if isinstance(container, dict): + return cast("dict[str, Any]", container).get(name) + return getattr(container, name, None) + + total = 0.0 + usage_cost = _field(usage, "cost") + if isinstance(usage_cost, int | float) and usage_cost > 0: + total += float(usage_cost) + + if bool(_field(usage, "is_byok")): + upstream = _field(_field(usage, "cost_details"), "upstream_inference_cost") + if isinstance(upstream, int | float) and upstream > 0: + total += float(upstream) + + return total if total > 0 else None + + +def _estimate_response_cost(kwargs: Any, completion_response: Any) -> float | None: + """Best-effort LiteLLM cost-map estimate when no provider-reported cost exists. + + LiteLLM strips provider cost fields when rebuilding streamed responses and + returns no ``response_cost`` for models missing from its cost map, so try + the provider-prefixed name, the raw name, and the bare model name. + """ + from litellm import completion_cost + + model = kwargs.get("model") if isinstance(kwargs, dict) else None + if not isinstance(model, str) or not model: + if isinstance(completion_response, dict): + model = cast("dict[str, Any]", completion_response).get("model") + else: + model = getattr(completion_response, "model", None) + if not isinstance(model, str) or not model: + return None + + provider = None + litellm_params = kwargs.get("litellm_params") if isinstance(kwargs, dict) else None + if isinstance(litellm_params, dict): + provider = litellm_params.get("custom_llm_provider") + + usage_payload = _usage_payload(completion_response) + if usage_payload is None: + return None + + candidates: list[str] = [] + if isinstance(provider, str) and provider and not model.startswith(f"{provider}/"): + candidates.append(f"{provider}/{model}") + candidates.append(model) + if "/" in model: + candidates.append(model.rsplit("/", 1)[-1]) + + for candidate in candidates: + try: + value = completion_cost( + completion_response={"model": candidate, "usage": usage_payload}, + model=candidate, + ) + except Exception: # nosec B112 # noqa: BLE001, S112 + continue + if isinstance(value, int | float) and value > 0: + return float(value) + return None + + +def _usage_payload(completion_response: Any) -> dict[str, Any] | None: + """Token counts as a plain dict, detached from the response's provider metadata.""" + usage: Any = getattr(completion_response, "usage", None) + if usage is None and isinstance(completion_response, dict): + usage = cast("dict[str, Any]", completion_response).get("usage") + if usage is None: + return None + if hasattr(usage, "model_dump"): + usage = usage.model_dump() + if not isinstance(usage, dict): + return None + payload = cast("dict[str, Any]", usage) + if not payload.get("total_tokens") and not ( + payload.get("prompt_tokens") or payload.get("completion_tokens") + ): + return None + return payload diff --git a/tests/test_cost_tracking.py b/tests/test_cost_tracking.py index fdbf7c17..543d3fd5 100644 --- a/tests/test_cost_tracking.py +++ b/tests/test_cost_tracking.py @@ -6,6 +6,7 @@ from types import SimpleNamespace from unittest.mock import MagicMock, patch import litellm +import pytest from strix.config.models import _configure_litellm_compatibility from strix.report.state import litellm_cost_callback @@ -42,3 +43,111 @@ def test_cost_callback_reads_usage_cost_from_mapping_response() -> None: litellm_cost_callback({}, response) report_state.record_observed_llm_cost.assert_called_once_with(0.125) + + +def test_cost_callback_reads_byok_upstream_inference_cost() -> None: + report_state = MagicMock() + response = SimpleNamespace( + usage=SimpleNamespace( + cost=0, + is_byok=True, + cost_details=SimpleNamespace(upstream_inference_cost=6.75e-06), + ), + _hidden_params={}, + ) + + with patch("strix.report.state.get_global_report_state", return_value=report_state): + litellm_cost_callback({"response_cost": None}, response) + + report_state.record_observed_llm_cost.assert_called_once_with(6.75e-06) + + +def test_cost_callback_sums_usage_cost_and_upstream_inference_cost() -> None: + report_state = MagicMock() + response = { + "usage": { + "cost": 0.01, + "is_byok": True, + "cost_details": {"upstream_inference_cost": 0.2}, + } + } + + with patch("strix.report.state.get_global_report_state", return_value=report_state): + litellm_cost_callback({}, response) + + report_state.record_observed_llm_cost.assert_called_once_with(pytest.approx(0.21)) + + +def test_cost_callback_ignores_upstream_cost_for_non_byok_responses() -> None: + report_state = MagicMock() + response = { + "usage": { + "cost": 0.05, + "is_byok": False, + "cost_details": {"upstream_inference_cost": 0.04}, + } + } + + with patch("strix.report.state.get_global_report_state", return_value=report_state): + litellm_cost_callback({}, response) + + report_state.record_observed_llm_cost.assert_called_once_with(0.05) + + +def test_cost_callback_estimates_cost_with_provider_prefixed_model() -> None: + report_state = MagicMock() + response = {"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}} + kwargs = { + "response_cost": None, + "model": "anthropic/claude-sonnet-4.5", + "litellm_params": {"custom_llm_provider": "openrouter"}, + } + + def fake_completion_cost(**kwargs: object) -> float: + if kwargs["model"] == "openrouter/anthropic/claude-sonnet-4.5": + return 0.5 + raise ValueError(kwargs["model"]) + + with ( + patch("strix.report.state.get_global_report_state", return_value=report_state), + patch("litellm.completion_cost", side_effect=fake_completion_cost), + ): + litellm_cost_callback(kwargs, response) + + report_state.record_observed_llm_cost.assert_called_once_with(0.5) + + +def test_cost_callback_estimates_cost_with_bare_model_fallback() -> None: + report_state = MagicMock() + response = {"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}} + kwargs = { + "response_cost": None, + "model": "openai/gpt-4o-mini", + "litellm_params": {"custom_llm_provider": "openrouter"}, + } + + def fake_completion_cost(**kwargs: object) -> float: + if kwargs["model"] == "gpt-4o-mini": + return 0.025 + raise ValueError(kwargs["model"]) + + with ( + patch("strix.report.state.get_global_report_state", return_value=report_state), + patch("litellm.completion_cost", side_effect=fake_completion_cost), + ): + litellm_cost_callback(kwargs, response) + + report_state.record_observed_llm_cost.assert_called_once_with(0.025) + + +def test_cost_callback_records_nothing_when_no_cost_available() -> None: + report_state = MagicMock() + response = {"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}} + + with ( + patch("strix.report.state.get_global_report_state", return_value=report_state), + patch("litellm.completion_cost", side_effect=ValueError("unknown model")), + ): + litellm_cost_callback({"response_cost": None, "model": "x/y"}, response) + + report_state.record_observed_llm_cost.assert_not_called()