diff --git a/openbb_platform/extensions/websockets/README.md b/openbb_platform/extensions/websockets/README.md index ee9c120a3ce..128d9010734 100644 --- a/openbb_platform/extensions/websockets/README.md +++ b/openbb_platform/extensions/websockets/README.md @@ -33,7 +33,7 @@ obb.websockets ``` > Except for, `get_results`, functions do not return the data or stream. Outputs will be a WebSocketConnectionStatus instance, or a string message. -> All functions, except `create_connection`, assume that a connection has already been establiehd and are referenced by parameters: +> All functions, except `create_connection`, assume that a connection has already been established and are referenced by parameters: > > |Parameter|Type | Required| Description | > |:-------|:-----|:--------:|------------:| @@ -52,7 +52,7 @@ All other endpoints require this to be used first. It is the only function mappi |:-------|:-----|:--------:|------------:| |provider |String |Yes |Name of the provider - i.e, `"polygon"`, `"fmp"`, `"tiingo"` | |name |String |Yes |Name to assign the connection. This is the 'name' parameter in the other endpoints.| -|auth_token |String |No |When supplied, the same token must be passed for all future interactions with the connection, or to read from the broadcast server. | +|auth_token |String |No |When supplied, the same token must be passed as a URL parameter to the broadcast server, and to interact with the client from the API. | |results_file |String |No |Absolute path to the file for continuous writing. Temp file is created by default. Unless 'save_results' is True, discarded on exit. | |save_results |Boolean |No |Whether to persist the file after the session ends, default is `False` | |table_name |String |No |Name of the SQL table to write the results to, consisting of an auto-increment ID and a serialized JSON string of the data. Default is `"records"`| @@ -107,7 +107,7 @@ conn.results.model_dump() ``` ```sh -{'name': 'crypto_tiingo', +{'status': {'name': 'crypto_tiingo', 'auth_required': False, 'subscribed_symbols': '*', 'is_running': True, @@ -117,9 +117,11 @@ conn.results.model_dump() 'broadcast_pid': 5813, 'results_file': '/var/folders/kc/j2lm7bkd5dsfqqnvz259gm6c0000gn/T/tmpwb4jslbg', 'table_name': 'records', - 'save_results': False} + 'save_results': False}} ``` +> From the Python interface, the client is also included in the results. Access it from `results.client` + All of the currently captured data can be dumped with the `get_results` endpoint. The return will be the typical data response object. ```python @@ -450,8 +452,6 @@ PROVIDER INFO: unsubscribed to: C.XAU/USD 'save_results': False} ``` - - ## Development ### Provider Interface @@ -609,7 +609,7 @@ class FmpWebSocketData(WebSocketData): This model is what we return from the `FmpWebSocketFetcher`. The provider will inherit two fields, only the `client` needs to be included on output. No additional fields should be defined. - client - - Instance of WebSocketClient, not returned to from the Router. + - Instance of WebSocketClient, not returned to the API. - status - Leave empty, the Router fills this and returns it. @@ -701,6 +701,8 @@ from openbb_websockets.client import WebSocketClient This is where things diverge slightly. Instead of returning `FmpWebSocketData`, it gets passed to the client connection for validating records as they are received. What gets returned by the Fetcher is the `WebSocketConnection`. +> + ```python class FmpWebSocketFetcher(Fetcher[FmpWebSocketQueryParams, FmpWebSocketConnection]): """FMP WebSocket model.""" @@ -715,7 +717,7 @@ class FmpWebSocketFetcher(Fetcher[FmpWebSocketQueryParams, FmpWebSocketConnectio query: FmpWebSocketQueryParams, credentials: Optional[dict[str, str]], **kwargs: Any, - ) -> WebSocketClient: + ) -> dict: """Extract data from the WebSocket.""" # pylint: disable=import-outside-toplevel import asyncio @@ -766,19 +768,19 @@ class FmpWebSocketFetcher(Fetcher[FmpWebSocketQueryParams, FmpWebSocketConnectio raise e from e # Check if the process is still running before returning. if client.is_running: - return client + return {"client": client} raise OpenBBError("Failed to connect to the WebSocket.") @staticmethod def transform_data( - data: WebSocketClient, + data: dict, query: FmpWebSocketQueryParams, **kwargs: Any, ) -> FmpWebSocketConnection: """Return the client as an instance of Data.""" # All we need to do here is return our client wrapped in the WebSocketConnection class. - return FmpWebSocketConnection(client=data) + return FmpWebSocketConnection(client=data["client"]) ``` #### Map To Router @@ -846,7 +848,7 @@ logger = get_logger("openbb.websocket.fmp") # A UUID gets attached to the name s #### `MessageQueue` -This is an async Queue with an input for the message handler. Create a second instance if a separate queue is required for the subscibe events. +This is an async Queue with an input for the message handler. Create a second instance if a separate queue is required for the subscribe events. Define your async message handler function, and create a task to run in the main event loop. @@ -875,7 +877,7 @@ message = await queue.dequeue() Before submitting the record to `write_to_db`, validate and transform the data with the WebSocketData that was created and imported. Use this function right before transmission, a failure will trigger a termination signal from the main application. ```python -# code above confirms that the message being processed is a data message and not an info message or error. +# code above confirms that the message being processed is a data message and not an info or error message. try: result = FmpWebSocketData.model_validate(message).model_dump_json( @@ -920,8 +922,8 @@ if __name__ == "__main__": for sig in (signal.SIGINT, signal.SIGTERM): loop.add_signal_handler(sig, handle_termination_signal, logger) - # asyncio.run_coroutine_threadsafe(some_connect_and_stream_function, loop) - # loop.run_forever() + asyncio.run_coroutine_threadsafe(some_connect_and_stream_function, loop) + loop.run_forever() ... ``` @@ -930,7 +932,7 @@ if __name__ == "__main__": The missing pieces that get created locally include: - Read `stdin` function for receiving subscribe/unsubscribe events while the connection is running. - - Messages to handle will always have the same format: `'{"event": "subscribe", "symbol": ticker}'` + - Messages to handle will always have the same format: `'{"event": "(un)subscribe", "symbol": ticker}'` - Converting for the symbology used by the provider needs to happen here. - Implementation depends on the requirements of the provider - i.e, how to structure send events. - Create the task before the initial `websockets.connect` block. diff --git a/openbb_platform/extensions/websockets/openbb_websockets/models.py b/openbb_platform/extensions/websockets/openbb_websockets/models.py index 8e53e72233e..ae90ff2352e 100644 --- a/openbb_platform/extensions/websockets/openbb_websockets/models.py +++ b/openbb_platform/extensions/websockets/openbb_websockets/models.py @@ -172,8 +172,8 @@ class WebSocketConnection(Data): @model_validator(mode="before") @classmethod - def _validate_inputs(cls, vaules): + def _validate_inputs(cls, values): """Validate the status.""" - if not vaules.get("status") and not vaules.get("client"): + if not values.get("status") and not values.get("client"): raise ValueError("Cannot initialize empty.") - return vaules + return values diff --git a/openbb_platform/providers/fmp/openbb_fmp/models/websocket_connection.py b/openbb_platform/providers/fmp/openbb_fmp/models/websocket_connection.py index b157e572548..c9fd7a5d765 100644 --- a/openbb_platform/providers/fmp/openbb_fmp/models/websocket_connection.py +++ b/openbb_platform/providers/fmp/openbb_fmp/models/websocket_connection.py @@ -140,7 +140,7 @@ class FmpWebSocketFetcher(Fetcher[FmpWebSocketQueryParams, FmpWebSocketConnectio query: FmpWebSocketQueryParams, credentials: Optional[dict[str, str]], **kwargs: Any, - ) -> WebSocketClient: + ) -> dict: """Extract data from the WebSocket.""" # pylint: disable=import-outside-toplevel import asyncio @@ -183,15 +183,15 @@ class FmpWebSocketFetcher(Fetcher[FmpWebSocketQueryParams, FmpWebSocketConnectio raise e from e if client.is_running: - return client + return {"client": client} raise OpenBBError("Failed to connect to the WebSocket.") @staticmethod def transform_data( - data: WebSocketClient, + data: dict, query: FmpWebSocketQueryParams, **kwargs: Any, ) -> FmpWebSocketConnection: """Return the client as an instance of Data.""" - return FmpWebSocketConnection(client=data) + return FmpWebSocketConnection(client=data["client"]) diff --git a/openbb_platform/providers/fmp/openbb_fmp/utils/websocket_client.py b/openbb_platform/providers/fmp/openbb_fmp/utils/websocket_client.py index b54890874bc..75bfb8915da 100644 --- a/openbb_platform/providers/fmp/openbb_fmp/utils/websocket_client.py +++ b/openbb_platform/providers/fmp/openbb_fmp/utils/websocket_client.py @@ -48,7 +48,7 @@ async def login(websocket, api_key): msg = message.get("message") logger.info("PROVIDER INFO: %s", msg) except Exception as e: # pylint: disable=broad-except - msg = f"PROVIDER ERROR: {e.__class__.__name__}: {e.__str__()}" + msg = f"PROVIDER ERROR: {e.__class__.__name__}: {e}" logger.error(msg) sys.exit(1) @@ -66,7 +66,7 @@ async def subscribe(websocket, symbol, event): await websocket.send(json.dumps(subscribe_event)) await asyncio.sleep(1) except Exception as e: # pylint: disable=broad-except - msg = f"PROVIDER ERROR: {e.__class__.__name__}: {e.__str__()}" + msg = f"PROVIDER ERROR: {e.__class__.__name__}: {e}" logger.error(msg) @@ -105,6 +105,10 @@ async def process_message(message, results_path, table_name, limit): if "you are not authorized" in message.get("message", "").lower(): msg = f"UnauthorizedError -> FMP Message: {message['message']}" logger.error(msg) + elif "Connected from another location" in message.get("message", ""): + msg = f"UnauthorizedError -> FMP Message: {message.get('message')}" + logger.info(msg) + sys.exit(0) else: msg = f"PROVIDER INFO: {message.get('message')}" logger.info(msg) @@ -156,9 +160,7 @@ async def connect_and_stream(url, symbol, api_key, results_path, table_name, lim await asyncio.shield(queue.enqueue(json.loads(message))) except websockets.ConnectionClosed as e: - msg = ( - f"PROVIDER INFO: The WebSocket connection was closed -> {e.__str__()}" - ) + msg = f"PROVIDER INFO: The WebSocket connection was closed -> {e}" logger.info(msg) # Attempt to reopen the connection logger.info("PROVIDER INFO: Attempting to reconnect after five seconds.") @@ -170,7 +172,7 @@ async def connect_and_stream(url, symbol, api_key, results_path, table_name, lim sys.exit(1) except Exception as e: - msg = f"PROVIDER ERROR: Unexpected error -> {e.__class__.__name__}: {e.__str__()}" + msg = f"PROVIDER ERROR: Unexpected error -> {e.__class__.__name__}: {e}" logger.error(msg) sys.exit(1) @@ -207,7 +209,7 @@ if __name__ == "__main__": logger.error("PROVIDER ERROR: WebSocket connection closed") except Exception as e: # pylint: disable=broad-except - msg = f"PROVIDER ERROR: {e.__class__.__name__}: {e.__str__()}" + msg = f"PROVIDER ERROR: {e.__class__.__name__}: {e}" logger.error(msg) finally: diff --git a/openbb_platform/providers/fmp/tests/test_fmp_fetchers.py b/openbb_platform/providers/fmp/tests/test_fmp_fetchers.py index 60074329bc6..42d48f6e8f3 100644 --- a/openbb_platform/providers/fmp/tests/test_fmp_fetchers.py +++ b/openbb_platform/providers/fmp/tests/test_fmp_fetchers.py @@ -70,6 +70,7 @@ from openbb_fmp.models.revenue_geographic import FMPRevenueGeographicFetcher from openbb_fmp.models.risk_premium import FMPRiskPremiumFetcher from openbb_fmp.models.share_statistics import FMPShareStatisticsFetcher from openbb_fmp.models.treasury_rates import FMPTreasuryRatesFetcher +from openbb_fmp.models.websocket_connection import FmpWebSocketFetcher from openbb_fmp.models.world_news import FMPWorldNewsFetcher from openbb_fmp.models.yield_curve import FMPYieldCurveFetcher @@ -777,3 +778,19 @@ def test_fmp_historical_market_cap_fetcher(credentials=test_credentials): fetcher = FmpHistoricalMarketCapFetcher() result = fetcher.test(params, credentials) assert result is None + + +@pytest.mark.record_screen +def test_fmp_websocket_fetcher(credentials=test_credentials): + """Test FMP Websocket fetcher.""" + import asyncio + + params = {"symbol": "btcusd", "name": "test", "limit": 10, "asset_type": "crypto"} + + try: + fetcher = FmpWebSocketFetcher() + result = fetcher.test(params, credentials) + response = asyncio.run(fetcher.fetch_data(params, credentials)) + assert result is None + finally: + response.client.disconnect() 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 0960bb4850d..48d3f708ffd 100644 --- a/openbb_platform/providers/polygon/openbb_polygon/models/websocket_connection.py +++ b/openbb_platform/providers/polygon/openbb_polygon/models/websocket_connection.py @@ -1219,7 +1219,7 @@ class PolygonWebSocketFetcher( query: PolygonWebSocketQueryParams, credentials: Optional[dict[str, str]], **kwargs: Any, - ) -> WebSocketClient: + ) -> dict: """Extract data from the WebSocket.""" api_key = credentials.get("polygon_api_key") if credentials else "" url = URL_MAP[query.asset_type] @@ -1261,15 +1261,15 @@ class PolygonWebSocketFetcher( raise client._exception from client._exception if client.is_running: - return client + return {"client": client} raise OpenBBError("Failed to connect to the WebSocket.") @staticmethod def transform_data( - data: WebSocketClient, + data: dict, query: PolygonWebSocketQueryParams, **kwargs: Any, ) -> PolygonWebSocketConnection: """Return the client as an instance of Data.""" - return PolygonWebSocketConnection(client=data) + return PolygonWebSocketConnection(client=data["client"]) 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 e22421abe2b..cdb706542ee 100644 --- a/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py +++ b/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py @@ -233,27 +233,23 @@ async def connect_and_stream(url, symbol, api_key, results_path, table_name, lim except websockets.InvalidStatusCode as e: if e.status_code == 404: - msg = f"PROVIDER ERROR: {e.__str__()}" + msg = f"PROVIDER ERROR: {e}" logger.error(msg) sys.exit(1) else: raise except websockets.InvalidURI as e: - msg = f"PROVIDER ERROR: {e.__str__()}" + msg = f"PROVIDER ERROR: {e}" logger.error(msg) sys.exit(1) except websockets.ConnectionClosedOK as e: - msg = ( - f"PROVIDER INFO: The WebSocket connection was closed -> {e.__str__()}" - ) + msg = f"PROVIDER INFO: The WebSocket connection was closed -> {e}" logger.info(msg) sys.exit(0) except websockets.ConnectionClosed as e: - msg = ( - f"PROVIDER INFO: The WebSocket connection was closed -> {e.__str__()}" - ) + msg = f"PROVIDER INFO: The WebSocket connection was closed -> {e}" logger.info(msg) # Attempt to reopen the connection logger.info("PROVIDER INFO: Attempting to reconnect after five seconds.") @@ -261,12 +257,12 @@ async def connect_and_stream(url, symbol, api_key, results_path, table_name, lim await connect_and_stream(url, symbol, api_key, results_path, table_name, limit) except websockets.WebSocketException as e: - msg = f"PROVIDER ERROR: WebSocketException -> {e.__str__()}" + msg = f"PROVIDER ERROR: WebSocketException -> {e}" logger.error(msg) sys.exit(1) except Exception as e: - msg = f"PROVIDER ERROR: Unexpected error -> {e.__class__.__name__}: {e.__str__()}" + msg = f"PROVIDER ERROR: Unexpected error -> {e.__class__.__name__}: {e}" logger.error(msg) sys.exit(1) diff --git a/openbb_platform/providers/tiingo/openbb_tiingo/models/websocket_connection.py b/openbb_platform/providers/tiingo/openbb_tiingo/models/websocket_connection.py index 3b083645952..ffeb3a3781d 100644 --- a/openbb_platform/providers/tiingo/openbb_tiingo/models/websocket_connection.py +++ b/openbb_platform/providers/tiingo/openbb_tiingo/models/websocket_connection.py @@ -233,7 +233,7 @@ class TiingoWebSocketFetcher( query: TiingoWebSocketQueryParams, credentials: Optional[dict[str, str]], **kwargs: Any, - ) -> WebSocketClient: + ) -> dict: """Initiailze the WebSocketClient and connect.""" # pylint: disable=import-outside-toplevel from asyncio import sleep @@ -291,16 +291,16 @@ class TiingoWebSocketFetcher( raise OpenBBError(exc) if client.is_running: - return client + return {"client": client} client.disconnect() raise OpenBBError("Failed to connect to the WebSocket.") @staticmethod def transform_data( - data: WebSocketClient, + data: dict, query: TiingoWebSocketQueryParams, **kwargs: Any, ) -> TiingoWebSocketConnection: """Return the client as an instance of Data.""" - return TiingoWebSocketConnection(client=data) + return TiingoWebSocketConnection(client=data["client"]) diff --git a/openbb_platform/providers/tiingo/openbb_tiingo/utils/websocket_client.py b/openbb_platform/providers/tiingo/openbb_tiingo/utils/websocket_client.py index 50ca68c09c1..25243c9dc55 100644 --- a/openbb_platform/providers/tiingo/openbb_tiingo/utils/websocket_client.py +++ b/openbb_platform/providers/tiingo/openbb_tiingo/utils/websocket_client.py @@ -236,9 +236,7 @@ async def connect_and_stream( sys.exit(1) except websockets.ConnectionClosed as e: - msg = ( - f"PROVIDER INFO: The WebSocket connection was closed -> {e.__str__()}" - ) + msg = f"PROVIDER INFO: The WebSocket connection was closed -> {e}" logger.info(msg) # Attempt to reopen the connection logger.info("PROVIDER INFO: Attempting to reconnect after five seconds...") @@ -248,11 +246,11 @@ async def connect_and_stream( ) except websockets.WebSocketException as e: - logger.error(str(e)) - sys.exit(1) + logger.info(str(e)) + sys.exit(0) except Exception as e: - msg = f"Unexpected error -> {e.__class__.__name__}: {e.__str__()}" + msg = f"Unexpected error -> {e.__class__.__name__}: {e}" logger.error(msg) sys.exit(1) @@ -286,7 +284,7 @@ if __name__ == "__main__": loop.run_forever() except Exception as e: # pylint: disable=broad-except - msg = f"Unexpected error -> {e.__class__.__name__}: {e.__str__()}" + msg = f"Unexpected error -> {e.__class__.__name__}: {e}" logger.error(msg) finally: