core/exchange.py

from pybit.unified_trading import HTTP
from dotenv import load_dotenv
import os
import yaml
import time
import logging

load_dotenv()
logger = logging.getLogger(__name__)

class BybitExchange:
    def __init__(self):
        self.settings = yaml.safe_load(open("config/settings.yaml"))

        API_KEY = os.getenv("BYBIT_API_KEY")
        API_SECRET = os.getenv("BYBIT_API_SECRET")
        TESTNET = os.getenv("BYBIT_TESTNET", "true").lower() == "true"

        if not API_KEY or not API_SECRET:
            raise Exception("API keys missing in .env")

        self.session = HTTP(
            api_key=API_KEY,
            api_secret=API_SECRET,
            testnet=TESTNET
        )

        self.max_retries = 3
        self.retry_delay = 1

    def _retry_request(self, func, *args, **kwargs):
        for attempt in range(self.max_retries):
            try:
                response = func(*args, **kwargs)
                if "retCode" in response and response["retCode"] == 0:
                    return response
                else:
                    logger.warning(f"API call failed: {response}")
                    if attempt < self.max_retries - 1:
                        time.sleep(self.retry_delay * (attempt + 1))
                        continue
                    raise Exception(f"API call failed after {self.max_retries} attempts: {response}")
            except Exception as e:
                logger.error(f"API request error: {e}")
                if attempt < self.max_retries - 1:
                    time.sleep(self.retry_delay * (attempt + 1))
                    continue
                raise e

    def get_balance(self):
        response = self._retry_request(self.session.get_wallet_balance, accountType="UNIFIED")
        return response["result"]["list"][0]["coin"]

    def get_positions(self, symbol=None):
        params = {"category": "linear"}
        if symbol:
            params["symbol"] = symbol

        response = self._retry_request(self.session.get_positions, **params)
        return response["result"]["list"]

    def set_leverage(self, symbol, leverage=None):
        lev = leverage or self.settings["trading"]["leverage"]
        response = self._retry_request(
            self.session.set_leverage,
            category="linear",
            symbol=symbol,
            buyLeverage=str(lev),
            sellLeverage=str(lev)
        )
        return response

    def market_order(self, symbol, side, qty, reduce_only=False):
        params = {
            "category": "linear",
            "symbol": symbol,
            "side": side.upper(),
            "orderType": "Market",
            "qty": str(qty)
        }

        if reduce_only:
            params["reduceOnly"] = True

        response = self._retry_request(self.session.place_order, **params)
        return response["result"]

    def limit_order(self, symbol, side, qty, price, reduce_only=False):
        params = {
            "category": "linear",
            "symbol": symbol,
            "side": side.upper(),
            "orderType": "Limit",
            "qty": str(qty),
            "price": str(price)
        }

        if reduce_only:
            params["reduceOnly"] = True

        response = self._retry_request(self.session.place_order, **params)
        return response["result"]

    def set_tp_sl(self, symbol, side, entry_price, tp_price, sl_price):
        response = self._retry_request(
            self.session.set_trading_stop,
            category="linear",
            symbol=symbol,
            takeProfit=str(tp_price),
            stopLoss=str(sl_price),
            tpTriggerBy="LastPrice",
            slTriggerBy="LastPrice"
        )
        return response

    def cancel_order(self, symbol, order_id=None):
        params = {"category": "linear", "symbol": symbol}
        if order_id:
            params["orderId"] = order_id

        response = self._retry_request(self.session.cancel_order, **params)
        return response["result"]

    def cancel_all_orders(self, symbol=None):
        params = {"category": "linear"}
        if symbol:
            params["symbol"] = symbol

        response = self._retry_request(self.session.cancel_all_orders, **params)
        return response["result"]

    def close_position(self, symbol):
        positions = self.get_positions(symbol)
        if not positions:
            return None

        position = positions[0]
        size = float(position["size"])

        if size > 0:

            return self.market_order(symbol, "Sell", size, reduce_only=True)
        elif size < 0:

            return self.market_order(symbol, "Buy", abs(size), reduce_only=True)

        return None

    def get_active_orders(self, symbol=None):
        params = {"category": "linear"}
        if symbol:
            params["symbol"] = symbol

        response = self._retry_request(self.session.get_open_orders, **params)
        return response["result"]["list"]

    def get_kline_data(self, symbol, interval="1", limit=200):
        response = self._retry_request(
            self.session.get_kline,
            category="linear",
            symbol=symbol,
            interval=interval,
            limit=limit
        )
        return response["result"]["list"]

    def get_orderbook(self, symbol, limit=25):
        response = self._retry_request(
            self.session.get_orderbook,
            category="linear",
            symbol=symbol,
            limit=limit
        )
        return response["result"]

    def get_ticker(self, symbol):
        response = self._retry_request(
            self.session.get_tickers,
            category="linear",
            symbol=symbol
        )
        return response["result"]["list"][0]

    def get_server_time(self):
        response = self._retry_request(self.session.get_server_time)
        return response["result"]["timeSecond"]