Upload websocket_hub.py

websocket_hub.py

@@ -45,7 +45,7 @@ from pathlib import Path
 from typing import Dict, List, Optional, Set
 
 import uvicorn
-from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from fastapi import FastAPI, Request, WebSocket, WebSocketDisconnect
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import FileResponse, JSONResponse
 
@@ -180,14 +180,13 @@ class ConnectionManager:
     })
 
     def __init__(self):
-        self._publishers:        Dict[str, WebSocket]       = {}
-        self._subscribers:       Set[WebSocket]              = set()
-        self._top3_subscribers:  Set[WebSocket]              = set()   # PATCH A
-        self._snapshots:         Dict[str, dict]             = {}
-        self._history:           Dict[str, deque]            = {}   # rolling per-space history
-        self._lock               = asyncio.Lock()
-        self._total_ingested:    int                         = 0
-        self._msg_counts:        Dict[str, Dict[str, int]]  = {}   # {space: {msg_type: count}}
+        self._publishers:     Dict[str, WebSocket]       = {}
+        self._subscribers:    Set[WebSocket]              = set()
+        self._snapshots:      Dict[str, dict]             = {}
+        self._history:        Dict[str, deque]            = {}   # rolling per-space history
+        self._lock            = asyncio.Lock()
+        self._total_ingested: int                         = 0
+        self._msg_counts:     Dict[str, Dict[str, int]]  = {}   # {space: {msg_type: count}}
 
     async def register_publisher(self, space_name: str, ws: WebSocket) -> None:
         await ws.accept()
@@ -270,7 +269,6 @@ class ConnectionManager:
             async with self._lock:
                 for ws in dead:
                     self._subscribers.discard(ws)
-        await self.broadcast_top3()          # PATCH C — push top-3 on every ingest
 
     async def send_initial_state(self, ws: WebSocket) -> None:
         async with self._lock:
@@ -313,88 +311,6 @@ class ConnectionManager:
             "subscriber_count": len(self._subscribers),
         }
 
-    # ── PATCH B — Top-3 subscriber management ─────────────────────────────────────
-
-    async def register_top3_subscriber(self, ws: WebSocket) -> None:
-        await ws.accept()
-        async with self._lock:
-            self._top3_subscribers.add(ws)
-        # Push current rankings immediately on connect
-        await self.broadcast_top3(target=ws)
-        logger.info(f"🏆 Top-3 subscriber connected (total={len(self._top3_subscribers)})")
-
-    async def unregister_top3_subscriber(self, ws: WebSocket) -> None:
-        async with self._lock:
-            self._top3_subscribers.discard(ws)
-        logger.info("🏆 Top-3 subscriber disconnected")
-
-    async def broadcast_top3(self, target: Optional[WebSocket] = None) -> None:
-        """
-        Compute top-3 ranked assets and push them to:
-          • a single WebSocket (target=ws) — used on initial connect
-          • all _top3_subscribers     (target=None) — used on every ingest
-        """
-        # ── Build ranked list (mirrors _compute_rankings() logic) ─────────────────
-        ranked = []
-        async with self._lock:
-            snapshots = dict(self._snapshots)
-
-        for name, snap in snapshots.items():
-            training = snap.get("training", {})
-            voting   = snap.get("voting",   {})
-            buy      = voting.get("buy_count",  0)
-            sell     = voting.get("sell_count", 0)
-            total    = buy + sell
-            sig_conf = (max(buy, sell) / total) if total > 0 else 0.0
-            avn_acc  = training.get("avn_accuracy", 0.0)
-            score    = round(sig_conf - avn_acc, 6)
-            ranked.append({
-                "rank":              0,             # filled in below
-                "space_name":        name,
-                "score":             score,
-                "signal_confidence": round(sig_conf, 6),
-                "avn_accuracy":      round(avn_acc,  6),
-                "dominant_signal":   voting.get("dominant_signal", "NEUTRAL"),
-                "buy_count":         buy,
-                "sell_count":        sell,
-                "last_price":        voting.get("last_price", 0.0),
-                "training_steps":    training.get("training_steps", 0),
-                "actor_loss":        training.get("actor_loss",     0.0),
-                "critic_loss":       training.get("critic_loss",    0.0),
-                "avn_loss":          training.get("avn_loss",       0.0),
-                "last_updated":      snap.get("last_updated",       0.0),
-            })
-
-        ranked.sort(key=lambda r: r["score"], reverse=True)
-        for i, r in enumerate(ranked):
-            r["rank"] = i + 1
-
-        message = json.dumps({
-            "type":          "top3_rankings",
-            "rankings":      ranked[:3],
-            "total_assets":  len(ranked),
-            "hub_timestamp": time.time(),
-        })
-
-        # ── Deliver ───────────────────────────────────────────────────────────────
-        if target is not None:
-            try:
-                await target.send_text(message)
-            except Exception:
-                pass
-            return
-
-        dead = []
-        for ws in list(self._top3_subscribers):
-            try:
-                await ws.send_text(message)
-            except Exception:
-                dead.append(ws)
-        if dead:
-            async with self._lock:
-                for ws in dead:
-                    self._top3_subscribers.discard(ws)
-
 
 # ══════════════════════════════════════════════════════════════════════════════════════
 # SECTION 3 — HUB TRADE STORE  (in-memory, fed by WebSocket messages)
