okx-codex-trader/okx_codex_trader/okx_client.py

import base64
import hashlib
import hmac
import json
from datetime import UTC, datetime
from decimal import Decimal, InvalidOperation, ROUND_DOWN
from math import isfinite
from typing import TypeAlias
from urllib.parse import urlencode

from okx_codex_trader.config import Config
from okx_codex_trader.models import Candle, InstrumentMeta, OrderResult, Position, TradeSignal


OkxRow: TypeAlias = dict[str, object] | list[object]


def _parse_finite_decimal(value: object) -> Decimal:
    if isinstance(value, bool):
        raise ValueError("contract sizing inputs are invalid")
    try:
        parsed = Decimal(str(value))
    except (InvalidOperation, TypeError, ValueError):
        raise ValueError("contract sizing inputs are invalid") from None
    if not parsed.is_finite():
        raise ValueError("contract sizing inputs are invalid")
    return parsed


def _parse_finite_float(value: object) -> float:
    if isinstance(value, bool):
        raise ValueError("okx response payload is invalid")
    try:
        parsed = float(value)
    except (TypeError, ValueError):
        raise ValueError("okx response payload is invalid") from None
    if not isfinite(parsed):
        raise ValueError("okx response payload is invalid")
    return parsed


def _parse_valid_leverage(value: object) -> int:
    if isinstance(value, bool) or not isinstance(value, int):
        raise ValueError("leverage is invalid")
    if value < 1 or value > 3:
        raise ValueError("leverage is invalid")
    return value


def build_contract_size(notional: float, price: float, metadata: InstrumentMeta) -> float:
    notional_decimal = _parse_finite_decimal(notional)
    price_decimal = _parse_finite_decimal(price)
    ct_val_decimal = _parse_finite_decimal(metadata.ct_val)
    lot_size = _parse_finite_decimal(metadata.lot_sz)
    min_size = _parse_finite_decimal(metadata.min_sz)
    if notional_decimal <= 0 or price_decimal <= 0 or ct_val_decimal <= 0 or lot_size <= 0 or min_size <= 0:
        raise ValueError("contract sizing inputs are invalid")
    raw_size = notional_decimal / (price_decimal * ct_val_decimal)
    size = (raw_size / lot_size).to_integral_value(rounding=ROUND_DOWN) * lot_size
    if size < min_size:
        raise ValueError("contract size below minimum")
    return float(size)


def _format_number(value: float) -> str:
    return format(Decimal(str(value)).normalize(), "f")


