"""Kalshi order-book tick collector. Connects to the Kalshi market-data websocket, subscribes to the ``orderbook_delta`` channel for the most liquid currently-tradeable markets, reconstructs each book, and appends a microstructure feature row on every update to date-partitioned Parquet files under ``data/ticks/``. This is the data-collection stage of the microstructure study: level-2 tick history is not downloadable, so it has to be accumulated live. Run it on an always-on machine — every row written is a future training example. python -m src.microstructure.collector Paths are relative to the current directory; run it from the repo root. """ import asyncio import json import logging import time from datetime import datetime, timezone from pathlib import Path import aiohttp import pandas as pd import websockets from src.microstructure.orderbook import KalshiOrderBook logger = logging.getLogger(__name__) KALSHI_REST = "https://api.elections.kalshi.com/trade-api/v2" KALSHI_WS = "wss://api.elections.kalshi.com/trade-api/ws/v2" DATA_DIR = Path("data/ticks") N_MARKETS = 30 # number of top markets to track FLUSH_SECONDS = 60 # write a Parquet file at least this often... FLUSH_ROWS = 5000 # ...or once this many rows are buffered async def discover_liquid_markets(session, limit=N_MARKETS): """Return the most liquid currently-open Kalshi market tickers. The default ``/markets`` feed leads with same-day multi-leg combo markets (``KXMVE*``); ``min_close_ts`` skips those. Markets are then ranked by 24h volume so the collector spends its websocket budget on books that move. """ min_close_ts = int(time.time()) + 86400 ranked = [] cursor = None for _ in range(4): params = {"limit": 200, "min_close_ts": min_close_ts} if cursor: params["cursor"] = cursor async with session.get( f"{KALSHI_REST}/markets", params=params, timeout=aiohttp.ClientTimeout(total=15), ) as resp: data = await resp.json() for m in data.get("markets", []): if m["ticker"].startswith("KXMVE"): continue try: bid = float(m.get("yes_bid_dollars")) ask = float(m.get("yes_ask_dollars")) except (TypeError, ValueError): continue if bid <= 0 or ask <= 0: continue ranked.append((float(m.get("volume_24h_fp") or 0), m["ticker"])) cursor = data.get("cursor") if not cursor: break ranked.sort(reverse=True) return [ticker for _, ticker in ranked[:limit]] class TickCollector: def __init__(self, n_markets=N_MARKETS): self.n_markets = n_markets self.books: dict[str, KalshiOrderBook] = {} self.buffer: list[dict] = [] self.rows_written = 0 self._last_seq = None self._last_flush = time.monotonic() def _apply_message(self, msg) -> bool: """Fold one parsed websocket message into the books. Returns ``True`` when a sequence gap is detected and the caller should reconnect (Kalshi will replay fresh snapshots). Touches only in-memory state, so it is unit-testable without a live socket. """ mtype = msg.get("type") if mtype not in ("orderbook_snapshot", "orderbook_delta"): return False # 'subscribed' / 'ok' / 'error' control frames body = msg.get("msg", {}) ticker = body.get("market_ticker") or body.get("market_id") seq = msg.get("seq") # A non-consecutive seq means we missed an update — the local book is # now wrong. Bail out so the connection is rebuilt from a snapshot. if (mtype == "orderbook_delta" and self._last_seq is not None and seq is not None and seq != self._last_seq + 1): logger.warning("seq gap %s -> %s; resyncing", self._last_seq, seq) return True self._last_seq = seq book = self.books.setdefault(ticker, KalshiOrderBook(ticker)) if mtype == "orderbook_snapshot": book.apply_snapshot(body.get("yes", []), body.get("no", [])) else: book.apply_delta(body["price"], body["delta"], body["side"]) feats = book.features() if feats is not None: feats["ts"] = datetime.now(timezone.utc).isoformat() feats["seq"] = seq self.buffer.append(feats) return False def _maybe_flush(self): due = (len(self.buffer) >= FLUSH_ROWS or time.monotonic() - self._last_flush >= FLUSH_SECONDS) if due: self.flush() def flush(self): """Write the buffered rows to a fresh Parquet file and clear the buffer.""" self._last_flush = time.monotonic() if not self.buffer: return now = datetime.now(timezone.utc) out_dir = DATA_DIR / f"date={now:%Y-%m-%d}" out_dir.mkdir(parents=True, exist_ok=True) path = out_dir / f"kalshi_{now:%H%M%S_%f}.parquet" pd.DataFrame(self.buffer).to_parquet(path, index=False) self.rows_written += len(self.buffer) logger.info("flushed %d rows -> %s (total %d)", len(self.buffer), path, self.rows_written) self.buffer.clear() async def _run_once(self, tickers): """One websocket session: subscribe and consume until it drops.""" self._last_seq = None async with websockets.connect(KALSHI_WS, ping_interval=10) as ws: await ws.send(json.dumps({ "id": 1, "cmd": "subscribe", "params": { "channels": ["orderbook_delta"], "market_tickers": tickers, }, })) logger.info("subscribed to %d markets", len(tickers)) async for raw in ws: if self._apply_message(json.loads(raw)): return # seq gap -> let the outer loop reconnect self._maybe_flush() async def run(self): async with aiohttp.ClientSession() as session: tickers = await discover_liquid_markets(session, self.n_markets) if not tickers: logger.error("no liquid markets found; nothing to collect") return logger.info("tracking %d markets: %s", len(tickers), ", ".join(tickers)) backoff = 1 while True: try: await self._run_once(tickers) backoff = 1 except (websockets.ConnectionClosed, OSError) as e: logger.warning("connection lost: %s", e) backoff = min(backoff * 2, 30) except Exception: logger.exception("unexpected collector error") backoff = min(backoff * 2, 30) self.flush() # keep buffered rows across the gap self.books.clear() # stale — a reconnect replays snapshots await asyncio.sleep(backoff) def main(): from src import config # noqa: F401 (configures root logging) logging.getLogger("websockets").setLevel(logging.WARNING) collector = TickCollector() try: asyncio.run(collector.run()) except KeyboardInterrupt: collector.flush() logger.info("collector stopped; %d rows written total", collector.rows_written) if __name__ == "__main__": main()