add polygon indices and send message to broadcast server from main client

openbb_platform/extensions/websockets/openbb_websockets/broadcast.py

@@ -25,6 +25,31 @@ AUTH_TOKEN = kwargs.pop("auth_token", None)
 app = FastAPI()
 
 
+async def read_stdin(broadcast_server):
+    """Read from stdin."""
+    while True:
+        line = await asyncio.get_event_loop().run_in_executor(None, sys.stdin.readline)
+        sys.stdin.flush()
+
+        if not line:
+            break
+
+        try:
+            command = (
+                json.loads(line.strip())
+                if line.strip().startswith("{") or line.strip().startswith("[")
+                else line.strip()
+            )
+            msg = (
+                "BROADCAST INFO:     Message received from parent process and relayed to active listeners ->"
+                + f" {json.dumps(command)}"
+            )
+            await broadcast_server.broadcast(json.dumps(command))
+            broadcast_server.logger.info(msg)
+        except json.JSONDecodeError:
+            broadcast_server.logger.error("Invalid JSON received from stdin")
+
+
 @app.websocket("/")
 async def websocket_endpoint(  # noqa: PLR0915
     websocket: WebSocket, auth_token: Optional[str] = None
@@ -61,29 +86,35 @@ async def websocket_endpoint(  # noqa: PLR0915
     connected_clients.add(broadcast_server)
 
     stream_task = asyncio.create_task(broadcast_server.stream_results())
+    stdin_task = asyncio.create_task(read_stdin(broadcast_server))
     try:
         await websocket.receive_text()
 
     except WebSocketDisconnect:
         pass
     except Exception as e:
-        broadcast_server.logger.error(f"Unexpected error: {e}")
+        msg = f"Unexpected error: {e.__class__.__name__} -> {e}"
+        broadcast_server.logger.error(msg)
         pass
     finally:
         if broadcast_server in connected_clients:
             connected_clients.remove(broadcast_server)
         stream_task.cancel()
+        stdin_task.cancel()
         try:
             await stream_task
+            await stdin_task
         except asyncio.CancelledError:
             broadcast_server.logger.info("Stream task cancelled")
         except Exception as e:
-            broadcast_server.logger.error(f"Error while cancelling stream task: {e}")
+            msg = f"Unexpected error while cancelling stream task: {e.__class__.__name__} -> {e}"
+            broadcast_server.logger.error(msg)
         if websocket.client_state != WebSocketState.DISCONNECTED:
             try:
                 await websocket.close()
             except RuntimeError as e:
-                broadcast_server.logger.error(f"Error while closing websocket: {e}")
+                msg = f"Unexpected error while closing websocket: {e.__class__.__name__} -> {e}"
+                broadcast_server.logger.error(msg)
 
 
 class BroadcastServer:
@@ -147,7 +178,7 @@ class BroadcastServer:
                                 await self.broadcast(json.dumps(json.loads(message)))
                             last_id = max(row[0] for row in rows)
                     else:
-                        self.logger.error(f"Results file not found: {file_path}")
+                        self.logger.error("Results file not found: %s", str(file_path))
                         break
 
                     await asyncio.sleep(self.sleep_time)
@@ -167,7 +198,8 @@ class BroadcastServer:
         except WebSocketDisconnect:
             pass
         except Exception as e:
-            self.logger.error(f"Unexpected error: {e}")
+            msg = f"Unexpected error: {e.__class__.__name__} -> {e}"
+            self.logger.error(msg)
         finally:
             return
 
@@ -220,9 +252,8 @@ def main():
             **kwargs,
         )
     except TypeError as e:
-        broadcast_server.logger.error(
-            f"Invalid keyword argument passed to unvicorn. -> {e.args[0]}\n"
-        )
+        msg = f"Invalid keyword argument passed to unvicorn. -> {e.args[0]}\n"
+        broadcast_server.logger.error(msg)
     except KeyboardInterrupt:
         broadcast_server.logger.info("Broadcast server terminated.")
     finally:

openbb_platform/extensions/websockets/openbb_websockets/client.py

@@ -224,7 +224,7 @@ class WebSocketClient:
                         self._thread.join()
                         err = ChildProcessError(output)
                         self._exception = err
-                        sys.stdout.write(msg + "\n")
+                        sys.stdout.write(output + "\n")
                         sys.stdout.flush()
                         break
 
@@ -672,6 +672,11 @@ def send_message(
     client, message, target: Literal["provider", "broadcast"] = "provider"
 ) -> None:
     """Send a message to the WebSocketConnection process."""
+    # pylint: disable=import-outside-toplevel
+    import json
+
+    if isinstance(message, (dict, list)):
+        message = json.dumps(message)
     try:
         if target == "provider":
             if client._process and client._process.stdin:

openbb_platform/providers/polygon/openbb_polygon/models/websocket_connection.py

@@ -30,11 +30,13 @@ from pydantic import Field, field_validator, model_validator
 URL_MAP = {
     "stock": "wss://socket.polygon.io/stocks",
     "stock_delayed": "wss://delayed.polygon.io/stocks",
+    "index": "wss://socket.polygon.io/indices",
+    "index_delayed": "wss://delayed.polygon.io/indices",
     "fx": "wss://socket.polygon.io/forex",
     "crypto": "wss://socket.polygon.io/crypto",
 }
 
-ASSET_CHOICES = ["stock", "stock_delayed", "fx", "crypto"]
+ASSET_CHOICES = ["stock", "stock_delayed", "fx", "crypto", "index", "index_delayed"]
 
 FEED_MAP = {
     "crypto": {
@@ -65,6 +67,16 @@ FEED_MAP = {
         "quote": "Q",
         "fmv": "FMV",
     },
+    "index": {
+        "aggs_min": "AM",
+        "aggs_sec": "AS",
+        "value": "V",
+    },
+    "index_delayed": {
+        "aggs_min": "AM",
+        "aggs_sec": "AS",
+        "value": "V",
+    },
 }
 
 
@@ -101,7 +113,9 @@ class PolygonWebSocketQueryParams(WebSocketQueryParams):
     symbol: str = Field(
         description=QUERY_DESCRIPTIONS.get("symbol", ""),
     )
-    asset_type: Literal["stock", "stock_delayed", "fx", "crypto"] = Field(
+    asset_type: Literal[
+        "stock", "stock_delayed", "fx", "crypto", "index", "index_delayed"
+    ] = Field(
         default="crypto",
         description="The asset type associated with the symbol(s)."
         + " Choose from: stock, stock_delayed, fx, crypto.",
@@ -150,7 +164,8 @@ class PolygonCryptoAggsWebSocketData(WebSocketData):
         description="The type of data.",
     )
     date: datetime = Field(
-        description="The start of the aggregate window.",
+        description=DATA_DESCRIPTIONS.get("date", "")
+        + "The end of the aggregate window.",
     )
     symbol: str = Field(
         description=DATA_DESCRIPTIONS.get("symbol", ""),
@@ -464,7 +479,8 @@ class PolygonStockAggsWebSocketData(WebSocketData):
         description="The type of data.",
     )
     date: datetime = Field(
-        description="The start of the aggregate window.",
+        description=DATA_DESCRIPTIONS.get("date", "")
+        + "The end of the aggregate window.",
     )
     symbol: str = Field(
         description=DATA_DESCRIPTIONS.get("symbol", ""),
@@ -547,7 +563,8 @@ class PolygonStockTradeWebSocketData(WebSocketData):
         description="The type of data.",
     )
     date: datetime = Field(
-        description="The start of the aggregate window.",
+        description=DATA_DESCRIPTIONS.get("date", "")
+        + "The end of the aggregate window.",
     )
     symbol: str = Field(
         description=DATA_DESCRIPTIONS.get("symbol", ""),
@@ -641,7 +658,8 @@ class PolygonStockQuoteWebSocketData(WebSocketData):
         description="The type of data.",
     )
     date: datetime = Field(
-        description="The start of the aggregate window.",
+        description=DATA_DESCRIPTIONS.get("date", "")
+        + "The end of the aggregate window.",
     )
     symbol: str = Field(
         description=DATA_DESCRIPTIONS.get("symbol", ""),
@@ -727,6 +745,94 @@ class PolygonStockQuoteWebSocketData(WebSocketData):
         return values
 
 
+class PolygonIndexAggsWebSocketData(WebSocketData):
+    """Polygon Index Aggregates WebSocket data model."""
+
+    __alias_dict__ = {
+        "type": "ev",
+        "symbol": "sym",
+        "date": "e",
+        "day_open": "op",
+        "open": "o",
+        "high": "h",
+        "low": "l",
+        "close": "c",
+    }
+
+    type: str = Field(
+        description="The type of data.",
+    )
+    date: datetime = Field(
+        description=DATA_DESCRIPTIONS.get("date", "")
+        + "The end of the aggregate window.",
+    )
+    symbol: str = Field(
+        description=DATA_DESCRIPTIONS.get("symbol", ""),
+    )
+    day_open: float = Field(
+        description="Today's official opening level.",
+    )
+    open: float = Field(
+        description=DATA_DESCRIPTIONS.get("open", "")
+        + " For the current aggregate window.",
+    )
+    high: float = Field(
+        description=DATA_DESCRIPTIONS.get("high", "")
+        + " For the current aggregate window.",
+    )
+    low: float = Field(
+        description=DATA_DESCRIPTIONS.get("low", "")
+        + " For the current aggregate window.",
+    )
+    close: float = Field(
+        description=DATA_DESCRIPTIONS.get("close", "")
+        + " For the current aggregate window.",
+    )
+
+    @field_validator("date", mode="before", check_fields=False)
+    @classmethod
+    def _validate_date(cls, v):
+        """Validate the date."""
+        return validate_date(cls, v)
+
+    @model_validator(mode="before")
+    @classmethod
+    def _validate_model(cls, values):
+        """Validate the model."""
+        _ = values.pop("s", None)
+        return values
+
+
+class PolygonIndexValueWebSocketData(WebSocketData):
+    """Polygon Index Value WebSocket data model."""
+
+    __alias_dict__ = {
+        "type": "ev",
+        "symbol": "T",
+        "date": "t",
+        "value": "val",
+    }
+
+    type: str = Field(
+        description="The type of data.",
+    )
+    date: datetime = Field(
+        description=DATA_DESCRIPTIONS.get("date", ""),
+    )
+    symbol: str = Field(
+        description=DATA_DESCRIPTIONS.get("symbol", ""),
+    )
+    value: float = Field(
+        description="The value of the index.",
+    )
+
+    @field_validator("date", mode="before", check_fields=False)
+    @classmethod
+    def _validate_date(cls, v):
+        """Validate the date."""
+        return validate_date(cls, v)
+
+
 class PolygonFairMarketValueData(WebSocketData):
     """Polygon Fair Market Value WebSocket Data."""
 
@@ -767,6 +873,8 @@ MODEL_MAP = {
     "AS": PolygonStockAggsWebSocketData,
     "T": PolygonStockTradeWebSocketData,
     "Q": PolygonStockQuoteWebSocketData,
+    "A": PolygonIndexAggsWebSocketData,
+    "V": PolygonIndexValueWebSocketData,
 }
 
 
@@ -775,7 +883,14 @@ class PolygonWebSocketData(Data):
 
     def __new__(cls, **data):
         """Create new instance of appropriate model type."""
-        model = MODEL_MAP.get(data.get("ev")) or MODEL_MAP.get(data.get("type"))
+        index_symbol = data.get("sym", "").startswith("I:") or data.get(
+            "symbol", ""
+        ).startswith("I:")
+        model = (
+            MODEL_MAP["A"]
+            if index_symbol
+            else MODEL_MAP.get(data.get("ev")) or MODEL_MAP.get(data.get("type"))
+        )
         if not model:
             return super().__new__(cls)
 

openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py

@@ -38,6 +38,11 @@ async def handle_symbol(symbol):
     new_symbols: list = []
     feed = FEED_MAP.get(ASSET_TYPE, {}).get(FEED)
     for s in symbols:
+
+        if s == "*":
+            new_symbols.append(f"{feed}.*")
+            continue
+
         if "." in s:
             _check = s.split(".")[0]
             if _check not in list(FEED_MAP.get(ASSET_TYPE, {}).values()):
@@ -48,28 +53,21 @@ async def handle_symbol(symbol):
                 )
                 continue
 
-        if s == "*":
-            new_symbols.append(f"{feed}.*")
-            continue
         ticker = s.upper()
-        if ticker and "." not in ticker and not ticker.startswith(feed):
+
+        if ticker and "." not in ticker:
             ticker = f"{feed}.{ticker}"
-        elif (
-            ASSET_TYPE == "crypto"
-            and "." not in ticker
-            and "-" not in ticker
-            and ticker != "*"
-        ):
+
+        if ASSET_TYPE == "crypto" and "-" not in ticker and ticker != "*":
             ticker = ticker[:3] + "-" + ticker[3:]
-        elif (
-            ASSET_TYPE == "fx"
-            and "." not in ticker
-            and "/" not in ticker
-            and ticker != "*"
-        ):
+        elif ASSET_TYPE == "fx" and "/" not in ticker and ticker != "*":
             ticker = ticker[:3] + "/" + ticker[3:]
         elif ASSET_TYPE == "fx" and "-" in ticker:
             ticker = ticker.replace("-", "/")
+        elif ASSET_TYPE == "index" and ":" not in ticker and ticker != "*":
+            _feed, _ticker = ticker.split(".") if "." in ticker else (feed, ticker)
+            ticker = f"{_feed}.I:{_ticker}"
+
         new_symbols.append(ticker)
 
     return ",".join(new_symbols)
@@ -87,12 +85,12 @@ async def login(websocket, api_key):
                 logger.info("PROVIDER INFO:      %s", msg.get("message"))
                 continue
             if msg.get("status") != "auth_success":
-                err = f"PROVIDER ERROR:    {msg.get('status')} -> {msg.get('message')}"
+                err = f"PROVIDER ERROR:     {msg.get('status')} -> {msg.get('message')}"
                 logger.error(err)
                 sys.exit(1)
             logger.info("PROVIDER INFO:      %s", msg.get("message"))
     except Exception as e:
-        logger.error("PROVIDER ERROR:    %s", e.args[0])
+        logger.error("PROVIDER ERROR:     %s -> %s", e.__class__.__name__, e.args[0])
         sys.exit(1)
 
 
@@ -103,7 +101,7 @@ async def subscribe(websocket, symbol, event):
     try:
         await websocket.send(subscribe_event)
     except Exception as e:
-        msg = f"PROVIDER ERROR:    {e}"
+        msg = f"PROVIDER ERROR:     {e.__class__.__name__} -> {e}"
         logger.error(msg)
 
 
@@ -146,6 +144,8 @@ async def process_message(message, results_path, table_name, limit):
 
             if result:
                 await write_to_db(result, results_path, table_name, limit)
+        else:
+            logger.info("PROVIDER INFO:      %s", msg)
 
 
 async def connect_and_stream(url, symbol, api_key, results_path, table_name, limit):
@@ -255,7 +255,7 @@ if __name__ == "__main__":
         logger.error("PROVIDER ERROR:     WebSocket connection closed")
 
     except Exception as e:  # pylint: disable=broad-except
-        msg = f"PROVIDER ERROR:    {e.args[0]}"
+        msg = f"PROVIDER ERROR:     {e.__class__.__name__} -> {e}"
         logger.error(msg)
 
     finally: