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.
357 lines
12 KiB
Python
357 lines
12 KiB
Python
"""Provider-scoped account balance / prepaid remaining.
|
|
|
|
Balance lookup is always bound to an explicit or default **backend**
|
|
(provider). Selecting a provider for a balance query must never mutate
|
|
library defaults (model or provider).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Callable, Mapping
|
|
from dataclasses import dataclass, field
|
|
from typing import Any, Protocol, runtime_checkable
|
|
|
|
from llm_connect._http import get_json
|
|
from llm_connect.config import LLMConfig, find_project_root, resolve_api_key
|
|
from llm_connect.exceptions import (
|
|
LLMAPIError,
|
|
LLMBalanceUnsupportedError,
|
|
LLMConfigurationError,
|
|
)
|
|
from llm_connect.fx import FxRate, resolve_fx_rate
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AccountBalance:
|
|
"""Normalised prepaid / key-limit snapshot for one backend."""
|
|
|
|
provider: str
|
|
currency: str
|
|
source: str
|
|
credits_total: float | None = None
|
|
credits_used: float | None = None
|
|
credits_remaining: float | None = None
|
|
limit: float | None = None
|
|
limit_remaining: float | None = None
|
|
limit_reset: str | None = None
|
|
usage: float | None = None
|
|
credits_remaining_eur: float | None = None
|
|
limit_remaining_eur: float | None = None
|
|
fx_source: str | None = None
|
|
raw: dict[str, Any] = field(default_factory=dict)
|
|
|
|
def __post_init__(self) -> None:
|
|
provider = str(self.provider or "").strip()
|
|
currency = str(self.currency or "").strip().upper()
|
|
source = str(self.source or "").strip()
|
|
if not provider:
|
|
raise ValueError("provider must be a non-empty string")
|
|
if not currency:
|
|
raise ValueError("currency must be a non-empty string")
|
|
if not source:
|
|
raise ValueError("source must be a non-empty string")
|
|
object.__setattr__(self, "provider", provider)
|
|
object.__setattr__(self, "currency", currency)
|
|
object.__setattr__(self, "source", source)
|
|
object.__setattr__(self, "raw", dict(self.raw or {}))
|
|
|
|
def to_dict(self) -> dict[str, Any]:
|
|
"""JSON-serialisable view (includes raw for automation)."""
|
|
return {
|
|
"provider": self.provider,
|
|
"currency": self.currency,
|
|
"source": self.source,
|
|
"credits_total": self.credits_total,
|
|
"credits_used": self.credits_used,
|
|
"credits_remaining": self.credits_remaining,
|
|
"limit": self.limit,
|
|
"limit_remaining": self.limit_remaining,
|
|
"limit_reset": self.limit_reset,
|
|
"usage": self.usage,
|
|
"credits_remaining_eur": self.credits_remaining_eur,
|
|
"limit_remaining_eur": self.limit_remaining_eur,
|
|
"fx_source": self.fx_source,
|
|
"raw": dict(self.raw),
|
|
}
|
|
|
|
|
|
@runtime_checkable
|
|
class ProviderBalanceClient(Protocol):
|
|
"""Backend that can report account / key prepaid remaining."""
|
|
|
|
provider_id: str
|
|
|
|
def get_balance(self, *, fx: FxRate | float | None = None) -> AccountBalance:
|
|
"""Fetch a live balance snapshot for this provider."""
|
|
...
|
|
|
|
|
|
BalanceClientFactory = Callable[[], ProviderBalanceClient]
|
|
|
|
|
|
def resolve_balance_provider(explicit: str | None = None) -> str:
|
|
"""Resolve which backend to query for balance.
|
|
|
|
1. Non-empty *explicit* provider id (one-shot override)
|
|
2. Else :attr:`LLMConfig.provider` default (today ``openrouter``)
|
|
|
|
Read-only: never mutates config, env, or module defaults.
|
|
"""
|
|
if explicit is not None and str(explicit).strip():
|
|
return str(explicit).strip().lower()
|
|
return str(LLMConfig().provider).strip().lower()
|
|
|
|
|
|
class BalanceClientRegistry:
|
|
"""Map provider id → balance client factory."""
|
|
|
|
def __init__(
|
|
self,
|
|
factories: Mapping[str, BalanceClientFactory] | None = None,
|
|
) -> None:
|
|
self._factories: dict[str, BalanceClientFactory] = {
|
|
str(key).strip().lower(): factory
|
|
for key, factory in (factories or {}).items()
|
|
}
|
|
|
|
def register(self, provider_id: str, factory: BalanceClientFactory) -> None:
|
|
"""Register or replace a factory for *provider_id*."""
|
|
key = str(provider_id).strip().lower()
|
|
if not key:
|
|
raise ValueError("provider_id must be a non-empty string")
|
|
self._factories[key] = factory
|
|
|
|
def supported(self) -> list[str]:
|
|
"""Return sorted provider ids that implement balance."""
|
|
return sorted(self._factories)
|
|
|
|
def create(self, provider_id: str) -> ProviderBalanceClient:
|
|
"""Instantiate a client for *provider_id* or raise unsupported."""
|
|
key = str(provider_id).strip().lower()
|
|
factory = self._factories.get(key)
|
|
if factory is None:
|
|
supported = self.supported()
|
|
raise LLMBalanceUnsupportedError(
|
|
f"Account balance is not supported for provider {key!r}. "
|
|
f"Supported: {', '.join(supported) if supported else '(none)'}",
|
|
provider=key,
|
|
supported=supported,
|
|
)
|
|
return factory()
|
|
|
|
@classmethod
|
|
def default(cls) -> "BalanceClientRegistry":
|
|
"""Built-in registry (OpenRouter first; more backends later)."""
|
|
return cls(
|
|
{
|
|
"openrouter": lambda: OpenRouterBalanceClient(),
|
|
}
|
|
)
|
|
|
|
|
|
def get_account_balance(
|
|
provider: str | None = None,
|
|
*,
|
|
registry: BalanceClientRegistry | None = None,
|
|
fx: FxRate | float | None = None,
|
|
) -> AccountBalance:
|
|
"""Fetch balance for *provider* or the current default backend.
|
|
|
|
*provider* is a one-shot selector and does not change library defaults.
|
|
"""
|
|
resolved = resolve_balance_provider(provider)
|
|
clients = registry or BalanceClientRegistry.default()
|
|
return clients.create(resolved).get_balance(fx=fx)
|
|
|
|
|
|
class OpenRouterBalanceClient:
|
|
"""OpenRouter prepaid credits + API-key limit remaining."""
|
|
|
|
provider_id = "openrouter"
|
|
|
|
def __init__(
|
|
self,
|
|
api_key: str | None = None,
|
|
api_base: str = "https://openrouter.ai/api/v1",
|
|
*,
|
|
timeout: int = 30,
|
|
get_json_fn: Callable[..., dict[str, Any]] | None = None,
|
|
) -> None:
|
|
self._api_base = api_base.rstrip("/")
|
|
self._timeout = timeout
|
|
self._get_json = get_json_fn or get_json
|
|
root = find_project_root()
|
|
key_file_paths = [root / "apikey-openrouter.txt"] if root else []
|
|
self._api_key = resolve_api_key(
|
|
explicit=api_key,
|
|
env_var="OPENROUTER_API_KEY",
|
|
key_file_paths=key_file_paths,
|
|
)
|
|
|
|
def get_balance(self, *, fx: FxRate | float | None = None) -> AccountBalance:
|
|
if not self._api_key:
|
|
raise LLMConfigurationError(
|
|
"OpenRouter API key not found for balance lookup "
|
|
"(set OPENROUTER_API_KEY or apikey-openrouter.txt)",
|
|
context={"provider": self.provider_id},
|
|
)
|
|
|
|
headers = {"Authorization": f"Bearer {self._api_key}"}
|
|
credits_payload: dict[str, Any] | None = None
|
|
key_payload: dict[str, Any] | None = None
|
|
errors: list[str] = []
|
|
sources: list[str] = []
|
|
|
|
try:
|
|
credits_payload = self._get_json(
|
|
f"{self._api_base}/credits",
|
|
headers=headers,
|
|
timeout=self._timeout,
|
|
)
|
|
sources.append("credits")
|
|
except LLMAPIError as exc:
|
|
errors.append(f"credits: {exc}")
|
|
|
|
try:
|
|
key_payload = self._get_json(
|
|
f"{self._api_base}/auth/key",
|
|
headers=headers,
|
|
timeout=self._timeout,
|
|
)
|
|
sources.append("auth/key")
|
|
except LLMAPIError as exc:
|
|
errors.append(f"auth/key: {exc}")
|
|
|
|
if credits_payload is None and key_payload is None:
|
|
raise LLMAPIError(
|
|
"OpenRouter balance lookup failed for both /credits and /auth/key: "
|
|
+ "; ".join(errors),
|
|
context={"provider": self.provider_id, "errors": errors},
|
|
)
|
|
|
|
credits_total: float | None = None
|
|
credits_used: float | None = None
|
|
credits_remaining: float | None = None
|
|
limit: float | None = None
|
|
limit_remaining: float | None = None
|
|
limit_reset: str | None = None
|
|
usage: float | None = None
|
|
|
|
if credits_payload is not None:
|
|
data = _unwrap_data(credits_payload)
|
|
credits_total = _optional_float(data.get("total_credits"))
|
|
credits_used = _optional_float(data.get("total_usage"))
|
|
if credits_total is not None and credits_used is not None:
|
|
credits_remaining = credits_total - credits_used
|
|
|
|
if key_payload is not None:
|
|
data = _unwrap_data(key_payload)
|
|
limit = _optional_float(data.get("limit"))
|
|
limit_remaining = _optional_float(data.get("limit_remaining"))
|
|
usage = _optional_float(data.get("usage"))
|
|
raw_reset = data.get("limit_reset")
|
|
if raw_reset is not None and str(raw_reset).strip():
|
|
limit_reset = str(raw_reset).strip()
|
|
|
|
fx_rate = resolve_fx_rate(fx)
|
|
credits_remaining_eur = None
|
|
limit_remaining_eur = None
|
|
fx_source = None
|
|
if fx_rate is not None:
|
|
fx_source = fx_rate.source
|
|
if credits_remaining is not None:
|
|
credits_remaining_eur = fx_rate.usd_to_eur(credits_remaining)
|
|
if limit_remaining is not None:
|
|
limit_remaining_eur = fx_rate.usd_to_eur(limit_remaining)
|
|
|
|
source_label = "openrouter:" + "+".join(
|
|
f"/api/v1/{part}" for part in sources
|
|
)
|
|
if errors:
|
|
source_label += f" (partial: {'; '.join(errors)})"
|
|
|
|
return AccountBalance(
|
|
provider=self.provider_id,
|
|
currency="USD",
|
|
source=source_label,
|
|
credits_total=credits_total,
|
|
credits_used=credits_used,
|
|
credits_remaining=credits_remaining,
|
|
limit=limit,
|
|
limit_remaining=limit_remaining,
|
|
limit_reset=limit_reset,
|
|
usage=usage,
|
|
credits_remaining_eur=credits_remaining_eur,
|
|
limit_remaining_eur=limit_remaining_eur,
|
|
fx_source=fx_source,
|
|
raw={
|
|
"credits": credits_payload,
|
|
"auth_key": key_payload,
|
|
"errors": errors,
|
|
},
|
|
)
|
|
|
|
|
|
def format_balance_human(balance: AccountBalance) -> str:
|
|
"""Multi-line human-readable balance report."""
|
|
lines = [f"provider: {balance.provider}"]
|
|
|
|
if balance.credits_remaining is not None:
|
|
eur = _fmt_eur(balance.credits_remaining_eur)
|
|
lines.append(
|
|
f"account credits remaining: {_fmt_money(balance.credits_remaining)} "
|
|
f"{balance.currency}{eur}"
|
|
)
|
|
if balance.credits_total is not None and balance.credits_used is not None:
|
|
lines.append(
|
|
f" (total={_fmt_money(balance.credits_total)} "
|
|
f"used={_fmt_money(balance.credits_used)})"
|
|
)
|
|
else:
|
|
lines.append("account credits remaining: unavailable")
|
|
|
|
if balance.limit_remaining is not None:
|
|
eur = _fmt_eur(balance.limit_remaining_eur)
|
|
extra = []
|
|
if balance.limit is not None:
|
|
extra.append(f"limit={_fmt_money(balance.limit)}")
|
|
if balance.limit_reset:
|
|
extra.append(f"reset={balance.limit_reset}")
|
|
suffix = f" [{', '.join(extra)}]" if extra else ""
|
|
lines.append(
|
|
f"key limit remaining: {_fmt_money(balance.limit_remaining)} "
|
|
f"{balance.currency}{eur}{suffix}"
|
|
)
|
|
else:
|
|
lines.append("key limit remaining: unavailable")
|
|
|
|
if balance.fx_source:
|
|
lines.append(f"fx: {balance.fx_source}")
|
|
lines.append(f"source: {balance.source}")
|
|
return "\n".join(lines)
|
|
|
|
|
|
def _unwrap_data(payload: Mapping[str, Any]) -> Mapping[str, Any]:
|
|
data = payload.get("data", payload)
|
|
if isinstance(data, Mapping):
|
|
return data
|
|
return {}
|
|
|
|
|
|
def _optional_float(value: Any) -> float | None:
|
|
if value is None:
|
|
return None
|
|
try:
|
|
return float(value)
|
|
except (TypeError, ValueError):
|
|
return None
|
|
|
|
|
|
def _fmt_money(value: float) -> str:
|
|
return f"{value:.6f}".rstrip("0").rstrip(".") if value != 0 else "0"
|
|
|
|
|
|
def _fmt_eur(value: float | None) -> str:
|
|
if value is None:
|
|
return ""
|
|
return f" (~{_fmt_money(value)} EUR)"
|