"""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())