Cerbero-mcp/services/mcp-alpaca/src/mcp_alpaca/client.py

from __future__ import annotations

import asyncio
import datetime as _dt
from typing import Any

from alpaca.data.historical import (
    CryptoHistoricalDataClient,
    OptionHistoricalDataClient,
    StockHistoricalDataClient,
)
from alpaca.data.requests import (
    CryptoBarsRequest,
    CryptoLatestQuoteRequest,
    CryptoLatestTradeRequest,
    OptionBarsRequest,
    OptionChainRequest,
    OptionLatestQuoteRequest,
    StockBarsRequest,
    StockLatestQuoteRequest,
    StockLatestTradeRequest,
    StockSnapshotRequest,
)
from alpaca.data.timeframe import TimeFrame, TimeFrameUnit
from alpaca.trading.client import TradingClient
from alpaca.trading.enums import (
    AssetClass,
    OrderSide,
    QueryOrderStatus,
    TimeInForce,
)
from alpaca.trading.requests import (
    ClosePositionRequest,
    GetAssetsRequest,
    GetOrdersRequest,
    LimitOrderRequest,
    MarketOrderRequest,
    ReplaceOrderRequest,
    StopOrderRequest,
)

_TF_MAP = {
    "1min": TimeFrame(1, TimeFrameUnit.Minute),
    "5min": TimeFrame(5, TimeFrameUnit.Minute),
    "15min": TimeFrame(15, TimeFrameUnit.Minute),
    "30min": TimeFrame(30, TimeFrameUnit.Minute),
    "1h": TimeFrame(1, TimeFrameUnit.Hour),
    "1d": TimeFrame(1, TimeFrameUnit.Day),
    "1w": TimeFrame(1, TimeFrameUnit.Week),
}

_ASSET_CLASSES = {"stocks", "crypto", "options"}


def _tf(interval: str) -> TimeFrame:
    if interval in _TF_MAP:
        return _TF_MAP[interval]
    raise ValueError(f"unsupported timeframe: {interval}")


def _asset_class_enum(ac: str) -> AssetClass:
    ac = ac.lower()
    if ac == "stocks":
        return AssetClass.US_EQUITY
    if ac == "crypto":
        return AssetClass.CRYPTO
    if ac == "options":
        return AssetClass.US_OPTION
    raise ValueError(f"invalid asset_class: {ac}")


def _serialize(obj: Any) -> Any:
    """Recursively convert pydantic/datetime objects → json-safe."""
    if obj is None or isinstance(obj, str | int | float | bool):
        return obj
    if isinstance(obj, _dt.datetime | _dt.date):
        return obj.isoformat()
    if isinstance(obj, dict):
        return {k: _serialize(v) for k, v in obj.items()}
    if isinstance(obj, list | tuple):
        return [_serialize(v) for v in obj]
    if hasattr(obj, "model_dump"):
        return _serialize(obj.model_dump())
    if hasattr(obj, "__dict__"):
        return _serialize(vars(obj))
    return str(obj)


