Add data module

polymarket_bot/data.py

@@ -0,0 +1,349 @@
+"""
+Module de données: récupération des marchés, carnets d'ordres, et WebSocket.
+"""
+import asyncio
+import json
+import logging
+import time
+from dataclasses import dataclass, field
+from typing import Optional
+
+import aiohttp
+import requests
+
+from .config import GAMMA_API_URL, CLOB_API_URL, WS_URL
+
+logger = logging.getLogger("polybot.data")
+
+
+# ══════════════════════════════════════════════════════════════════
+# DATA MODELS
+# ══════════════════════════════════════════════════════════════════
+@dataclass
+class Token:
+    token_id: str
+    outcome: str        # "Yes" or "No"
+    price: float = 0.0
+
+
+@dataclass
+class Market:
+    market_id: str
+    condition_id: str
+    question: str
+    tokens: list        # [Token(YES), Token(NO)]
+    volume: float = 0.0
+    active: bool = True
+    closed: bool = False
+    end_date: str = ""
+    tags: list = field(default_factory=list)
+    neg_risk: bool = False
+    neg_risk_market_id: str = ""
+    last_updated: float = 0.0
+
+    @property
+    def yes_token(self) -> Optional[Token]:
+        for t in self.tokens:
+            if t.outcome.lower() == "yes":
+                return t
+        return None
+
+    @property
+    def no_token(self) -> Optional[Token]:
+        for t in self.tokens:
+            if t.outcome.lower() == "no":
+                return t
+        return None
+
+    @property
+    def spread(self) -> float:
+        """Différence entre YES+NO et 1.0. Positif = arbitrage long possible."""
+        yes = self.yes_token
+        no = self.no_token
+        if yes and no and yes.price > 0 and no.price > 0:
+            return 1.0 - (yes.price + no.price)
+        return 0.0
+
+
+@dataclass
+class OrderBookLevel:
+    price: float
+    size: float
+
+
+@dataclass
+class OrderBook:
+    token_id: str
+    bids: list  # [OrderBookLevel]
+    asks: list  # [OrderBookLevel]
+    timestamp: float = 0.0
+
+    @property
+    def best_bid(self) -> Optional[float]:
+        return self.bids[0].price if self.bids else None
+
+    @property
+    def best_ask(self) -> Optional[float]:
+        return self.asks[0].price if self.asks else None
+
+    @property
+    def mid_price(self) -> Optional[float]:
+        if self.best_bid is not None and self.best_ask is not None:
+            return (self.best_bid + self.best_ask) / 2
+        return None
+
+    @property
+    def total_bid_depth(self) -> float:
+        return sum(l.price * l.size for l in self.bids)
+
+    @property
+    def total_ask_depth(self) -> float:
+        return sum(l.price * l.size for l in self.asks)
+
+
+# ══════════════════════════════════════════════════════════════════
+# GAMMA API (Public Market Data)
+# ══════════════════════════════════════════════════════════════════
+class GammaClient:
+    """Client pour l'API Gamma (données de marché publiques)."""
+
+    def __init__(self, base_url: str = GAMMA_API_URL):
+        self.base_url = base_url
+        self._session: Optional[aiohttp.ClientSession] = None
+
+    async def _get_session(self) -> aiohttp.ClientSession:
+        if self._session is None or self._session.closed:
+            self._session = aiohttp.ClientSession()
+        return self._session
+
+    async def close(self):
+        if self._session and not self._session.closed:
+            await self._session.close()
+
+    async def get_active_markets(self, limit: int = 100, offset: int = 0) -> list[dict]:
+        """Récupère les marchés actifs depuis Gamma API."""
+        session = await self._get_session()
+        params = {
+            "active": "true",
+            "closed": "false",
+            "limit": limit,
+            "offset": offset,
+        }
+        try:
+            async with session.get(f"{self.base_url}/markets", params=params) as resp:
+                if resp.status == 200:
+                    data = await resp.json()
+                    return data if isinstance(data, list) else data.get("data", [])
+                else:
+                    logger.error(f"Gamma API error: {resp.status}")
+                    return []
+        except Exception as e:
+            logger.error(f"Gamma API request failed: {e}")
+            return []
+
+    async def get_all_active_markets(self, max_markets: int = 500) -> list[Market]:
+        """Récupère tous les marchés actifs avec pagination."""
+        all_markets = []
+        offset = 0
+        batch_size = 100
+
+        while offset < max_markets:
+            raw = await self.get_active_markets(limit=batch_size, offset=offset)
+            if not raw:
+                break
+
+            for m in raw:
+                try:
+                    tokens = []
+                    # Parse outcome prices (can be JSON string or list)
+                    outcome_prices_raw = m.get("outcomePrices", "[]")
+                    if isinstance(outcome_prices_raw, str):
+                        outcome_prices = json.loads(outcome_prices_raw)
+                    else:
+                        outcome_prices = outcome_prices_raw
+
+                    # Parse outcomes (can be JSON string or list)
+                    outcomes_raw = m.get("outcomes", "[]")
+                    if isinstance(outcomes_raw, str):
+                        outcomes = json.loads(outcomes_raw)
+                    else:
+                        outcomes = outcomes_raw
+
+                    # Parse token IDs from clobTokenIds
+                    clob_token_ids_raw = m.get("clobTokenIds", "[]")
+                    if isinstance(clob_token_ids_raw, str):
+                        clob_token_ids = json.loads(clob_token_ids_raw)
+                    else:
+                        clob_token_ids = clob_token_ids_raw
+
+                    # Build tokens from outcomes + clobTokenIds + prices
+                    if outcomes and clob_token_ids and outcome_prices:
+                        for i in range(min(len(outcomes), len(clob_token_ids), len(outcome_prices))):
+                            tokens.append(Token(
+                                token_id=str(clob_token_ids[i]),
+                                outcome=outcomes[i],
+                                price=float(outcome_prices[i]),
+                            ))
+
+                    market = Market(
+                        market_id=m.get("id", ""),
+                        condition_id=m.get("conditionId", ""),
+                        question=m.get("question", ""),
+                        tokens=tokens,
+                        volume=float(m.get("volumeNum", m.get("volume", 0))),
+                        active=m.get("active", True),
+                        closed=m.get("closed", False),
+                        end_date=m.get("endDate", ""),
+                        tags=[t.get("label", t) if isinstance(t, dict) else str(t) for t in (m.get("events", [{}])[0].get("tags", []) if m.get("events") else [])],
+                        neg_risk=m.get("negRisk", False),
+                        neg_risk_market_id=m.get("negRiskRequestID", ""),
+                        last_updated=time.time(),
+                    )
+                    all_markets.append(market)
+                except Exception as e:
+                    logger.warning(f"Failed to parse market: {e}")
+                    continue
+
+            offset += batch_size
+            if len(raw) < batch_size:
+                break
+
+        logger.info(f"Fetched {len(all_markets)} active markets from Gamma API")
+        return all_markets
+
+    async def get_events(self) -> list[dict]:
+        """Récupère les événements (groupes de marchés)."""
+        session = await self._get_session()
+        try:
+            async with session.get(f"{self.base_url}/events", params={"active": "true"}) as resp:
+                if resp.status == 200:
+                    return await resp.json()
+                return []
+        except Exception as e:
+            logger.error(f"Failed to get events: {e}")
+            return []
+
+
+# ══════════════════════════════════════════════════════════════════
+# CLOB API (Order Book Data)
+# ══════════════════════════════════════════════════════════════════
+class CLOBDataClient:
+    """Client pour les données du CLOB (carnet d'ordres)."""
+
+    def __init__(self, base_url: str = CLOB_API_URL):
+        self.base_url = base_url
+        self._session: Optional[aiohttp.ClientSession] = None
+
+    async def _get_session(self) -> aiohttp.ClientSession:
+        if self._session is None or self._session.closed:
+            self._session = aiohttp.ClientSession()
+        return self._session
+
+    async def close(self):
+        if self._session and not self._session.closed:
+            await self._session.close()
+
+    async def get_order_book(self, token_id: str) -> Optional[OrderBook]:
+        """Récupère le carnet d'ordres pour un token."""
+        session = await self._get_session()
+        try:
+            async with session.get(f"{self.base_url}/book", params={"token_id": token_id}) as resp:
+                if resp.status == 200:
+                    data = await resp.json()
+                    bids = [OrderBookLevel(float(b["price"]), float(b["size"]))
+                            for b in data.get("bids", [])]
+                    asks = [OrderBookLevel(float(a["price"]), float(a["size"]))
+                            for a in data.get("asks", [])]
+                    # Trier: bids décroissant, asks croissant
+                    bids.sort(key=lambda x: x.price, reverse=True)
+                    asks.sort(key=lambda x: x.price)
+                    return OrderBook(token_id=token_id, bids=bids, asks=asks, timestamp=time.time())
+                else:
+                    logger.warning(f"Order book fetch failed for {token_id}: {resp.status}")
+                    return None
+        except Exception as e:
+            logger.error(f"Order book request failed: {e}")
+            return None
+
+    async def get_midpoint(self, token_id: str) -> Optional[float]:
+        """Récupère le prix médian pour un token."""
+        session = await self._get_session()
+        try:
+            async with session.get(f"{self.base_url}/midpoint", params={"token_id": token_id}) as resp:
+                if resp.status == 200:
+                    data = await resp.json()
+                    return float(data.get("mid", 0))
+                return None
+        except Exception as e:
+            logger.error(f"Midpoint request failed: {e}")
+            return None
+
+    async def get_price(self, token_id: str, side: str = "BUY") -> Optional[float]:
+        """Récupère le prix pour un côté (BUY/SELL)."""
+        session = await self._get_session()
+        try:
+            params = {"token_id": token_id, "side": side}
+            async with session.get(f"{self.base_url}/price", params=params) as resp:
+                if resp.status == 200:
+                    data = await resp.json()
+                    return float(data.get("price", 0))
+                return None
+        except Exception as e:
+            logger.error(f"Price request failed: {e}")
+            return None
+
+
+# ══════════════════════════════════════════════════════════════════
+# WEBSOCKET REAL-TIME FEED
+# ══════════════════════════════════════════════════════════════════
+class WebSocketFeed:
+    """Flux WebSocket temps réel pour les carnets d'ordres."""
+
+    def __init__(self, url: str = WS_URL):
+        self.url = url
+        self._callbacks: dict = {}  # token_id -> [callback]
+        self._running = False
+        self._ws = None
+
+    def subscribe(self, token_id: str, callback):
+        """Enregistre un callback pour les mises à jour d'un token."""
+        if token_id not in self._callbacks:
+            self._callbacks[token_id] = []
+        self._callbacks[token_id].append(callback)
+
+    async def start(self, token_ids: list[str]):
+        """Démarre le flux WebSocket."""
+        import websockets
+
+        self._running = True
+        while self._running:
+            try:
+                async with websockets.connect(self.url) as ws:
+                    self._ws = ws
+                    # Subscribe to all markets
+                    subscribe_msg = {
+                        "auth": {},
+                        "markets": [{"asset_id": tid, "type": "book"} for tid in token_ids]
+                    }
+                    await ws.send(json.dumps(subscribe_msg))
+                    logger.info(f"WebSocket connected, subscribed to {len(token_ids)} tokens")
+
+                    async for msg in ws:
+                        try:
+                            data = json.loads(msg)
+                            asset_id = data.get("asset_id", "")
+                            if asset_id in self._callbacks:
+                                for cb in self._callbacks[asset_id]:
+                                    await cb(data)
+                        except json.JSONDecodeError:
+                            continue
+
+            except Exception as e:
+                logger.error(f"WebSocket error: {e}, reconnecting in 5s...")
+                await asyncio.sleep(5)
+
+    async def stop(self):
+        """Arrête le flux WebSocket."""
+        self._running = False
+        if self._ws:
+            await self._ws.close()