wehub-resource-sync
315 lines
12 KiB
Python
315 lines
12 KiB
Python
"""Cache-aware, provenance-labeled cost estimates in the usage tracker and turn loop.
|
|
|
|
Covers the issue-490 fix: unbilled token buckets are priced with the
|
|
four-bucket ``estimate_cost`` at the layered-resolver price, every
|
|
cost-fields dict discloses ``estimateBasis``/``priceSource``, and a model
|
|
mixing billed and unbilled calls reports billed + estimate instead of
|
|
collapsing to billed-only.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import AsyncIterator
|
|
from typing import Any
|
|
|
|
import pytest
|
|
|
|
from opensquilla.engine import Agent, AgentConfig
|
|
from opensquilla.engine.turn_runner.turn_finalizer_stage import _turn_usage_payload
|
|
from opensquilla.engine.types import DoneEvent
|
|
from opensquilla.engine.usage import SessionUsage, UsageTracker, usage_scope
|
|
from opensquilla.provider import DoneEvent as ProviderDone
|
|
from opensquilla.provider import TextDeltaEvent as ProviderText
|
|
|
|
|
|
def test_issue_490_row_estimate_is_cache_aware(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("OPENSQUILLA_OPENROUTER_LIVE_PRICING", "0")
|
|
su = SessionUsage()
|
|
su.add(
|
|
11_559_964,
|
|
262_086,
|
|
model_id="deepseek/deepseek-v4-pro",
|
|
cache_read_tokens=10_313_958,
|
|
provider="deepseek",
|
|
)
|
|
row = su.model_breakdown[0]
|
|
# fresh 1,246,006 * 0.435 + read 10,313,958 * 0.003625 + out 262,086 * 0.87 (per M)
|
|
assert row["costUsd"] == pytest.approx(0.807416, abs=1e-4)
|
|
assert row["estimateBasis"] == "cache_aware"
|
|
# The vendored models.dev snapshot carries deepseek-v4-pro's own cost
|
|
# keys, so the provider-qualified lookup resolves through the catalog
|
|
# layer (same rates as the static table).
|
|
assert row["priceSource"] == "catalog"
|
|
assert row["costSource"] == "opensquilla_estimate"
|
|
# Session cost property delegates to the same cache-aware math.
|
|
assert su.cost == pytest.approx(0.807416, abs=1e-4)
|
|
assert su.total_cost == pytest.approx(0.807416, abs=1e-4)
|
|
|
|
|
|
def test_unknown_cache_rate_labels_cache_blind(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("OPENSQUILLA_OPENROUTER_LIVE_PRICING", "0")
|
|
su = SessionUsage()
|
|
su.add(1_000_000, 0, model_id="kimi-latest", cache_read_tokens=900_000, provider="moonshot")
|
|
row = su.model_breakdown[0]
|
|
assert row["estimateBasis"] == "cache_blind"
|
|
assert row["priceSource"] == "default"
|
|
# Cache-blind fallback prices the full input at the default rate.
|
|
assert row["costUsd"] == pytest.approx(3.0)
|
|
|
|
|
|
def test_billed_plus_unbilled_same_model_is_mixed_not_collapsed(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setenv("OPENSQUILLA_OPENROUTER_LIVE_PRICING", "0")
|
|
su = SessionUsage()
|
|
su.add(1000, 100, model_id="deepseek/deepseek-v4-pro", billed_cost=0.05, provider="openrouter")
|
|
su.add(2_000_000, 0, model_id="deepseek/deepseek-v4-pro", provider="deepseek")
|
|
row = su.model_breakdown[0]
|
|
assert row["costSource"] == "mixed"
|
|
assert row["costUsd"] == pytest.approx(0.05 + 2 * 0.435, abs=1e-4)
|
|
assert row["billedCostUsd"] == pytest.approx(0.05)
|
|
# total_cost keeps the row invariant: billed + estimate-of-unbilled.
|
|
assert su.total_cost == pytest.approx(0.05 + 2 * 0.435, abs=1e-4)
|
|
breakdown_sum = sum(item["costUsd"] for item in su.model_breakdown)
|
|
assert breakdown_sum == pytest.approx(su.total_cost, abs=1e-6)
|
|
|
|
|
|
def test_billed_only_row_has_null_basis_but_price_source(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setenv("OPENSQUILLA_OPENROUTER_LIVE_PRICING", "0")
|
|
su = SessionUsage()
|
|
su.add(1000, 50, model_id="claude-opus-4-7", billed_cost=0.05)
|
|
row = su.model_breakdown[0]
|
|
assert row["costSource"] == "provider_billed"
|
|
assert row["costUsd"] == pytest.approx(0.05)
|
|
assert row["estimatedCostUsd"] == 0.0
|
|
assert row["estimateBasis"] is None
|
|
assert row["estimate_basis"] is None
|
|
assert row["priceSource"] == "static_table"
|
|
assert row["price_source"] == "static_table"
|
|
|
|
|
|
def test_local_free_row_labels_free_basis(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("OPENSQUILLA_OPENROUTER_LIVE_PRICING", "0")
|
|
su = SessionUsage()
|
|
su.add(10_000, 500, model_id="qwen3:4b", provider="ollama")
|
|
row = su.model_breakdown[0]
|
|
assert row["costUsd"] == 0.0
|
|
assert row["estimateBasis"] == "free"
|
|
assert row["priceSource"] == "local_free"
|
|
|
|
|
|
def test_unbilled_counters_accumulate_only_when_call_is_unbilled(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setenv("OPENSQUILLA_OPENROUTER_LIVE_PRICING", "0")
|
|
su = SessionUsage()
|
|
su.add(1000, 100, model_id="m", billed_cost=0.05, cache_read_tokens=300, cache_write_tokens=10)
|
|
su.add(2000, 200, model_id="m", cache_read_tokens=500, cache_write_tokens=20)
|
|
|
|
assert su._per_model is not None
|
|
mu = su._per_model["m"]
|
|
# Totals accumulate everything.
|
|
assert mu.input_tokens == 3000
|
|
assert mu.output_tokens == 300
|
|
assert mu.cache_read_tokens == 800
|
|
assert mu.cache_write_tokens == 30
|
|
# Unbilled buckets carry only the unbilled call.
|
|
assert mu.unbilled_input_tokens == 2000
|
|
assert mu.unbilled_output_tokens == 200
|
|
assert mu.unbilled_cache_read_tokens == 500
|
|
assert mu.unbilled_cache_write_tokens == 20
|
|
|
|
|
|
def test_scoped_usage_accumulates_unbilled_counters(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("OPENSQUILLA_OPENROUTER_LIVE_PRICING", "0")
|
|
tracker = UsageTracker()
|
|
with usage_scope("subagent:researcher"):
|
|
tracker.add("s1", 1000, 50, model_id="m", cache_read_tokens=400)
|
|
scoped = tracker.get_scope("s1", "subagent:researcher")
|
|
assert scoped is not None
|
|
assert scoped._per_model is not None
|
|
mu = scoped._per_model["m"]
|
|
assert mu.unbilled_input_tokens == 1000
|
|
assert mu.unbilled_output_tokens == 50
|
|
assert mu.unbilled_cache_read_tokens == 400
|
|
|
|
|
|
def test_checkpoint_clone_preserves_unbilled_counters(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("OPENSQUILLA_OPENROUTER_LIVE_PRICING", "0")
|
|
tracker = UsageTracker()
|
|
tracker.add("s1", 1000, 50, model_id="m", cache_read_tokens=400, cache_write_tokens=30)
|
|
checkpoint = tracker.session_checkpoint("s1")
|
|
assert checkpoint is not None
|
|
assert checkpoint._per_model is not None
|
|
mu = checkpoint._per_model["m"]
|
|
assert mu.unbilled_input_tokens == 1000
|
|
assert mu.unbilled_output_tokens == 50
|
|
assert mu.unbilled_cache_read_tokens == 400
|
|
assert mu.unbilled_cache_write_tokens == 30
|
|
|
|
# Mutating the live session must not leak into the checkpoint.
|
|
tracker.add("s1", 500, 5, model_id="m")
|
|
assert mu.unbilled_input_tokens == 1000
|
|
|
|
|
|
def test_delta_snapshot_prices_unbilled_delta_cache_aware(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setenv("OPENSQUILLA_OPENROUTER_LIVE_PRICING", "0")
|
|
tracker = UsageTracker()
|
|
sk = "agent:main:delta"
|
|
tracker.add(
|
|
sk,
|
|
1000,
|
|
100,
|
|
model_id="deepseek/deepseek-v4-pro",
|
|
billed_cost=0.05,
|
|
provider="deepseek",
|
|
)
|
|
checkpoint = tracker.session_checkpoint(sk)
|
|
tracker.add(
|
|
sk,
|
|
1_000_000,
|
|
0,
|
|
model_id="deepseek/deepseek-v4-pro",
|
|
cache_read_tokens=900_000,
|
|
provider="deepseek",
|
|
)
|
|
|
|
delta = tracker.session_delta_snapshot(sk, checkpoint)
|
|
assert delta is not None
|
|
assert delta.input_tokens == 1_000_000
|
|
assert delta.cache_read_tokens == 900_000
|
|
assert delta.billed_cost == 0.0
|
|
# fresh 100,000 * 0.435 + read 900,000 * 0.003625 (per M) = 0.0467625
|
|
assert delta.cost_usd == pytest.approx(0.0467625, abs=1e-6)
|
|
|
|
|
|
def test_delta_snapshot_mixed_turn_adds_billed_and_estimate(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setenv("OPENSQUILLA_OPENROUTER_LIVE_PRICING", "0")
|
|
tracker = UsageTracker()
|
|
sk = "agent:main:mixed-delta"
|
|
checkpoint = tracker.session_checkpoint(sk)
|
|
tracker.add(
|
|
sk,
|
|
1000,
|
|
100,
|
|
model_id="deepseek/deepseek-v4-pro",
|
|
billed_cost=0.05,
|
|
provider="deepseek",
|
|
)
|
|
tracker.add(sk, 2_000_000, 0, model_id="deepseek/deepseek-v4-pro", provider="deepseek")
|
|
|
|
delta = tracker.session_delta_snapshot(sk, checkpoint)
|
|
assert delta is not None
|
|
assert delta.billed_cost == pytest.approx(0.05)
|
|
# Billed delta + estimate of the unbilled delta — no billed-only collapse.
|
|
assert delta.cost_usd == pytest.approx(0.05 + 2 * 0.435, abs=1e-6)
|
|
|
|
|
|
class _CachedUsageProvider:
|
|
provider_name = "fake"
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
input_tokens: int,
|
|
output_tokens: int,
|
|
cached_tokens: int = 0,
|
|
model: str = "",
|
|
) -> None:
|
|
self._input_tokens = input_tokens
|
|
self._output_tokens = output_tokens
|
|
self._cached_tokens = cached_tokens
|
|
self._model = model
|
|
|
|
def chat(self, messages: Any, tools: Any = None, config: Any = None) -> AsyncIterator[Any]:
|
|
return self._stream()
|
|
|
|
async def _stream(self) -> AsyncIterator[Any]:
|
|
yield ProviderText(text="done")
|
|
yield ProviderDone(
|
|
stop_reason="stop",
|
|
input_tokens=self._input_tokens,
|
|
output_tokens=self._output_tokens,
|
|
cached_tokens=self._cached_tokens,
|
|
model=self._model,
|
|
)
|
|
|
|
async def list_models(self) -> list[Any]:
|
|
return []
|
|
|
|
|
|
async def test_agent_turn_end_estimate_is_cache_aware(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setenv("OPENSQUILLA_OPENROUTER_LIVE_PRICING", "0")
|
|
provider = _CachedUsageProvider(
|
|
input_tokens=1_000_000,
|
|
output_tokens=0,
|
|
cached_tokens=900_000,
|
|
model="deepseek/deepseek-v4-pro",
|
|
)
|
|
agent = Agent(
|
|
provider=provider,
|
|
config=AgentConfig(model_id="deepseek/deepseek-v4-pro", provider_id="deepseek"),
|
|
)
|
|
events = [event async for event in agent.run_turn("hello")]
|
|
done = next(event for event in events if event.kind == "done")
|
|
# fresh 100,000 * 0.435 + read 900,000 * 0.003625 (per M) = 0.0467625
|
|
assert done.cost_usd == pytest.approx(0.0467625, abs=1e-6)
|
|
assert done.cost_source == "opensquilla_static_estimate"
|
|
assert done.estimate_basis == "cache_aware"
|
|
|
|
|
|
async def test_agent_turn_end_billed_cost_has_null_basis(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
monkeypatch.setenv("OPENSQUILLA_OPENROUTER_LIVE_PRICING", "0")
|
|
|
|
class _BilledProvider(_CachedUsageProvider):
|
|
async def _stream(self) -> AsyncIterator[Any]:
|
|
yield ProviderText(text="done")
|
|
yield ProviderDone(
|
|
stop_reason="stop",
|
|
input_tokens=self._input_tokens,
|
|
output_tokens=self._output_tokens,
|
|
billed_cost=0.02,
|
|
model=self._model,
|
|
)
|
|
|
|
agent = Agent(
|
|
provider=_BilledProvider(input_tokens=1000, output_tokens=10, model="claude-opus-4-7"),
|
|
config=AgentConfig(model_id="claude-opus-4-7"),
|
|
)
|
|
events = [event async for event in agent.run_turn("hello")]
|
|
done = next(event for event in events if event.kind == "done")
|
|
assert done.cost_source == "provider_billed"
|
|
assert done.cost_usd == pytest.approx(0.02)
|
|
assert done.estimate_basis is None
|
|
|
|
|
|
def test_turn_usage_payload_carries_estimate_basis() -> None:
|
|
done = DoneEvent(input_tokens=5, output_tokens=3, estimate_basis="cache_aware")
|
|
payload = _turn_usage_payload(done, resolved_model="m")
|
|
assert payload is not None
|
|
assert payload["estimate_basis"] == "cache_aware"
|
|
|
|
legacy = _turn_usage_payload(DoneEvent(input_tokens=5, output_tokens=3), resolved_model="m")
|
|
assert legacy is not None
|
|
assert legacy["estimate_basis"] is None
|
|
|
|
|
|
def test_turn_usage_payload_carries_decision_id() -> None:
|
|
"""The feedback loop's client entry point: decisionId on the wire."""
|
|
done = DoneEvent(input_tokens=5, output_tokens=3, decision_id="a" * 32)
|
|
payload = _turn_usage_payload(done, resolved_model="m")
|
|
assert payload is not None
|
|
assert payload["decision_id"] == "a" * 32
|
|
|
|
legacy = _turn_usage_payload(DoneEvent(input_tokens=5, output_tokens=3), resolved_model="m")
|
|
assert legacy is not None
|
|
assert legacy["decision_id"] is None
|