@@ -539,6 +455,13 @@ _LOG_DIR    = os.environ.get("RANKER_LOG_DIR", "/app/ranker_logs")
 _hub_trades = HubTradeStore()
 logger.info("✅ HubTradeStore initialised — awaiting trade_opened/trade_closed WS messages")
 
+# ── AXRVI live rankings store ─────────────────────────────────────────────────────────
+# Populated by POST /api/axrvi/rankings from the Executo ranker after every
+# rank_and_gate() cycle (~every 5s). Falls back to hub-snapshot scoring when stale.
+_axrvi_rankings:    List[dict] = []
+_axrvi_rankings_ts: float      = 0.0
+_AXRVI_RANKINGS_TTL: float     = 30.0   # seconds before falling back to snapshot scoring
+
 
 # ══════════════════════════════════════════════════════════════════════════════════════
 # SECTION 4 — FASTAPI APPLICATION
@@ -1079,6 +1002,51 @@ _HTML_PATH = Path(os.environ.get(
 
 
 def _compute_rankings() -> List[dict]:
+    """
+    Build the rankings list served by /api/state.
+
+    Priority order:
+      1. Live AXRVI rankings pushed by the Executo ranker via
+         POST /api/axrvi/rankings (within the last 30 s).
+         These contain real softmax-Shreve priorities from AXRVINet.
+      2. Fallback: hub-snapshot vote-ratio scoring used before the
+         ranker connects or if the push is stale (e.g. ranker restart).
+    """
+    global _axrvi_rankings, _axrvi_rankings_ts
+
+    # ── Path 1: fresh AXRVI rankings ────────────────────────────────────────
+    if _axrvi_rankings and (time.time() - _axrvi_rankings_ts) < _AXRVI_RANKINGS_TTL:
+        snapshots = manager.get_all_snapshots()
+        merged: List[dict] = []
+        for r in _axrvi_rankings:
+            name     = r.get("space_name", "")
+            snap     = snapshots.get(name, {})
+            training = snap.get("training", {})
+            voting   = snap.get("voting",   {})
+            buy      = voting.get("buy_count",  r.get("buy_count",  0))
+            sell     = voting.get("sell_count", r.get("sell_count", 0))
+            merged.append({
+                # Core AXRVI fields — these are the live ranker values
+                "rank":              r.get("rank",             0),
+                "space_name":        name,
+                "score":             r.get("score",            0.0),
+                "final_priority":    r.get("final_priority",   r.get("score", 0.0)),
+                "signal_confidence": r.get("signal_confidence",0.0),
+                "dominant_signal":   r.get("dominant_signal",  "NEUTRAL"),
+                "avn_accuracy":      r.get("avn_accuracy",     0.0),
+                "epistemic_std":     r.get("epistemic_std",    0.0),
+                "training_steps":    r.get("training_steps",   training.get("training_steps", 0)),
+                # Hub-snapshot fields merged in (latest available)
+                "actor_loss":        training.get("actor_loss",  0.0),
+                "critic_loss":       training.get("critic_loss", 0.0),
+                "avn_loss":          training.get("avn_loss",    0.0),
+                "buy_count":         buy,
+                "sell_count":        sell,
+                "last_updated":      snap.get("last_updated",   _axrvi_rankings_ts),
+            })
+        return merged
+
+    # ── Path 2: fallback hub-snapshot scoring ────────────────────────────────
     ranked: List[dict] = []
     for name, snap in manager.get_all_snapshots().items():
         training = snap.get("training", {})
@@ -1093,6 +1061,7 @@ def _compute_rankings() -> List[dict]:
             "rank":              0,
             "space_name":        name,
             "score":             score,
+            "final_priority":    score,
             "signal_confidence": round(sig_conf, 6),
             "avn_accuracy":      round(avn_acc,  6),
             "dominant_signal":   voting.get("dominant_signal", "NEUTRAL"),
@@ -1103,6 +1072,7 @@ def _compute_rankings() -> List[dict]:
             "critic_loss":       training.get("critic_loss",    0.0),
             "avn_loss":          training.get("avn_loss",       0.0),
             "last_updated":      snap.get("last_updated",       0.0),
+            "epistemic_std":     0.0,
         })
     ranked.sort(key=lambda r: r["score"], reverse=True)
     for i, r in enumerate(ranked):
