okx-codex-trader/okx_codex_trader/cli.py

import argparse
import json
from dataclasses import asdict
from pathlib import Path
from typing import Callable, Sequence

from okx_codex_trader.backtest import run_backtest
from okx_codex_trader.codex_analyzer import analyze_with_codex
from okx_codex_trader.config import Config, load_config
from okx_codex_trader.okx_client import OkxClient
from okx_codex_trader.strategy import validate_signal


SUPPORTED_SYMBOLS = ("BTC-USDT-SWAP", "ETH-USDT-SWAP")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(prog="okx-codex-trader")
    subparsers = parser.add_subparsers(dest="command", required=True)

    fetch_history = subparsers.add_parser("fetch-history")
    fetch_history.add_argument("--symbol", choices=SUPPORTED_SYMBOLS, required=True)
    fetch_history.add_argument("--bar", required=True)
    fetch_history.add_argument("--limit", type=int, required=True)

    backtest = subparsers.add_parser("backtest")
    backtest.add_argument("--symbol", choices=SUPPORTED_SYMBOLS, required=True)
    backtest.add_argument("--bar", required=True)
    backtest.add_argument("--limit", type=int, required=True)
    backtest.add_argument("--leverage", type=int, choices=(1, 2, 3), required=True)

    analyze = subparsers.add_parser("analyze")
    analyze.add_argument("--symbol", choices=SUPPORTED_SYMBOLS, required=True)
    analyze.add_argument("--bar", required=True)
    analyze.add_argument("--limit", type=int, required=True)
    analyze.add_argument("--output-file", required=True)

    paper_order = subparsers.add_parser("paper-order")
    paper_order.add_argument("--symbol", choices=SUPPORTED_SYMBOLS, required=True)
    paper_order.add_argument("--signal-file", required=True)
    paper_order.add_argument("--margin-usdt", type=float, required=True)

    positions = subparsers.add_parser("positions")
    positions.add_argument("--symbol", choices=SUPPORTED_SYMBOLS, required=True)

    return parser


def _write_text(path: str, text: str) -> None:
    Path(path).write_text(text)


def _dump_json(payload: object) -> str:
    return json.dumps(payload, indent=2)


def main_factory(
    *,
    load_config: Callable[[], Config] = load_config,
    client_factory: Callable[[Config], OkxClient] = OkxClient,
    analyze_fn: Callable = analyze_with_codex,
    write_text: Callable[[str, str], None] = _write_text,
):
    def main(argv: Sequence[str] | None = None) -> int:
        parser = build_parser()
        args = parser.parse_args(argv)
        config = load_config()
        client = client_factory(config)

        if args.command == "fetch-history":
            candles = client.get_candles(args.symbol, args.bar, args.limit)
            print(_dump_json([asdict(candle) for candle in candles]))
            return 0

        if args.command == "backtest":
            candles = client.get_candles(args.symbol, args.bar, args.limit)
            print(_dump_json(run_backtest(candles=candles, leverage=args.leverage).to_dict()))
            return 0

        if args.command == "analyze":
            candles = client.get_candles(args.symbol, args.bar, args.limit)
            signal = analyze_fn(candles=candles, symbol=args.symbol, bar=args.bar)
            output = _dump_json(asdict(signal))
            write_text(args.output_file, output)
            print(output)
            return 0

        if args.command == "paper-order":
            signal = validate_signal(json.loads(Path(args.signal_file).read_text()))
            print(_dump_json(asdict(client.place_demo_order(args.symbol, signal, args.margin_usdt))))
            return 0

        positions = client.get_positions(args.symbol)
        print(_dump_json([asdict(position) for position in positions]))
        return 0

    return main


main = main_factory()


if __name__ == "__main__":
    raise SystemExit(main())