class OkxClient:
    base_url = "https://www.okx.com"
    request_timeout = 10.0

    def __init__(self, config: Config, session=None):
        self.config = config
        if session is None:
            import requests

            session = requests.Session()
        self.session = session

    def _invalid_payload(self) -> ValueError:
        return ValueError("okx response payload is invalid")

    def _transport_error(self) -> ValueError:
        return ValueError("okx transport error")

    def _first_item(self, data: list[OkxRow]) -> dict[str, object]:
        if not data:
            raise self._invalid_payload()
        item = data[0]
        if not isinstance(item, dict):
            raise self._invalid_payload()
        return item

    def _request(
        self,
        method: str,
        path: str,
        *,
        params: dict[str, object] | None = None,
        json_body: dict[str, object] | None = None,
    ) -> list[OkxRow]:
        timestamp = datetime.now(UTC).isoformat(timespec="milliseconds").replace("+00:00", "Z")
        query = urlencode(params or {})
        path_with_query = path if not query else f"{path}?{query}"
        body = "" if json_body is None else json.dumps(json_body, separators=(",", ":"))
        signature = base64.b64encode(
            hmac.new(
                self.config.api_secret.encode(),
                f"{timestamp}{method.upper()}{path_with_query}{body}".encode(),
                hashlib.sha256,
            ).digest()
        ).decode()
        headers = {
            "OK-ACCESS-KEY": self.config.api_key,
            "OK-ACCESS-SIGN": signature,
            "OK-ACCESS-TIMESTAMP": timestamp,
            "OK-ACCESS-PASSPHRASE": self.config.api_passphrase,
            "x-simulated-trading": "1",
        }
        if json_body is not None:
            headers["Content-Type"] = "application/json"
        try:
            response = self.session.request(
                method.upper(),
                f"{self.base_url}{path}",
                headers=headers,
                params=params,
                data=body if json_body is not None else None,
                timeout=self.request_timeout,
            )
        except Exception:
            raise self._transport_error() from None
        try:
            payload = response.json()
        except Exception:
            raise self._invalid_payload() from None
        if not isinstance(payload, dict):
            raise self._invalid_payload()
        if getattr(response, "status_code", 200) >= 400:
            raise ValueError(str(payload.get("msg") or "okx http error"))
        if payload.get("code") != "0":
            raise ValueError(str(payload.get("msg") or payload.get("code") or "okx api error"))
        data = payload.get("data")
        if not isinstance(data, list):
            raise self._invalid_payload()
        return data

    def get_candles(self, symbol: str, bar: str, limit: int) -> list[Candle]:
        data = self._request(
            "GET",
            "/api/v5/market/history-candles",
            params={"instId": symbol, "bar": bar, "limit": limit},
        )
        try:
            candles = [
                Candle(
                    symbol=symbol,
                    ts=int(entry[0]),
                    open=_parse_finite_float(entry[1]),
                    high=_parse_finite_float(entry[2]),
                    low=_parse_finite_float(entry[3]),
                    close=_parse_finite_float(entry[4]),
                    volume=_parse_finite_float(entry[5]),
                )
                for entry in data
            ]
            return sorted(candles, key=lambda candle: candle.ts)
        except (IndexError, KeyError, TypeError, ValueError):
            raise self._invalid_payload() from None

    def get_instrument_meta(self, symbol: str) -> InstrumentMeta:
        data = self._request(
            "GET",
            "/api/v5/public/instruments",
            params={"instType": "SWAP", "instId": symbol},
        )
        instrument = self._first_item(data)
        try:
            if instrument.get("instId") != symbol or instrument.get("instType") != "SWAP":
                raise self._invalid_payload()
            return InstrumentMeta(
                ct_val=_parse_finite_float(instrument["ctVal"]),
                lot_sz=_parse_finite_float(instrument["lotSz"]),
                min_sz=_parse_finite_float(instrument["minSz"]),
            )
        except (KeyError, TypeError, ValueError):
            raise self._invalid_payload() from None

    def get_last_price(self, symbol: str) -> float:
        data = self._request("GET", "/api/v5/market/ticker", params={"instId": symbol})
        ticker = self._first_item(data)
        try:
            if ticker.get("instId") != symbol:
                raise self._invalid_payload()
            return _parse_finite_float(ticker["last"])
        except (KeyError, TypeError, ValueError):
            raise self._invalid_payload() from None

    def ensure_hedge_mode(self) -> None:
        data = self._request("GET", "/api/v5/account/config")
        config = self._first_item(data)
        pos_mode = config.get("posMode")
        if not isinstance(pos_mode, str):
            raise self._invalid_payload()
        if pos_mode != "long_short_mode":
            raise ValueError("hedge mode is required")

    def set_leverage(self, symbol: str, leverage: int, pos_side: str) -> None:
        if not symbol.endswith("-SWAP"):
            raise ValueError("swap instrument is required")
        leverage = _parse_valid_leverage(leverage)
        if pos_side not in {"long", "short"}:
            raise ValueError("pos_side is invalid")
        self._request(
            "POST",
            "/api/v5/account/set-leverage",
            json_body={
                "instId": symbol,
                "lever": str(leverage),
                "mgnMode": "isolated",
                "posSide": pos_side,
            },
        )

    def place_demo_order(self, symbol: str, signal: TradeSignal, margin_usdt: float) -> OrderResult:
        if signal.action == "flat":
            return OrderResult(
                status="noop",
                order_id=None,
                symbol=symbol,
                side=None,
                pos_side=None,
                order_type=None,
                size=None,
            )
        if signal.action not in {"long", "short"}:
            raise ValueError("action is invalid")
        if not symbol.endswith("-SWAP"):
            raise ValueError("swap instrument is required")
        leverage = _parse_valid_leverage(signal.leverage)
        try:
            margin_value = _parse_finite_float(margin_usdt)
            margin_decimal = _parse_finite_decimal(margin_usdt)
        except ValueError:
            raise ValueError("margin_usdt is invalid") from None
        if margin_value <= 0:
            raise ValueError("margin_usdt is invalid")
        self.ensure_hedge_mode()
        metadata = self.get_instrument_meta(symbol)
        price = signal.entry_price if signal.entry_price is not None else self.get_last_price(symbol)
        side = "buy" if signal.action == "long" else "sell"
        pos_side = "long" if signal.action == "long" else "short"
        size = build_contract_size(margin_decimal * leverage, price, metadata)
        self.set_leverage(symbol, leverage, pos_side)
        order_type = "market" if signal.entry_price is None else "limit"
        request_body = {
            "instId": symbol,
            "tdMode": "isolated",
            "side": side,
            "posSide": pos_side,
            "ordType": order_type,
            "sz": _format_number(size),
        }
        if signal.entry_price is not None:
            request_body["px"] = _format_number(signal.entry_price)
        data = self._request("POST", "/api/v5/trade/order", json_body=request_body)
        order = self._first_item(data)
        order_id = str(order.get("ordId") or "")
        if not order_id:
            raise self._invalid_payload()
        return OrderResult(
            status="placed",
            order_id=order_id,
            symbol=symbol,
            side=side,
            pos_side=pos_side,
            order_type=order_type,
            size=size,
        )

    def get_positions(self, symbol: str) -> list[Position]:
        requested_symbol = symbol
        data = self._request("GET", "/api/v5/account/positions", params={"instId": requested_symbol})
        if not data:
            return []
        try:
            positions = []
            for entry in data:
                size = _parse_finite_float(entry["pos"])
                if size == 0.0:
                    continue
                symbol = entry["instId"]
                pos_side = entry["posSide"]
                if not isinstance(symbol, str) or not isinstance(pos_side, str):
                    raise self._invalid_payload()
                if symbol != requested_symbol:
                    raise self._invalid_payload()
                if pos_side not in {"long", "short"}:
                    raise self._invalid_payload()
                positions.append(
                    Position(
                        symbol=symbol,
                        pos_side=pos_side,
                        size=size,
                        avg_price=_parse_finite_float(entry["avgPx"]),
                    )
                )
            return positions
        except (KeyError, TypeError, ValueError):
            raise self._invalid_payload() from None