tegwick
ed7c632155
Add pluggable balance registry, OpenRouter credits/key limit client, and llm-connect balance with one-shot --provider that does not change library defaults.
238 lines
7.5 KiB
Python
238 lines
7.5 KiB
Python
import json
|
|
from datetime import datetime, timezone
|
|
from zoneinfo import ZoneInfo
|
|
|
|
from llm_connect.cli import main
|
|
from llm_connect.quality import QualityLedger, QualityObservation
|
|
from llm_connect.usage import UsageEvent, UsageLedger
|
|
|
|
|
|
def test_rates_show_json_outputs_default_registry(capsys):
|
|
assert main(["rates", "show", "--json"]) == 0
|
|
|
|
payload = json.loads(capsys.readouterr().out)
|
|
|
|
assert payload["openai/gpt-4o-mini"]["prompt_per_1k"] == 0.00015
|
|
assert payload["moonshotai/kimi-k3"]["prompt_per_1k"] == 0.003
|
|
|
|
|
|
def test_classes_show_lists_builtins(capsys):
|
|
assert main(["classes", "show"]) == 0
|
|
|
|
output = capsys.readouterr().out
|
|
|
|
assert "chunk-summarization" in output
|
|
assert "entity-extraction" in output
|
|
|
|
|
|
def test_classes_fit_reads_quality_ledger(tmp_path, capsys):
|
|
ledger = QualityLedger(tmp_path / "quality.jsonl")
|
|
for _ in range(3):
|
|
ledger.append(
|
|
QualityObservation(
|
|
task_type="extract",
|
|
adapter_id="openrouter",
|
|
model_id="openai/gpt-4o-mini",
|
|
cost_usd=0.001,
|
|
quality_score=0.9,
|
|
latency_ms=100,
|
|
tokens_in=500,
|
|
tokens_out=350,
|
|
recorded_at=datetime(2026, 5, 19, tzinfo=timezone.utc),
|
|
tags={
|
|
"problem_class": "entity-extraction",
|
|
"dimensions": {
|
|
"chunk_words": 300,
|
|
"template_words": 100,
|
|
"expected_entities": 5,
|
|
},
|
|
},
|
|
)
|
|
)
|
|
|
|
assert main(["classes", "fit", str(ledger.path), "--class", "entity-extraction", "--json"]) == 0
|
|
|
|
payload = json.loads(capsys.readouterr().out)
|
|
|
|
assert payload["entity-extraction"]["params"]["tokens_per_entity"] == 70
|
|
|
|
|
|
def test_run_with_mock_reports_cost_and_writes_ledger(tmp_path, capsys):
|
|
ledger = tmp_path / "usage.jsonl"
|
|
assert (
|
|
main(
|
|
[
|
|
"run",
|
|
"hello world",
|
|
"--provider",
|
|
"mock",
|
|
"--model",
|
|
"moonshotai/kimi-k3",
|
|
"--ledger",
|
|
str(ledger),
|
|
"--eur-per-usd",
|
|
"1.0",
|
|
"--json",
|
|
]
|
|
)
|
|
== 0
|
|
)
|
|
|
|
payload = json.loads(capsys.readouterr().out)
|
|
assert payload["content"]
|
|
assert payload["usage"]["total_tokens"] > 0
|
|
assert payload["cost_usd"] is not None
|
|
assert payload["cost_eur"] is not None
|
|
events = UsageLedger(ledger).read_all()
|
|
assert len(events) == 1
|
|
assert events[0].source == "cli"
|
|
assert events[0].model_id == "moonshotai/kimi-k3"
|
|
|
|
|
|
def test_cost_estimate_cli(capsys):
|
|
assert (
|
|
main(
|
|
[
|
|
"cost",
|
|
"estimate",
|
|
"--model",
|
|
"moonshotai/kimi-k3",
|
|
"--prompt-tokens",
|
|
"1000",
|
|
"--completion-tokens",
|
|
"1000",
|
|
"--eur-per-usd",
|
|
"1",
|
|
"--json",
|
|
]
|
|
)
|
|
== 0
|
|
)
|
|
payload = json.loads(capsys.readouterr().out)
|
|
assert payload["cost_usd"] == 0.018
|
|
assert payload["cost_eur"] == 0.018
|
|
|
|
|
|
def test_balance_cli_json(monkeypatch, capsys):
|
|
from llm_connect.balance import AccountBalance
|
|
|
|
def fake_get_balance(provider=None, *, registry=None, fx=None):
|
|
assert provider is None # default path
|
|
return AccountBalance(
|
|
provider="openrouter",
|
|
currency="USD",
|
|
source="openrouter:/api/v1/credits+auth/key",
|
|
credits_remaining=12.5,
|
|
credits_remaining_eur=11.5,
|
|
limit_remaining=9.5,
|
|
limit_remaining_eur=8.74,
|
|
limit=10.0,
|
|
limit_reset="monthly",
|
|
fx_source="explicit",
|
|
)
|
|
|
|
monkeypatch.setattr("llm_connect.cli.get_account_balance", fake_get_balance)
|
|
assert main(["balance", "--json"]) == 0
|
|
payload = json.loads(capsys.readouterr().out)
|
|
assert payload["provider"] == "openrouter"
|
|
assert payload["credits_remaining"] == 12.5
|
|
assert payload["limit_remaining"] == 9.5
|
|
|
|
|
|
def test_balance_cli_explicit_provider_one_shot(monkeypatch, capsys):
|
|
from llm_connect.balance import AccountBalance
|
|
from llm_connect.config import LLMConfig
|
|
from llm_connect.openrouter import _DEFAULT_MODEL
|
|
|
|
seen = {}
|
|
|
|
def fake_get_balance(provider=None, *, registry=None, fx=None):
|
|
seen["provider"] = provider
|
|
return AccountBalance(
|
|
provider="openrouter",
|
|
currency="USD",
|
|
source="test",
|
|
credits_remaining=1.0,
|
|
)
|
|
|
|
before_provider = LLMConfig().provider
|
|
before_model = _DEFAULT_MODEL
|
|
monkeypatch.setattr("llm_connect.cli.get_account_balance", fake_get_balance)
|
|
assert main(["balance", "--provider", "openrouter"]) == 0
|
|
assert seen["provider"] == "openrouter"
|
|
assert LLMConfig().provider == before_provider
|
|
assert _DEFAULT_MODEL == before_model
|
|
assert "account credits remaining" in capsys.readouterr().out
|
|
|
|
|
|
def test_balance_cli_unsupported_provider(monkeypatch, capsys):
|
|
from llm_connect.exceptions import LLMBalanceUnsupportedError
|
|
|
|
def boom(provider=None, *, registry=None, fx=None):
|
|
raise LLMBalanceUnsupportedError(
|
|
"not supported",
|
|
provider=provider or "",
|
|
supported=["openrouter"],
|
|
)
|
|
|
|
monkeypatch.setattr("llm_connect.cli.get_account_balance", boom)
|
|
assert main(["balance", "--provider", "gemini"]) == 2
|
|
err = capsys.readouterr().err
|
|
assert "not supported" in err
|
|
|
|
|
|
def test_spend_week_current_and_last(tmp_path, capsys, monkeypatch):
|
|
ledger_path = tmp_path / "usage.jsonl"
|
|
ledger = UsageLedger(ledger_path)
|
|
berlin = ZoneInfo("Europe/Berlin")
|
|
current_start = datetime(2026, 8, 3, 0, 0, tzinfo=berlin)
|
|
current_end = datetime(2026, 8, 5, 12, 0, tzinfo=berlin)
|
|
last_start = datetime(2026, 7, 27, 0, 0, tzinfo=berlin)
|
|
last_end = current_start
|
|
|
|
def _fake_week_window(which="current", *, now=None, tz=None):
|
|
if which == "last":
|
|
return last_start, last_end
|
|
return current_start, current_end
|
|
|
|
monkeypatch.setattr("llm_connect.cli.week_window", _fake_week_window)
|
|
|
|
ledger.append(
|
|
UsageEvent(
|
|
provider="mock",
|
|
model_id="moonshotai/kimi-k3",
|
|
prompt_tokens=100,
|
|
completion_tokens=20,
|
|
total_tokens=120,
|
|
cost_usd=0.01,
|
|
cost_eur=0.009,
|
|
cost_source="rate_table:moonshotai/kimi-k3",
|
|
source="cli",
|
|
recorded_at=datetime(2026, 8, 4, 9, 0, tzinfo=berlin),
|
|
)
|
|
)
|
|
ledger.append(
|
|
UsageEvent(
|
|
provider="mock",
|
|
model_id="moonshotai/kimi-k3",
|
|
prompt_tokens=50,
|
|
completion_tokens=10,
|
|
total_tokens=60,
|
|
cost_usd=0.005,
|
|
cost_eur=0.0045,
|
|
cost_source="rate_table:moonshotai/kimi-k3",
|
|
source="cli",
|
|
recorded_at=datetime(2026, 7, 29, 9, 0, tzinfo=berlin), # previous week
|
|
)
|
|
)
|
|
|
|
assert main(["spend", "week", "--ledger", str(ledger_path), "--json"]) == 0
|
|
current = json.loads(capsys.readouterr().out)
|
|
assert current["event_count"] == 1
|
|
assert current["total_tokens"] == 120
|
|
assert current["cost_eur"] == 0.009
|
|
|
|
assert main(["spend", "week", "--last", "--ledger", str(ledger_path), "--json"]) == 0
|
|
last = json.loads(capsys.readouterr().out)
|
|
assert last["event_count"] == 1
|
|
assert last["total_tokens"] == 60
|