mirror of
https://github.com/usestrix/strix.git
synced 2026-11-01 07:02:32 +00:00
245 lines
7.8 KiB
Python
245 lines
7.8 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING
|
|
from unittest.mock import patch
|
|
|
|
import litellm
|
|
from agents.usage import Usage
|
|
|
|
from strix.report.pricing import resolve_litellm_model
|
|
from strix.report.usage import LLMUsageLedger
|
|
|
|
|
|
if TYPE_CHECKING:
|
|
import pytest
|
|
|
|
|
|
def test_resolves_common_bare_model_names() -> None:
|
|
resolve_litellm_model.cache_clear()
|
|
assert resolve_litellm_model("deepseek-v4-flash") == "deepseek/deepseek-v4-flash"
|
|
assert resolve_litellm_model("openai/deepseek-v4-flash") == "deepseek/deepseek-v4-flash"
|
|
# LiteLLM prices grok-4.5 identically under xai/, perplexity/xai/ and
|
|
# openrouter/x-ai/, so the resolver may return any of them; the contract
|
|
# is a provider-qualified name LiteLLM can price.
|
|
grok = resolve_litellm_model("grok-4.5")
|
|
assert grok is not None
|
|
assert grok.endswith("/grok-4.5")
|
|
assert grok in litellm.model_cost
|
|
# MiniMax-M3 is sold by several LiteLLM providers at different prices, so
|
|
# the resolver must not guess from its bare name. A provider-qualified
|
|
# model remains deterministic.
|
|
assert resolve_litellm_model("minimax/MiniMax-M3") == "minimax/MiniMax-M3"
|
|
|
|
|
|
def test_resolver_returns_none_for_unresolvable_model() -> None:
|
|
resolve_litellm_model.cache_clear()
|
|
assert resolve_litellm_model("provider/not-a-real-model") is None
|
|
|
|
|
|
def test_ledger_uses_estimate_when_routed_provider_reports_no_cost() -> None:
|
|
usage = Usage()
|
|
usage.requests = 1
|
|
usage.input_tokens = 1000
|
|
usage.output_tokens = 200
|
|
usage.total_tokens = 1200
|
|
ledger = LLMUsageLedger()
|
|
|
|
with patch("litellm.completion_cost", return_value=0.42):
|
|
ledger.record(agent_id="a", usage=usage, model="openai/deepseek-v4-flash")
|
|
|
|
assert ledger.total_cost == 0.42
|
|
|
|
|
|
def test_ledger_prefers_observed_cost_over_estimate() -> None:
|
|
usage = Usage()
|
|
usage.requests = 1
|
|
usage.input_tokens = 1000
|
|
usage.output_tokens = 200
|
|
usage.total_tokens = 1200
|
|
ledger = LLMUsageLedger()
|
|
|
|
with patch("litellm.completion_cost", return_value=0.42):
|
|
ledger.record(agent_id="a", usage=usage, model="openai/deepseek-v4-flash")
|
|
ledger.record_observed_cost(0.17)
|
|
|
|
assert ledger.total_cost == 0.17
|
|
|
|
|
|
def test_unpriced_model_without_manual_rates_costs_zero(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
for name in (
|
|
"STRIX_LLM_INPUT_COST_PER_MTOK",
|
|
"STRIX_LLM_OUTPUT_COST_PER_MTOK",
|
|
"STRIX_LLM_CACHED_INPUT_COST_PER_MTOK",
|
|
):
|
|
monkeypatch.delenv(name, raising=False)
|
|
usage = Usage()
|
|
usage.requests = 1
|
|
usage.input_tokens = 1000
|
|
usage.output_tokens = 200
|
|
usage.total_tokens = 1200
|
|
ledger = LLMUsageLedger()
|
|
|
|
ledger.record(agent_id="a", usage=usage, model="gateway/unpriced-model")
|
|
|
|
assert ledger.total_cost == 0.0
|
|
|
|
|
|
def test_manual_rates_price_usage_and_override_observed_cost(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setenv("STRIX_LLM_INPUT_COST_PER_MTOK", "2")
|
|
monkeypatch.setenv("STRIX_LLM_OUTPUT_COST_PER_MTOK", "8")
|
|
monkeypatch.delenv("STRIX_LLM_CACHED_INPUT_COST_PER_MTOK", raising=False)
|
|
usage = Usage()
|
|
usage.requests = 1
|
|
usage.input_tokens = 1000
|
|
usage.output_tokens = 200
|
|
usage.total_tokens = 1200
|
|
ledger = LLMUsageLedger()
|
|
|
|
ledger.record(agent_id="a", usage=usage, model="gateway/unpriced-model")
|
|
ledger.record_observed_cost(0.17)
|
|
|
|
assert ledger.total_cost == 0.0036
|
|
assert ledger.to_record()["agents"][0]["cost"] == 0.0036
|
|
|
|
|
|
def test_cached_tokens_fall_back_to_input_rate(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("STRIX_LLM_INPUT_COST_PER_MTOK", "5")
|
|
monkeypatch.setenv("STRIX_LLM_OUTPUT_COST_PER_MTOK", "7")
|
|
monkeypatch.delenv("STRIX_LLM_CACHED_INPUT_COST_PER_MTOK", raising=False)
|
|
usage = Usage()
|
|
usage.requests = 1
|
|
usage.input_tokens = 10
|
|
usage.output_tokens = 1
|
|
usage.total_tokens = 11
|
|
usage.input_tokens_details.cached_tokens = 2
|
|
ledger = LLMUsageLedger()
|
|
|
|
ledger.record(agent_id="a", usage=usage, model="gateway/unpriced-model")
|
|
|
|
assert ledger.total_cost == 0.000057
|
|
|
|
|
|
def test_hydrated_usage_uses_configured_rates_not_persisted_cost(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setenv("STRIX_LLM_INPUT_COST_PER_MTOK", "2")
|
|
monkeypatch.setenv("STRIX_LLM_OUTPUT_COST_PER_MTOK", "8")
|
|
usage = Usage()
|
|
usage.requests = 1
|
|
usage.input_tokens = 1000
|
|
usage.output_tokens = 200
|
|
usage.total_tokens = 1200
|
|
original = LLMUsageLedger()
|
|
original.record(agent_id="a", usage=usage, model="gateway/unpriced-model")
|
|
persisted = original.to_record()
|
|
persisted["cost"] = 42.0
|
|
hydrated = LLMUsageLedger()
|
|
|
|
hydrated.hydrate(persisted)
|
|
|
|
assert hydrated.total_cost == 0.0036
|
|
|
|
|
|
def test_invalid_manual_rate_is_ignored(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
caplog: pytest.LogCaptureFixture,
|
|
) -> None:
|
|
monkeypatch.setenv("STRIX_LLM_INPUT_COST_PER_MTOK", "invalid")
|
|
monkeypatch.setenv("STRIX_LLM_OUTPUT_COST_PER_MTOK", "10")
|
|
monkeypatch.delenv("STRIX_LLM_CACHED_INPUT_COST_PER_MTOK", raising=False)
|
|
usage = Usage()
|
|
usage.requests = 1
|
|
usage.input_tokens = 1000
|
|
usage.output_tokens = 200
|
|
usage.total_tokens = 1200
|
|
ledger = LLMUsageLedger()
|
|
|
|
ledger.record(agent_id="a", usage=usage, model="gateway/unpriced-model")
|
|
|
|
assert ledger.total_cost == 0.002
|
|
assert "Ignoring invalid STRIX_LLM_INPUT_COST_PER_MTOK" in caplog.text
|
|
|
|
|
|
def test_zero_cost_overrides_manual_rates(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("STRIX_LLM_INPUT_COST_PER_MTOK", "10")
|
|
usage = Usage()
|
|
usage.requests = 1
|
|
usage.input_tokens = 1000
|
|
usage.total_tokens = 1000
|
|
ledger = LLMUsageLedger()
|
|
ledger.zero_cost = True
|
|
|
|
ledger.record(agent_id="a", usage=usage, model="gateway/unpriced-model")
|
|
|
|
assert ledger.total_cost == 0.0
|
|
|
|
|
|
def test_hydrated_estimate_continues_accumulating_new_estimates() -> None:
|
|
usage = Usage()
|
|
usage.requests = 1
|
|
usage.input_tokens = 1000
|
|
usage.output_tokens = 200
|
|
usage.total_tokens = 1200
|
|
ledger = LLMUsageLedger()
|
|
ledger.hydrate({"cost": 0.42})
|
|
|
|
with patch("litellm.completion_cost", return_value=0.17):
|
|
ledger.record(agent_id="a", usage=usage, model="openai/deepseek-v4-flash")
|
|
|
|
assert ledger.total_cost == 0.59
|
|
|
|
|
|
def test_zero_cost_disables_both_observed_and_estimated_costs() -> None:
|
|
usage = Usage()
|
|
usage.requests = 1
|
|
usage.input_tokens = 1000
|
|
usage.output_tokens = 200
|
|
usage.total_tokens = 1200
|
|
ledger = LLMUsageLedger()
|
|
ledger.zero_cost = True
|
|
|
|
with patch("litellm.completion_cost", return_value=0.42) as estimate:
|
|
ledger.record(agent_id="a", usage=usage, model="deepseek-v4-flash")
|
|
ledger.record_observed_cost(1.0)
|
|
|
|
estimate.assert_not_called()
|
|
assert ledger.total_cost == 0.0
|
|
|
|
|
|
def test_resolver_uses_provider_when_bare_entry_has_one() -> None:
|
|
original = litellm.model_cost
|
|
litellm.model_cost = {
|
|
"example": {
|
|
"litellm_provider": "example-provider",
|
|
"input_cost_per_token": 1.0,
|
|
"output_cost_per_token": 2.0,
|
|
}
|
|
}
|
|
try:
|
|
resolve_litellm_model.cache_clear()
|
|
assert resolve_litellm_model("example") == "example-provider/example"
|
|
finally:
|
|
litellm.model_cost = original
|
|
resolve_litellm_model.cache_clear()
|
|
|
|
|
|
def test_resolver_does_not_guess_between_differently_priced_providers() -> None:
|
|
original = litellm.model_cost
|
|
litellm.model_cost = {
|
|
"provider-a/example": {
|
|
"input_cost_per_token": 1.0,
|
|
"output_cost_per_token": 2.0,
|
|
},
|
|
"provider-b/example": {
|
|
"input_cost_per_token": 3.0,
|
|
"output_cost_per_token": 4.0,
|
|
},
|
|
}
|
|
try:
|
|
resolve_litellm_model.cache_clear()
|
|
assert resolve_litellm_model("example") is None
|
|
finally:
|
|
litellm.model_cost = original
|
|
resolve_litellm_model.cache_clear()
|