@@ -1131,6 +1101,36 @@ async def serve_dashboard():
     )
 
 
+@app.post("/api/axrvi/rankings")
+async def receive_axrvi_rankings(request: Request):
+    """
+    Called by the Executo ranker after every rank_and_gate() cycle (~5 s).
+    Stores the live AXRVI-scored ranking list so _compute_rankings() can serve
+    it from /api/state instead of the stale hub-snapshot vote-ratio fallback.
+
+    Expected body:
+      {"rankings": [{"space_name": "V75", "score": 0.24, "rank": 1, ...}, ...]}
+    """
+    global _axrvi_rankings, _axrvi_rankings_ts
+    try:
+        body = await request.json()
+    except Exception as e:
+        return JSONResponse({"ok": False, "error": f"Bad JSON: {e}"}, status_code=400)
+
+    rankings = body.get("rankings", [])
+    if not isinstance(rankings, list):
+        return JSONResponse({"ok": False, "error": "rankings must be a list"}, status_code=400)
+
+    _axrvi_rankings    = rankings
+    _axrvi_rankings_ts = time.time()
+    logger.debug(
+        f"[AXRVI Rankings] Received {len(rankings)} assets | "
+        f"top={rankings[0].get('space_name','?')} score={rankings[0].get('score',0):.4f}"
+        if rankings else "[AXRVI Rankings] Received empty list"
+    )
+    return JSONResponse({"ok": True, "count": len(rankings), "ts": _axrvi_rankings_ts})
+
+
 @app.get("/api/state")
 async def api_state():
     """Full dashboard state polled by hub_dashboard.html every 2 s."""
@@ -1162,58 +1162,6 @@ async def api_state():
 _START_TIME = time.time()
 
 
-# ══════════════════════════════════════════════════════════════════════════════════════
-# SECTION 8b — TOP-3 RANKINGS ENDPOINTS  (PATCH D)
-# ══════════════════════════════════════════════════════════════════════════════════════
-
-@app.get("/api/rankings/top3")
-async def api_top3_rankings():
-    """
-    GET /api/rankings/top3
-
-    Returns the top 3 assets scored by:
-      score = signal_confidence − avn_accuracy
-      (higher = stronger signal conviction with lower training noise)
-
-    Suitable for one-shot polling from your IDE or any HTTP client.
-    For real-time streaming connect to ws://<host>/ws/top3 instead.
-    """
-    rankings = _compute_rankings()[:3]
-    return JSONResponse({
-        "top3":         rankings,
-        "total_assets": len(_compute_rankings()),
-        "timestamp":    datetime.utcnow().isoformat() + "Z",
-    })
-
-
-@app.websocket("/ws/top3")
-async def ws_top3_endpoint(websocket: WebSocket):
-    """
-    WS /ws/top3
-
-    Pushes a 'top3_rankings' message to the client every time any asset
-    publishes a metrics update.  The payload is identical to /api/rankings/top3
-    but delivered in real time.
-
-    Message shape:
-      {
-        "type":          "top3_rankings",
-        "rankings":      [ { rank, space_name, score, dominant_signal, … }, … ],
-        "total_assets":  <int>,
-        "hub_timestamp": <unix_float>
-      }
-    """
-    await manager.register_top3_subscriber(websocket)
-    try:
-        while True:
-            # Keep connection alive; we only send, never read from clients here.
-            await websocket.receive_text()
-    except Exception:
-        pass
-    finally:
-        await manager.unregister_top3_subscriber(websocket)
-
-
 # ══════════════════════════════════════════════════════════════════════════════════════
 # SECTION 9 — ENTRY POINT
 # ══════════════════════════════════════════════════════════════════════════════════════