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.
410 lines
15 KiB
Python
410 lines
15 KiB
Python
"""Command-line helpers for llm-connect registries, run, and spend."""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import json
|
||
import sys
|
||
from collections.abc import Iterable, Mapping
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from llm_connect.balance import (
|
||
format_balance_human,
|
||
get_account_balance,
|
||
)
|
||
from llm_connect.costs import estimate_cost
|
||
from llm_connect.exceptions import LLMBalanceUnsupportedError, LLMConfigurationError, LLMError
|
||
from llm_connect.factory import create_adapter
|
||
from llm_connect.fx import resolve_fx_rate
|
||
from llm_connect.models import RunConfig
|
||
from llm_connect.problem_classes import ProblemClass, ProblemClassRegistry
|
||
from llm_connect.quality import QualityLedger
|
||
from llm_connect.rates import ModelRateRegistry
|
||
from llm_connect.usage import (
|
||
UsageLedger,
|
||
default_usage_ledger_path,
|
||
event_from_response,
|
||
format_cost_footer,
|
||
suppress_auto_usage_record,
|
||
week_window,
|
||
)
|
||
|
||
|
||
def main(argv: list[str] | None = None) -> int:
|
||
"""Run the ``llm-connect`` command."""
|
||
parser = _build_parser()
|
||
args = parser.parse_args(argv)
|
||
return int(args.func(args))
|
||
|
||
|
||
def _build_parser() -> argparse.ArgumentParser:
|
||
parser = argparse.ArgumentParser(prog="llm-connect")
|
||
commands = parser.add_subparsers(dest="command", required=True)
|
||
|
||
rates = commands.add_parser("rates", help="Inspect model rate registries")
|
||
rate_commands = rates.add_subparsers(dest="rates_command", required=True)
|
||
rate_show = rate_commands.add_parser("show", help="Show model rates")
|
||
rate_show.add_argument("--rates", type=Path, help="YAML registry overlay")
|
||
rate_show.add_argument("--json", action="store_true", help="Emit JSON")
|
||
rate_show.set_defaults(func=_rates_show)
|
||
|
||
classes = commands.add_parser("classes", help="Inspect problem classes")
|
||
class_commands = classes.add_subparsers(dest="classes_command", required=True)
|
||
class_show = class_commands.add_parser("show", help="Show problem classes")
|
||
class_show.add_argument("--json", action="store_true", help="Emit JSON")
|
||
class_show.set_defaults(func=_classes_show)
|
||
|
||
class_fit = class_commands.add_parser("fit", help="Fit problem-class params from a ledger")
|
||
class_fit.add_argument("ledger", type=Path, help="QualityLedger JSONL path")
|
||
class_fit.add_argument("--class", dest="class_name", help="Fit one class by name")
|
||
class_fit.add_argument("--min-observations", type=int, default=3)
|
||
class_fit.add_argument("--json", action="store_true", help="Emit JSON")
|
||
class_fit.set_defaults(func=_classes_fit)
|
||
|
||
run = commands.add_parser("run", help="Execute a prompt and report tokens + cost")
|
||
run.add_argument("prompt", help="Prompt text to send")
|
||
run.add_argument("--provider", default="openrouter", help="Provider key (default openrouter)")
|
||
run.add_argument("--model", default=None, help="Model id (provider default if omitted)")
|
||
run.add_argument("--temperature", type=float, default=0.7)
|
||
run.add_argument("--max-tokens", type=int, default=2000)
|
||
run.add_argument("--timeout", type=int, default=300, help="Timeout seconds")
|
||
run.add_argument("--ledger", type=Path, default=None, help="Usage ledger JSONL path")
|
||
run.add_argument("--no-ledger", action="store_true", help="Do not append to the usage ledger")
|
||
run.add_argument("--json", action="store_true", help="Emit JSON response + cost")
|
||
run.add_argument(
|
||
"--eur-per-usd",
|
||
type=float,
|
||
default=None,
|
||
help="Override FX rate (euros per one USD)",
|
||
)
|
||
run.set_defaults(func=_run_prompt)
|
||
|
||
cost = commands.add_parser("cost", help="Cost estimation helpers")
|
||
cost_commands = cost.add_subparsers(dest="cost_command", required=True)
|
||
cost_estimate = cost_commands.add_parser(
|
||
"estimate",
|
||
help="Estimate cost for token counts without calling a provider",
|
||
)
|
||
cost_estimate.add_argument("--model", required=True, help="Model id")
|
||
cost_estimate.add_argument("--prompt-tokens", type=int, required=True)
|
||
cost_estimate.add_argument("--completion-tokens", type=int, default=0)
|
||
cost_estimate.add_argument("--rates", type=Path, help="YAML registry overlay")
|
||
cost_estimate.add_argument("--eur-per-usd", type=float, default=None)
|
||
cost_estimate.add_argument("--json", action="store_true", help="Emit JSON")
|
||
cost_estimate.set_defaults(func=_cost_estimate)
|
||
|
||
spend = commands.add_parser("spend", help="Inspect recorded spend")
|
||
spend_commands = spend.add_subparsers(dest="spend_command", required=True)
|
||
spend_week = spend_commands.add_parser(
|
||
"week",
|
||
help="Weekly spend: current Mon→now, or --last for previous Mon–Sun",
|
||
)
|
||
spend_week.add_argument(
|
||
"--last",
|
||
action="store_true",
|
||
help="Previous calendar week (Mon 00:00 → Sun end / this Mon 00:00)",
|
||
)
|
||
spend_week.add_argument("--ledger", type=Path, default=None, help="Usage ledger JSONL path")
|
||
spend_week.add_argument("--tz", default=None, help="Timezone (default Europe/Berlin)")
|
||
spend_week.add_argument("--json", action="store_true", help="Emit JSON")
|
||
spend_week.set_defaults(func=_spend_week)
|
||
|
||
balance = commands.add_parser(
|
||
"balance",
|
||
help=(
|
||
"Show prepaid/account remaining for a backend "
|
||
"(default: current default provider; --provider is one-shot only)"
|
||
),
|
||
)
|
||
balance.add_argument(
|
||
"--provider",
|
||
default=None,
|
||
help=(
|
||
"Backend to query for this command only (does not change library defaults). "
|
||
"Omit to use the current default provider."
|
||
),
|
||
)
|
||
balance.add_argument("--json", action="store_true", help="Emit JSON")
|
||
balance.add_argument(
|
||
"--eur-per-usd",
|
||
type=float,
|
||
default=None,
|
||
help="Override FX rate (euros per one USD)",
|
||
)
|
||
balance.set_defaults(func=_balance_show)
|
||
|
||
return parser
|
||
|
||
|
||
def _rates_show(args: argparse.Namespace) -> int:
|
||
registry = ModelRateRegistry.default()
|
||
if args.rates:
|
||
registry = registry.merged_with(ModelRateRegistry.from_yaml(args.rates))
|
||
rates = registry.all()
|
||
if args.json:
|
||
print(
|
||
json.dumps(
|
||
{
|
||
model_id: {
|
||
"prompt_per_1k": rate.prompt_per_1k,
|
||
"completion_per_1k": rate.completion_per_1k,
|
||
"currency": rate.currency,
|
||
"source_url": rate.source_url,
|
||
"captured_at": rate.captured_at,
|
||
}
|
||
for model_id, rate in sorted(rates.items())
|
||
},
|
||
indent=2,
|
||
sort_keys=True,
|
||
)
|
||
)
|
||
return 0
|
||
|
||
print("model_id\tprompt_per_1k\tcompletion_per_1k\tcurrency\tcaptured_at")
|
||
for model_id, rate in sorted(rates.items()):
|
||
print(
|
||
f"{model_id}\t{rate.prompt_per_1k:g}\t{rate.completion_per_1k:g}\t"
|
||
f"{rate.currency}\t{rate.captured_at}"
|
||
)
|
||
return 0
|
||
|
||
|
||
def _classes_show(args: argparse.Namespace) -> int:
|
||
classes = ProblemClassRegistry.default().all()
|
||
if args.json:
|
||
print(json.dumps(_classes_payload(classes.values()), indent=2, sort_keys=True))
|
||
return 0
|
||
|
||
print("name\tdimensions\ttunable_params\tcurrent_params")
|
||
for problem_class in sorted(classes.values(), key=lambda item: item.name):
|
||
print(
|
||
f"{problem_class.name}\t{', '.join(problem_class.base_dimensions)}\t"
|
||
f"{', '.join(problem_class.tunable_params)}\t{_format_params(problem_class.params)}"
|
||
)
|
||
return 0
|
||
|
||
|
||
def _classes_fit(args: argparse.Namespace) -> int:
|
||
if args.min_observations <= 0:
|
||
raise SystemExit("--min-observations must be positive")
|
||
registry = ProblemClassRegistry.default()
|
||
classes = registry.all()
|
||
if args.class_name:
|
||
problem_class = registry.get(args.class_name)
|
||
if problem_class is None:
|
||
raise SystemExit(f"Unknown problem class: {args.class_name}")
|
||
selected: list[ProblemClass] = [problem_class]
|
||
else:
|
||
selected = list(classes.values())
|
||
|
||
observations = QualityLedger(args.ledger).read_all()
|
||
fitted: list[ProblemClass] = [
|
||
problem_class.fit(observations, min_observations=args.min_observations)
|
||
for problem_class in selected
|
||
]
|
||
if args.json:
|
||
print(json.dumps(_classes_payload(fitted), indent=2, sort_keys=True))
|
||
return 0
|
||
|
||
print("name\tfitted_params\tconfidence")
|
||
for problem_class in sorted(fitted, key=lambda item: item.name):
|
||
confidence = getattr(problem_class, "confidence", 0.5)
|
||
print(f"{problem_class.name}\t{_format_params(problem_class.params)}\t{confidence:g}")
|
||
return 0
|
||
|
||
|
||
def _run_prompt(args: argparse.Namespace) -> int:
|
||
fx = resolve_fx_rate(args.eur_per_usd)
|
||
adapter_kwargs: dict[str, Any] = {}
|
||
if args.model is not None:
|
||
adapter_kwargs["model"] = args.model
|
||
adapter = create_adapter(args.provider, **adapter_kwargs)
|
||
config = RunConfig(
|
||
model_name=args.model or "gpt-4",
|
||
temperature=args.temperature,
|
||
max_tokens=args.max_tokens,
|
||
timeout_seconds=args.timeout,
|
||
)
|
||
|
||
with suppress_auto_usage_record():
|
||
response = adapter.execute_prompt(args.prompt, config)
|
||
|
||
usage = response.usage or {}
|
||
prompt_tokens = int(usage.get("prompt_tokens") or 0)
|
||
completion_tokens = int(usage.get("completion_tokens") or 0)
|
||
total_tokens = int(usage.get("total_tokens") or (prompt_tokens + completion_tokens))
|
||
model_id = response.model or args.model or "unknown"
|
||
estimate = estimate_cost(
|
||
model_id,
|
||
prompt_tokens,
|
||
completion_tokens,
|
||
fx=fx,
|
||
)
|
||
|
||
if not args.no_ledger:
|
||
ledger_path = args.ledger or default_usage_ledger_path()
|
||
event = event_from_response(
|
||
response,
|
||
provider=args.provider,
|
||
source="cli",
|
||
fx=fx,
|
||
)
|
||
UsageLedger(ledger_path).append(event)
|
||
|
||
footer = format_cost_footer(
|
||
estimate,
|
||
prompt_tokens=prompt_tokens,
|
||
completion_tokens=completion_tokens,
|
||
total_tokens=total_tokens,
|
||
model_id=model_id,
|
||
)
|
||
|
||
if args.json:
|
||
print(
|
||
json.dumps(
|
||
{
|
||
"content": response.content,
|
||
"model": model_id,
|
||
"usage": {
|
||
"prompt_tokens": prompt_tokens,
|
||
"completion_tokens": completion_tokens,
|
||
"total_tokens": total_tokens,
|
||
},
|
||
"cost_usd": estimate.cost_usd,
|
||
"cost_eur": estimate.cost_eur,
|
||
"cost_source": estimate.cost_source,
|
||
"fx_source": estimate.fx_source,
|
||
"finish_reason": response.finish_reason,
|
||
},
|
||
indent=2,
|
||
sort_keys=True,
|
||
)
|
||
)
|
||
return 0
|
||
|
||
print(response.content)
|
||
print(footer, file=sys.stderr)
|
||
return 0
|
||
|
||
|
||
def _cost_estimate(args: argparse.Namespace) -> int:
|
||
if args.prompt_tokens < 0 or args.completion_tokens < 0:
|
||
raise SystemExit("token counts must be non-negative")
|
||
registry = ModelRateRegistry.default()
|
||
if args.rates:
|
||
registry = registry.merged_with(ModelRateRegistry.from_yaml(args.rates))
|
||
estimate = estimate_cost(
|
||
args.model,
|
||
args.prompt_tokens,
|
||
args.completion_tokens,
|
||
registry=registry,
|
||
fx=args.eur_per_usd,
|
||
)
|
||
total = args.prompt_tokens + args.completion_tokens
|
||
if args.json:
|
||
print(
|
||
json.dumps(
|
||
{
|
||
"model": args.model,
|
||
"prompt_tokens": args.prompt_tokens,
|
||
"completion_tokens": args.completion_tokens,
|
||
"total_tokens": total,
|
||
"cost_usd": estimate.cost_usd,
|
||
"cost_eur": estimate.cost_eur,
|
||
"cost_source": estimate.cost_source,
|
||
"fx_source": estimate.fx_source,
|
||
},
|
||
indent=2,
|
||
sort_keys=True,
|
||
)
|
||
)
|
||
return 0
|
||
|
||
print(
|
||
format_cost_footer(
|
||
estimate,
|
||
prompt_tokens=args.prompt_tokens,
|
||
completion_tokens=args.completion_tokens,
|
||
total_tokens=total,
|
||
model_id=args.model,
|
||
)
|
||
)
|
||
return 0
|
||
|
||
|
||
def _spend_week(args: argparse.Namespace) -> int:
|
||
which = "last" if args.last else "current"
|
||
start, end = week_window(which, tz=args.tz)
|
||
ledger_path = args.ledger or default_usage_ledger_path()
|
||
summary = UsageLedger(ledger_path).sum_range(start, end)
|
||
|
||
if args.json:
|
||
payload = summary.to_dict()
|
||
payload["which"] = which
|
||
payload["tz"] = args.tz or "Europe/Berlin"
|
||
payload["ledger"] = str(ledger_path)
|
||
print(json.dumps(payload, indent=2, sort_keys=True))
|
||
return 0
|
||
|
||
label = "last week (Mon–Sun)" if which == "last" else "current week (Mon–now)"
|
||
print(f"spend {label}")
|
||
print(f"window: {_fmt_local(start)} → {_fmt_local(end)}")
|
||
print(f"ledger: {ledger_path}")
|
||
print(f"events: {summary.event_count}")
|
||
print(
|
||
f"tokens: in={summary.prompt_tokens} out={summary.completion_tokens} "
|
||
f"total={summary.total_tokens}"
|
||
)
|
||
usd = "unknown" if summary.cost_usd is None else f"{summary.cost_usd:.6f}"
|
||
eur = "unknown" if summary.cost_eur is None else f"{summary.cost_eur:.6f}"
|
||
print(f"cost: eur={eur} usd={usd}")
|
||
if summary.unknown_cost_events:
|
||
print(f"unknown_cost_events: {summary.unknown_cost_events}")
|
||
return 0
|
||
|
||
|
||
def _fmt_local(dt: datetime) -> str:
|
||
return dt.isoformat()
|
||
|
||
|
||
def _balance_show(args: argparse.Namespace) -> int:
|
||
try:
|
||
balance = get_account_balance(args.provider, fx=args.eur_per_usd)
|
||
except LLMBalanceUnsupportedError as exc:
|
||
print(str(exc), file=sys.stderr)
|
||
return 2
|
||
except LLMConfigurationError as exc:
|
||
print(str(exc), file=sys.stderr)
|
||
return 2
|
||
except LLMError as exc:
|
||
print(str(exc), file=sys.stderr)
|
||
return 1
|
||
|
||
if args.json:
|
||
print(json.dumps(balance.to_dict(), indent=2, sort_keys=True))
|
||
return 0
|
||
|
||
print(format_balance_human(balance))
|
||
return 0
|
||
|
||
|
||
def _classes_payload(classes: Iterable[ProblemClass]) -> dict[str, dict[str, Any]]:
|
||
return {
|
||
problem_class.name: {
|
||
"base_dimensions": list(problem_class.base_dimensions),
|
||
"tunable_params": list(problem_class.tunable_params),
|
||
"params": dict(problem_class.params),
|
||
"confidence": getattr(problem_class, "confidence", 0.5),
|
||
}
|
||
for problem_class in sorted(classes, key=lambda item: item.name)
|
||
}
|
||
|
||
|
||
def _format_params(params: Mapping[str, float]) -> str:
|
||
return ", ".join(f"{key}={value:g}" for key, value in sorted(dict(params).items()))
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|