mirror of
https://github.com/usestrix/strix.git
synced 2026-08-23 11:22:37 +02:00
refactor(cost): encapsulate streamed OpenRouter cost cache, clear per run
This commit is contained in:
+19
-14
@@ -8,21 +8,22 @@ from unittest.mock import MagicMock, patch
|
||||
import litellm
|
||||
import pytest
|
||||
|
||||
import strix.report.state as state_module
|
||||
from strix.config.models import (
|
||||
_configure_litellm_compatibility,
|
||||
_install_openrouter_stream_cost_capture,
|
||||
)
|
||||
from strix.report.state import (
|
||||
ReportState,
|
||||
litellm_cost_callback,
|
||||
openrouter_stream_cost,
|
||||
remember_streamed_openrouter_cost,
|
||||
set_global_report_state,
|
||||
streamed_openrouter_costs,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _clear_streamed_costs() -> None:
|
||||
state_module._streamed_openrouter_costs.clear()
|
||||
streamed_openrouter_costs.clear()
|
||||
|
||||
|
||||
def test_streaming_logging_stays_enabled_for_cost_callback() -> None:
|
||||
@@ -181,7 +182,7 @@ def test_openrouter_stream_cost_extracts_plain_and_byok_totals() -> None:
|
||||
|
||||
def test_cost_callback_recovers_streamed_openrouter_cost_by_response_id() -> None:
|
||||
report_state = MagicMock()
|
||||
remember_streamed_openrouter_cost("gen-abc", {"cost": 0.42})
|
||||
streamed_openrouter_costs.remember("gen-abc", {"cost": 0.42})
|
||||
# LiteLLM strips cost from the rebuilt streamed usage; only the id survives.
|
||||
response = SimpleNamespace(id="gen-abc", usage=SimpleNamespace(cost=None), _hidden_params={})
|
||||
|
||||
@@ -193,12 +194,12 @@ def test_cost_callback_recovers_streamed_openrouter_cost_by_response_id() -> Non
|
||||
|
||||
report_state.record_observed_llm_cost.assert_called_once_with(0.42)
|
||||
# The entry is consumed so a later response cannot double-count it.
|
||||
assert "gen-abc" not in state_module._streamed_openrouter_costs
|
||||
assert streamed_openrouter_costs.take(response) is None
|
||||
|
||||
|
||||
def test_streamed_openrouter_cost_prefers_provider_report_over_estimate() -> None:
|
||||
report_state = MagicMock()
|
||||
remember_streamed_openrouter_cost("gen-xyz", {"cost": 0.9})
|
||||
streamed_openrouter_costs.remember("gen-xyz", {"cost": 0.9})
|
||||
response = SimpleNamespace(
|
||||
id="gen-xyz",
|
||||
usage=SimpleNamespace(prompt_tokens=10, completion_tokens=5, total_tokens=15),
|
||||
@@ -215,14 +216,16 @@ def test_streamed_openrouter_cost_prefers_provider_report_over_estimate() -> Non
|
||||
estimate.assert_not_called()
|
||||
|
||||
|
||||
def test_remember_streamed_openrouter_cost_evicts_oldest_over_limit() -> None:
|
||||
limit = state_module._STREAMED_OPENROUTER_COST_LIMIT
|
||||
for i in range(limit + 5):
|
||||
remember_streamed_openrouter_cost(f"gen-{i}", {"cost": 0.001})
|
||||
def test_streamed_openrouter_costs_ignores_entries_without_cost() -> None:
|
||||
streamed_openrouter_costs.remember("gen-none", {"prompt_tokens": 10})
|
||||
streamed_openrouter_costs.remember("", {"cost": 0.5})
|
||||
assert streamed_openrouter_costs.take(SimpleNamespace(id="gen-none")) is None
|
||||
|
||||
assert len(state_module._streamed_openrouter_costs) == limit
|
||||
assert "gen-0" not in state_module._streamed_openrouter_costs
|
||||
assert f"gen-{limit + 4}" in state_module._streamed_openrouter_costs
|
||||
|
||||
def test_streamed_openrouter_costs_cleared_on_new_run() -> None:
|
||||
streamed_openrouter_costs.remember("gen-stale", {"cost": 0.7})
|
||||
set_global_report_state(ReportState.__new__(ReportState))
|
||||
assert streamed_openrouter_costs.take(SimpleNamespace(id="gen-stale")) is None
|
||||
|
||||
|
||||
def test_openrouter_stream_handler_records_cost() -> None:
|
||||
@@ -243,4 +246,6 @@ def test_openrouter_stream_handler_records_cost() -> None:
|
||||
handler = handler_cls(streaming_response=iter([]), sync_stream=True)
|
||||
handler.chunk_parser(chunk)
|
||||
|
||||
assert state_module._streamed_openrouter_costs["gen-stream"] == pytest.approx(0.0035055)
|
||||
assert streamed_openrouter_costs.take(SimpleNamespace(id="gen-stream")) == pytest.approx(
|
||||
0.0035055
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user