class AlpacaClient:
    def __init__(
        self,
        api_key: str,
        secret_key: str,
        paper: bool = True,
        trading: Any | None = None,
        stock_data: Any | None = None,
        crypto_data: Any | None = None,
        option_data: Any | None = None,
    ) -> None:
        self.api_key = api_key
        self.secret_key = secret_key
        self.paper = paper
        self._trading = trading or TradingClient(
            api_key=api_key, secret_key=secret_key, paper=paper
        )
        self._stock = stock_data or StockHistoricalDataClient(
            api_key=api_key, secret_key=secret_key
        )
        self._crypto = crypto_data or CryptoHistoricalDataClient(
            api_key=api_key, secret_key=secret_key
        )
        self._option = option_data or OptionHistoricalDataClient(
            api_key=api_key, secret_key=secret_key
        )

    async def _run(self, fn, /, *args, **kwargs):
        return await asyncio.to_thread(fn, *args, **kwargs)

    # ── Account / positions ──────────────────────────────────────

    async def get_account(self) -> dict:
        acc = await self._run(self._trading.get_account)
        return _serialize(acc)

    async def get_positions(self) -> list[dict]:
        pos = await self._run(self._trading.get_all_positions)
        return [_serialize(p) for p in pos]

    async def get_activities(self, limit: int = 50) -> list[dict]:
        acts = await self._run(self._trading.get_account_activities)
        data = [_serialize(a) for a in acts]
        return data[:limit]

    # ── Assets ──────────────────────────────────────────────────

    async def get_assets(
        self, asset_class: str = "stocks", status: str = "active"
    ) -> list[dict]:
        req = GetAssetsRequest(
            asset_class=_asset_class_enum(asset_class),
            status=status,
        )
        assets = await self._run(self._trading.get_all_assets, req)
        return [_serialize(a) for a in assets[:500]]

    # ── Market data ─────────────────────────────────────────────

    async def get_ticker(self, symbol: str, asset_class: str = "stocks") -> dict:
        ac = asset_class.lower()
        if ac == "stocks":
            req = StockLatestTradeRequest(symbol_or_symbols=symbol)
            data = await self._run(self._stock.get_stock_latest_trade, req)
            trade = data.get(symbol)
            q_req = StockLatestQuoteRequest(symbol_or_symbols=symbol)
            qdata = await self._run(self._stock.get_stock_latest_quote, q_req)
            quote = qdata.get(symbol)
            return {
                "symbol": symbol,
                "asset_class": "stocks",
                "last_price": getattr(trade, "price", None),
                "bid": getattr(quote, "bid_price", None),
                "ask": getattr(quote, "ask_price", None),
                "bid_size": getattr(quote, "bid_size", None),
                "ask_size": getattr(quote, "ask_size", None),
                "timestamp": _serialize(getattr(trade, "timestamp", None)),
            }
        if ac == "crypto":
            req = CryptoLatestTradeRequest(symbol_or_symbols=symbol)
            data = await self._run(self._crypto.get_crypto_latest_trade, req)
            trade = data.get(symbol)
            q_req = CryptoLatestQuoteRequest(symbol_or_symbols=symbol)
            qdata = await self._run(self._crypto.get_crypto_latest_quote, q_req)
            quote = qdata.get(symbol)
            return {
                "symbol": symbol,
                "asset_class": "crypto",
                "last_price": getattr(trade, "price", None),
                "bid": getattr(quote, "bid_price", None),
                "ask": getattr(quote, "ask_price", None),
                "timestamp": _serialize(getattr(trade, "timestamp", None)),
            }
        if ac == "options":
            req = OptionLatestQuoteRequest(symbol_or_symbols=symbol)
            data = await self._run(self._option.get_option_latest_quote, req)
            quote = data.get(symbol)
            return {
                "symbol": symbol,
                "asset_class": "options",
                "bid": getattr(quote, "bid_price", None),
                "ask": getattr(quote, "ask_price", None),
                "timestamp": _serialize(getattr(quote, "timestamp", None)),
            }
        raise ValueError(f"invalid asset_class: {asset_class}")

    async def get_bars(
        self,
        symbol: str,
        asset_class: str = "stocks",
        interval: str = "1d",
        start: str | None = None,
        end: str | None = None,
        limit: int = 1000,
    ) -> dict:
        tf = _tf(interval)
        start_dt = _dt.datetime.fromisoformat(start) if start else (
            _dt.datetime.now(_dt.UTC) - _dt.timedelta(days=30)
        )
        end_dt = _dt.datetime.fromisoformat(end) if end else _dt.datetime.now(_dt.UTC)
        ac = asset_class.lower()
        if ac == "stocks":
            req = StockBarsRequest(
                symbol_or_symbols=symbol, timeframe=tf,
                start=start_dt, end=end_dt, limit=limit,
            )
            data = await self._run(self._stock.get_stock_bars, req)
        elif ac == "crypto":
            req = CryptoBarsRequest(
                symbol_or_symbols=symbol, timeframe=tf,
                start=start_dt, end=end_dt, limit=limit,
            )
            data = await self._run(self._crypto.get_crypto_bars, req)
        elif ac == "options":
            req = OptionBarsRequest(
                symbol_or_symbols=symbol, timeframe=tf,
                start=start_dt, end=end_dt, limit=limit,
            )
            data = await self._run(self._option.get_option_bars, req)
        else:
            raise ValueError(f"invalid asset_class: {asset_class}")
        bars_dict = getattr(data, "data", {}) or {}
        rows = bars_dict.get(symbol, []) or []
        bars = [
            {
                "timestamp": _serialize(getattr(b, "timestamp", None)),
                "open": getattr(b, "open", None),
                "high": getattr(b, "high", None),
                "low": getattr(b, "low", None),
                "close": getattr(b, "close", None),
                "volume": getattr(b, "volume", None),
            }
            for b in rows
        ]
        return {"symbol": symbol, "asset_class": ac, "interval": interval, "bars": bars}

    async def get_snapshot(self, symbol: str) -> dict:
        req = StockSnapshotRequest(symbol_or_symbols=symbol)
        data = await self._run(self._stock.get_stock_snapshot, req)
        return _serialize(data.get(symbol))

    async def get_option_chain(
        self,
        underlying: str,
        expiry: str | None = None,
    ) -> dict:
        kwargs: dict[str, Any] = {"underlying_symbol": underlying}
        if expiry:
            kwargs["expiration_date"] = _dt.date.fromisoformat(expiry)
        req = OptionChainRequest(**kwargs)
        data = await self._run(self._option.get_option_chain, req)
        return {
            "underlying": underlying,
            "expiry": expiry,
            "contracts": _serialize(data),
        }

    # ── Orders ──────────────────────────────────────────────────

    async def get_open_orders(self, limit: int = 50) -> list[dict]:
        req = GetOrdersRequest(status=QueryOrderStatus.OPEN, limit=limit)
        orders = await self._run(self._trading.get_orders, filter=req)
        return [_serialize(o) for o in orders]

    async def place_order(
        self,
        symbol: str,
        side: str,
        qty: float | None = None,
        notional: float | None = None,
        order_type: str = "market",
        limit_price: float | None = None,
        stop_price: float | None = None,
        tif: str = "day",
        asset_class: str = "stocks",
    ) -> dict:
        side_enum = OrderSide.BUY if side.lower() == "buy" else OrderSide.SELL
        tif_enum = TimeInForce(tif.lower())
        ot = order_type.lower()
        common = {
            "symbol": symbol,
            "side": side_enum,
            "time_in_force": tif_enum,
        }
        if qty is not None:
            common["qty"] = qty
        if notional is not None:
            common["notional"] = notional
        if ot == "market":
            req = MarketOrderRequest(**common)
        elif ot == "limit":
            if limit_price is None:
                raise ValueError("limit_price required for limit order")
            req = LimitOrderRequest(**common, limit_price=limit_price)
        elif ot == "stop":
            if stop_price is None:
                raise ValueError("stop_price required for stop order")
            req = StopOrderRequest(**common, stop_price=stop_price)
        else:
            raise ValueError(f"unsupported order_type: {order_type}")
        order = await self._run(self._trading.submit_order, req)
        return _serialize(order)

    async def amend_order(
        self,
        order_id: str,
        qty: float | None = None,
        limit_price: float | None = None,
        stop_price: float | None = None,
        tif: str | None = None,
    ) -> dict:
        kwargs: dict[str, Any] = {}
        if qty is not None:
            kwargs["qty"] = qty
        if limit_price is not None:
            kwargs["limit_price"] = limit_price
        if stop_price is not None:
            kwargs["stop_price"] = stop_price
        if tif is not None:
            kwargs["time_in_force"] = TimeInForce(tif.lower())
        req = ReplaceOrderRequest(**kwargs)
        order = await self._run(self._trading.replace_order_by_id, order_id, req)
        return _serialize(order)

    async def cancel_order(self, order_id: str) -> dict:
        await self._run(self._trading.cancel_order_by_id, order_id)
        return {"order_id": order_id, "canceled": True}

    async def cancel_all_orders(self) -> list[dict]:
        resp = await self._run(self._trading.cancel_orders)
        return [_serialize(r) for r in resp]

    # ── Position close ──────────────────────────────────────────

    async def close_position(
        self, symbol: str, qty: float | None = None, percentage: float | None = None
    ) -> dict:
        req = None
        if qty is not None or percentage is not None:
            kwargs: dict[str, Any] = {}
            if qty is not None:
                kwargs["qty"] = str(qty)
            if percentage is not None:
                kwargs["percentage"] = str(percentage)
            req = ClosePositionRequest(**kwargs)
        order = await self._run(
            self._trading.close_position, symbol, close_options=req
        )
        return _serialize(order)

    async def close_all_positions(self, cancel_orders: bool = True) -> list[dict]:
        resp = await self._run(
            self._trading.close_all_positions, cancel_orders=cancel_orders
        )
        return [_serialize(r) for r in resp]

    # ── Clock / calendar ────────────────────────────────────────

    async def get_clock(self) -> dict:
        clock = await self._run(self._trading.get_clock)
        return _serialize(clock)

    async def get_calendar(
        self, start: str | None = None, end: str | None = None
    ) -> list[dict]:
        from alpaca.trading.requests import GetCalendarRequest

        kwargs: dict[str, Any] = {}
        if start:
            kwargs["start"] = _dt.date.fromisoformat(start)
        if end:
            kwargs["end"] = _dt.date.fromisoformat(end)
        req = GetCalendarRequest(**kwargs) if kwargs else None
        cal = await self._run(
            self._trading.get_calendar, filters=req
        ) if req else await self._run(self._trading.get_calendar)
        return [_serialize(c) for c in cal]