diff --git a/openbb_platform/extensions/websockets/openbb_websockets/broadcast.py b/openbb_platform/extensions/websockets/openbb_websockets/broadcast.py index 9235a6b4bdf..29bc7b739b3 100644 --- a/openbb_platform/extensions/websockets/openbb_websockets/broadcast.py +++ b/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: diff --git a/openbb_platform/extensions/websockets/openbb_websockets/client.py b/openbb_platform/extensions/websockets/openbb_websockets/client.py index b045aec1896..743d58d51c7 100644 --- a/openbb_platform/extensions/websockets/openbb_websockets/client.py +++ b/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: diff --git a/openbb_platform/providers/polygon/openbb_polygon/models/websocket_connection.py b/openbb_platform/providers/polygon/openbb_polygon/models/websocket_connection.py index 8c4ac14a3e6..f7d8d06a07b 100644 --- a/openbb_platform/providers/polygon/openbb_polygon/models/websocket_connection.py +++ b/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) diff --git a/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py b/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py index ba8feeee0e1..39cbf0dd6a7 100644 --- a/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py +++ b/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: