llm-connect/llm_connect/cli.py

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