Implement LLM-WP-0008: provider-scoped account balance CLI
Add pluggable balance registry, OpenRouter credits/key limit client, and llm-connect balance with one-shot --provider that does not change library defaults.
This commit is contained in:
parent
c0d5c4ad08
commit
ed7c632155
11 changed files with 783 additions and 14 deletions
141
tests/test_balance.py
Normal file
141
tests/test_balance.py
Normal file
|
|
@ -0,0 +1,141 @@
|
|||
import pytest
|
||||
|
||||
from llm_connect.balance import (
|
||||
AccountBalance,
|
||||
BalanceClientRegistry,
|
||||
OpenRouterBalanceClient,
|
||||
format_balance_human,
|
||||
get_account_balance,
|
||||
resolve_balance_provider,
|
||||
)
|
||||
from llm_connect.config import LLMConfig
|
||||
from llm_connect.exceptions import LLMAPIError, LLMBalanceUnsupportedError, LLMConfigurationError
|
||||
from llm_connect.openrouter import _DEFAULT_MODEL as OPENROUTER_DEFAULT_MODEL
|
||||
|
||||
|
||||
def test_resolve_balance_provider_explicit_and_default():
|
||||
assert resolve_balance_provider("OpenRouter") == "openrouter"
|
||||
assert resolve_balance_provider(None) == LLMConfig().provider
|
||||
assert resolve_balance_provider("") == LLMConfig().provider
|
||||
assert resolve_balance_provider(" gemini ") == "gemini"
|
||||
|
||||
|
||||
def test_resolve_balance_provider_does_not_mutate_defaults():
|
||||
before_provider = LLMConfig().provider
|
||||
before_model = OPENROUTER_DEFAULT_MODEL
|
||||
resolve_balance_provider("gemini")
|
||||
resolve_balance_provider("openai")
|
||||
assert LLMConfig().provider == before_provider
|
||||
assert OPENROUTER_DEFAULT_MODEL == before_model
|
||||
|
||||
|
||||
def test_registry_unsupported_provider():
|
||||
registry = BalanceClientRegistry.default()
|
||||
with pytest.raises(LLMBalanceUnsupportedError, match="gemini") as exc_info:
|
||||
registry.create("gemini")
|
||||
assert "openrouter" in exc_info.value.supported
|
||||
|
||||
|
||||
def test_openrouter_balance_client_maps_credits_and_key_limit():
|
||||
def fake_get(url, headers=None, timeout=30):
|
||||
if url.endswith("/credits"):
|
||||
return {"data": {"total_credits": 25.0, "total_usage": 12.5}}
|
||||
if url.endswith("/auth/key"):
|
||||
return {
|
||||
"data": {
|
||||
"limit": 10.0,
|
||||
"limit_remaining": 9.5,
|
||||
"limit_reset": "monthly",
|
||||
"usage": 1.5,
|
||||
}
|
||||
}
|
||||
raise AssertionError(f"unexpected url {url}")
|
||||
|
||||
client = OpenRouterBalanceClient(
|
||||
api_key="sk-test",
|
||||
get_json_fn=fake_get,
|
||||
)
|
||||
balance = client.get_balance(fx=1.0)
|
||||
|
||||
assert balance.provider == "openrouter"
|
||||
assert balance.credits_remaining == pytest.approx(12.5)
|
||||
assert balance.credits_total == pytest.approx(25.0)
|
||||
assert balance.credits_used == pytest.approx(12.5)
|
||||
assert balance.limit_remaining == pytest.approx(9.5)
|
||||
assert balance.limit == pytest.approx(10.0)
|
||||
assert balance.limit_reset == "monthly"
|
||||
assert balance.credits_remaining_eur == pytest.approx(12.5)
|
||||
assert balance.limit_remaining_eur == pytest.approx(9.5)
|
||||
assert "credits" in balance.source
|
||||
assert "auth/key" in balance.source
|
||||
|
||||
|
||||
def test_openrouter_partial_failure_still_returns_one_signal():
|
||||
def fake_get(url, headers=None, timeout=30):
|
||||
if url.endswith("/credits"):
|
||||
raise LLMAPIError("down", status_code=500)
|
||||
if url.endswith("/auth/key"):
|
||||
return {"data": {"limit": 5, "limit_remaining": 4, "limit_reset": "monthly"}}
|
||||
raise AssertionError(url)
|
||||
|
||||
balance = OpenRouterBalanceClient(api_key="sk-test", get_json_fn=fake_get).get_balance(
|
||||
fx=0.92
|
||||
)
|
||||
assert balance.credits_remaining is None
|
||||
assert balance.limit_remaining == pytest.approx(4.0)
|
||||
assert balance.limit_remaining_eur == pytest.approx(3.68)
|
||||
assert "partial" in balance.source
|
||||
|
||||
|
||||
def test_openrouter_both_endpoints_fail():
|
||||
def fake_get(url, headers=None, timeout=30):
|
||||
raise LLMAPIError("nope", status_code=503)
|
||||
|
||||
with pytest.raises(LLMAPIError, match="both"):
|
||||
OpenRouterBalanceClient(api_key="sk-test", get_json_fn=fake_get).get_balance()
|
||||
|
||||
|
||||
def test_openrouter_missing_key():
|
||||
client = OpenRouterBalanceClient(api_key=None, get_json_fn=lambda *a, **k: {})
|
||||
# Force empty key even if env has one
|
||||
client._api_key = None
|
||||
with pytest.raises(LLMConfigurationError, match="API key"):
|
||||
client.get_balance()
|
||||
|
||||
|
||||
def test_get_account_balance_uses_registry():
|
||||
class Stub:
|
||||
provider_id = "stub"
|
||||
|
||||
def get_balance(self, *, fx=None):
|
||||
return AccountBalance(
|
||||
provider="stub",
|
||||
currency="USD",
|
||||
source="stub:test",
|
||||
credits_remaining=1.0,
|
||||
)
|
||||
|
||||
registry = BalanceClientRegistry({"stub": lambda: Stub()})
|
||||
balance = get_account_balance("stub", registry=registry)
|
||||
assert balance.credits_remaining == 1.0
|
||||
|
||||
|
||||
def test_format_balance_human_mentions_both_signals():
|
||||
text = format_balance_human(
|
||||
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,
|
||||
limit_reset="monthly",
|
||||
fx_source="explicit",
|
||||
)
|
||||
)
|
||||
assert "account credits remaining" in text
|
||||
assert "key limit remaining" in text
|
||||
assert "EUR" in text
|
||||
assert "openrouter" in text
|
||||
Loading…
Add table
Add a link
Reference in a new issue