diff --git a/openbb_platform/extensions/websockets/openbb_websockets/client.py b/openbb_platform/core/openbb_core/provider/utils/websockets/client.py similarity index 76% rename from openbb_platform/extensions/websockets/openbb_websockets/client.py rename to openbb_platform/core/openbb_core/provider/utils/websockets/client.py index 82b04254c1c..35accd184b9 100644 --- a/openbb_platform/extensions/websockets/openbb_websockets/client.py +++ b/openbb_platform/core/openbb_core/provider/utils/websockets/client.py @@ -1,48 +1,60 @@ -"""WebSocket Client module for interacting with a provider websocket in a non-blocking pattern.""" +"""Module for running OpenBB Provider websocket connection scripts.""" # pylint: disable=too-many-statements,protected-access # flake8: noqa: PLR0915 -import logging from typing import TYPE_CHECKING, Any, Literal, Optional, Union if TYPE_CHECKING: - from openbb_core.provider.abstract.data import Data + import logging + + from pydantic import BaseModel class WebSocketClient: # pylint: disable=too-many-instance-attributes - """Client for interacting with a websocket server in a non-blocking pattern. + """ + Client for interacting with a websocket server in a non-blocking pattern, and handling the subprocesses. Parameters ---------- name : str - Name to assign the WebSocket connection. Used to identify and manage multiple instances. + Name to assign the WebSocket connection. Used to identify and manage multiple instances from the API. module : str - The Python module for the provider websocket_client module. Runs in a separate thread. - Example: 'openbb_fmp.utils.websocket_client'. Pass additional keyword arguments by including kwargs. + The Python module for the provider websocket_client module. + Runs in a separate thread, and is an equivalent to 'python -m module'. + Example: 'openbb_fmp.utils.websocket_client'. + Pass additional keyword arguments to the script by including **kwargs. symbol : Optional[str] - The symbol(s) requested to subscribe. Enter multiple symbols separated by commas without spaces. + The symbol(s) requested to subscribe on start. Enter multiple symbols separated by commas, without spaces. + Where supported by the provider, * represents all symbols within the feed. limit : Optional[int] The limit of records to hold in memory. Once the limit is reached, the oldest records are removed. - Default is 1000. Set to None to keep all records. + Default is 5000. Set to None to keep all records. results_file : Optional[str] Absolute path to the file for continuous writing. By default, a temporary file is created. + File is discarded when the Python session ends unless 'save_results' is set to True. + The connection can be re-established with the same results file to continue writing. + EACH NEW CONNECTION SHOULD HAVE A UNIQUE RESULTS FILE. + If the intention is to permanently store the results for historical records, + save the current session to a new file and copy new records at periodic intervals into the master. table_name : Optional[str] SQL table name to store serialized data messages. By default, 'records'. save_results : bool Whether to persist the results after the main Python session ends. Default is False. - data_model : Optional[Data] + data_model : Optional[BaseModel] Pydantic data model to validate the results before storing them in the database. Also used to deserialize the results from the database. auth_token : Optional[str] - The authentication token to use for the WebSocket connection. Default is None. - Only used for API and Python application endpoints. + Used to limit access to the broadcast stream. When provided, listeners should supply as a URL parameter. + Example: 'ws://127.0.0.1:6666/?auth_token=SOME_TOKEN'. + When provided, the auth_token is required to interact with the instance of this class + from the REST API and Python application endpoints. logger : Optional[logging.Logger] - The pre-configured logger instance to use for this connection. By default, a new logger is created. + A pre-configured logger to use for this instance. By default, a new logger is created. kwargs : Optional[dict] Additional keyword arguments to pass to the target provider module. Keywords and values must not contain spaces. To pass items to 'websocket.connect()', include them in the 'kwargs' dictionary as, - {'connect_kwargs': {'key': 'value'}}. + {'api_key': 'MY_API_KEY', 'connect_kwargs': {'key': 'value'}}. Properties ---------- @@ -55,10 +67,9 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes is_broadcasting : bool Check if the broadcast server process is running. broadcast_address : str - URI address for the results broadcast server. + URI address for connecting to the broadcast stream. results : list[Data] - All stored results from the provider's WebSocket stream. - Results are stored in a SQLite database as a serialized JSON string, this property deserializes the results. + All stored results from the provider connection. Clear the results by deleting the property. e.g., del client.results Methods @@ -68,15 +79,17 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes disconnect Disconnect from the provider WebSocket. subscribe - Subscribe to a new symbol or list of symbols. + Send a subscribe message to the provider connection. unsubscribe - Unsubscribe from a symbol or list of symbols. + Send an unsubscribe message to the provider connection. start_broadcasting - Start the broadcast server to stream results over a network connection. + Start the broadcast server to stream results over a network. stop_broadcasting Stop the broadcast server and disconnect all listening clients. send_message Send a message to the WebSocket process. Messages can be sent to "provider" or "broadcast" targets. + query_database + Run a SELECT query to the database. Returns a list of deserialized results, or validated models. """ def __init__( # noqa: PLR0913 # pylint: disable=too-many-arguments,too-many-positional-arguments,too-many-locals @@ -84,13 +97,13 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes name: str, module: str, symbol: Optional[str] = None, - limit: Optional[int] = 1000, + limit: Optional[int] = 5000, results_file: Optional[str] = None, table_name: Optional[str] = None, save_results: bool = False, - data_model: Optional["Data"] = None, + data_model: Optional["BaseModel"] = None, auth_token: Optional[str] = None, - logger: Optional[logging.Logger] = None, + logger: Optional["logging.Logger"] = None, **kwargs, ) -> None: """Initialize the WebSocketClient class.""" @@ -162,20 +175,21 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes results_file=self.results_file, table_name=self.table_name, limit=self._limit, - keep_results=self.save_results, logger=self.logger, data_model=self.data_model, ) except Exception as e: # pylint: disable=broad-except msg = ( "Unexpected error setting up the SQLite database and table ->" - f" {e.__class__.__name__ if hasattr(e, '__class__') else e.__name__}: {e}" + f" {e.__class__.__name__ if hasattr(e, '__class__') else e.__name__} -> {e}" ) self.logger.error(msg) self._exception = OpenBBError(msg) + self._atexit() + raise OpenBBError(msg) from e def _atexit(self) -> None: - """Clean up the WebSocket client processes at exit.""" + """Clean up the running processes at exit.""" # pylint: disable=import-outside-toplevel import os @@ -310,13 +324,23 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes output = output.replace("INFO:", "BROADCAST INFO:") + "\n" output = output[0] if isinstance(output, tuple) else output output = clean_message(output) + if ( + output.startswith("BROADCAST ERROR:") + or "unexpected error" in output.lower() + ): + self._psutil_broadcast_process.kill() + self._broadcast_process.wait() + self._broadcast_thread.join() + sys.stdout.write(output + "\n") + sys.stdout.flush() + break sys.stdout.write(output + "\n") sys.stdout.flush() except queue.Empty: continue def connect(self) -> None: # pylint: disable=too-many-locals - """Connect to the provider WebSocket.""" + """Connect to the provider client connection.""" # pylint: disable=import-outside-toplevel import os # noqa import psutil @@ -367,35 +391,44 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes if kwarg not in command: command.extend([kwarg]) - self._process = subprocess.Popen( # noqa # pylint: disable=consider-using-with - command, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - stdin=subprocess.PIPE, - env=os.environ, - text=True, - bufsize=1, - ) - self._psutil_process = psutil.Process(self._process.pid) + try: + self._process = ( + subprocess.Popen( # noqa # pylint: disable=consider-using-with + command, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + stdin=subprocess.PIPE, + env=os.environ, + text=True, + bufsize=1, + ) + ) + self._psutil_process = psutil.Process(self._process.pid) - log_output_queue: queue.Queue = queue.Queue() - self._thread = threading.Thread( - target=non_blocking_websocket, - args=( - self, - log_output_queue, - self._provider_message_queue, - ), - ) - self._thread.daemon = True - self._thread.start() + log_output_queue: queue.Queue = queue.Queue() + self._thread = threading.Thread( + target=non_blocking_websocket, + args=( + self, + log_output_queue, + self._provider_message_queue, + ), + ) + self._thread.daemon = True + self._thread.start() - self._log_thread = threading.Thread( - target=self._log_provider_output, - args=(log_output_queue,), - ) - self._log_thread.daemon = True - self._log_thread.start() + self._log_thread = threading.Thread( + target=self._log_provider_output, + args=(log_output_queue,), + ) + self._log_thread.daemon = True + self._log_thread.start() + + except Exception as e: # pylint: disable=broad-except + msg = f"Unexpected error -> {e.__class__.__name__ if hasattr(e, '__class__') else e.__name__} -> {e}" + self.logger.error(msg) + self._atexit() + raise OpenBBError(msg) from e # Give it some startup time to allow the connection to be established and for exceptions to populate. time.sleep(2) @@ -406,12 +439,14 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes raise OpenBBError(exc) if not self.is_running: - self.logger.error("The provider server failed to start.") + self.logger.error( + "Unexpected error -> Provider connection process failed to start." + ) def send_message( self, message, target: Literal["provider", "broadcast"] = "provider" ) -> None: - """Send a message to the WebSocket process.""" + """Write to the provider, or broadcast, process stdin.""" if target == "provider": self._provider_message_queue.put(message) read_message_queue(self, self._provider_message_queue) @@ -420,10 +455,10 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes read_message_queue(self, self._broadcast_message_queue, target="broadcast") def disconnect(self) -> None: - """Disconnect from the provider WebSocket.""" + """Disconnect from the provider connection.""" self._stop_log_thread_event.set() if self._process is None or self.is_running is False: - self.logger.info("Not connected to the provider WebSocket.") + self.logger.info("Provider client connection is not running.") return if ( self._psutil_process is not None @@ -435,13 +470,18 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes self._thread.join() self._log_thread.join() self._stop_log_thread_event.clear() - self.logger.info("Disconnected from the provider WebSocket.") + self.logger.info("Disconnected from the provider server.") if hasattr(self, "_exception") and self._exception: raise self._exception return def subscribe(self, symbol) -> None: - """Subscribe to a new symbol or list of symbols.""" + """ + Send a subscribe message to the active provider connection. + + Messages are sent as JSON strings formatted as: + {"event": "subscribe", "symbol": "AAPL,MSFT"} + """ # pylint: disable=import-outside-toplevel import json # noqa import time @@ -463,7 +503,12 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes self._symbol = ",".join(new_symbols) def unsubscribe(self, symbol) -> None: - """Unsubscribe from a symbol or list of symbols.""" + """ + Unsubscribe from a symbol or list of symbols. + + Messages are sent as JSON strings formatted as: + {"event": "unsubscribe", "symbol": "AAPL,MSFT"} + """ # pylint: disable=import-outside-toplevel import json # noqa import time @@ -500,23 +545,33 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes @property def results(self) -> list: - """Retrieve the deserialized results from the results file.""" + """ + Retrieve the deserialized results from the active Database. + + Clear the results by deleting the property. e.g., del client.results + """ # pylint: disable=import-outside-toplevel from openbb_core.app.model.abstract.error import OpenBBError try: return self.database.fetch_all() except Exception as e: # pylint: disable=broad-except - raise OpenBBError(f"Error retrieving results: {e}") from e + msg = ( + "Error retrieving results:" + f" {e.__class__.__name__ if hasattr(e,'__class__') else e.__name__} -> {e}" + ) + raise OpenBBError(msg) from e @results.deleter def results(self): - """Clear results stored from the WebSocket stream.""" - # pylint: disable=import-outside-toplevel + """Clear results stored by the active WebSocket stream.""" try: self.database.clear_results() except Exception as e: # pylint: disable=broad-except - msg = f"Error clearing results: {e.__class__.__name__}: {e}" + msg = ( + "Error clearing results:" + f" {e.__class__.__name__ if hasattr(e, '__class__') else e.__name__} -> {e}" + ) self.logger.error(msg) @property @@ -560,13 +615,63 @@ class WebSocketClient: # pylint: disable=too-many-instance-attributes else None ) + def query_database( + self, + sql: Optional[str] = None, + limit: Optional[int] = 25, + ) -> list: + """ + Make a SELECT query to the database for results. + + The database always contains two columns: + "id" - an auto-incrementing ID + "message" - a JSON serialized row of data + + Parameters + ---------- + sql : Optional[str] + SQL query to execute. Default is None. + limit : Optional[int] + Limit the number of records returned, by most recent. Default is 25, set to None to return all records. + + Returns + ------- + list + A list of deserialized results from the database. + If a 'data_model' was supplied at initialization, it will be a list of validated models. + """ + if not sql: + query = f"SELECT message FROM {self.table_name} ORDER BY id DESC" # noqa + if limit is not None: + query += f" LIMIT {limit};" + else: + query = ( + sql.replace(";", "") + f" LIMIT {limit}" + if limit is not None and "LIMIT" not in sql.upper() + else sql + ) + + return self.database.query(query) + def start_broadcasting( # pylint: disable=too-many-locals self, host: str = "127.0.0.1", port: int = 6666, **kwargs, ) -> None: - """Broadcast results over a network connection.""" + """ + Broadcast results over a network connection. + + Parameters + ---------- + host : str + The host address to broadcast results to. Default is 127.0.0.1 + port : int + The port to broadcast results to. Default is 6666 + If the port is already in use, the next available port is used. + **kwargs: dict + Additional keyword arguments to pass to the `uvicorn.run`. + """ # pylint: disable=import-outside-toplevel import os # noqa import subprocess @@ -692,9 +797,10 @@ def non_blocking_websocket(client, output_queue, provider_message_queue) -> None if output: output_queue.put(output.strip()) - except Exception as e: + except Exception as e: # pylint: disable=broad-except msg = ( - f"Unexpected error in non_blocking_websocket: {e.__class__.__name__} -> {e}" + "Unexpected error in non_blocking_websocket:" + f" {e.__class__.__name__ if hasattr(e, '__class__') else e.__name__} -> {e}" ) client.logger.error(msg) raise e from e @@ -710,7 +816,7 @@ def send_message( # pylint: disable=import-outside-toplevel import json - if isinstance(message, (dict, list)): + if not isinstance(message, str): message = json.dumps(message) try: if target == "provider": @@ -726,7 +832,10 @@ def send_message( else: client.logger.error("Broadcast process is not running.") except Exception as e: # pylint: disable=broad-except - msg = f"Error sending message to the {target} process: {e.__class__.__name__} -> {e}" + msg = ( + f"Error sending message to the {target} process:" + f" {e.__class__.__name__ if hasattr(e, '__class__') else e.__name__} -> {e}" + ) client.logger.error(msg) @@ -747,8 +856,8 @@ def read_message_queue( send_message(client, message, target="broadcast") except Exception as e: # pylint: disable=broad-except err = ( - f"Error while attempting to transmit from the outgoing message queue: {e.__class__.__name__} " - f"-> {e} -> {message}" + "Error while attempting to transmit from the outgoing message queue:" + f"{e.__class__.__name__ if hasattr(e, '__class__') else e.__name__} -> {e} -> {message}" ) client.logger.error(err) @@ -767,7 +876,8 @@ def non_blocking_broadcast(client, output_queue, broadcast_message_queue) -> Non output_queue.put(output.strip()) except Exception as e: # pylint: disable=broad-except err = ( - f"Unexpected error in non_blocking_broadcast: {e.__class__.__name__} -> {e}" + f"Unexpected error in non_blocking_broadcast:" + f" {e.__class__.__name__ if hasattr(e, '__class__') else e.__name__} -> {e}" ) client.logger.error(err) finally: diff --git a/openbb_platform/core/openbb_core/provider/utils/websockets/database.py b/openbb_platform/core/openbb_core/provider/utils/websockets/database.py index 4765cc06806..19968f88e68 100644 --- a/openbb_platform/core/openbb_core/provider/utils/websockets/database.py +++ b/openbb_platform/core/openbb_core/provider/utils/websockets/database.py @@ -66,7 +66,6 @@ class Database: table_name: Optional[str] = None, data_model: Optional["BaseModel"] = None, limit: Optional[int] = None, - keep_results: Optional[bool] = False, logger: Optional["logging.Logger"] = None, loop: Optional["asyncio.AbstractEventLoop"] = None, **kwargs, @@ -87,6 +86,10 @@ class Database: temp_file_path = temp_file.name self.results_path = Path(temp_file_path).absolute() self.results_file = temp_file_path + elif results_file and "://" in results_file: + self.results_file = results_file + self.results_path = results_file + kwargs["uri"] = True else: self.results_path = Path(results_file).absolute() self.results_file = results_file diff --git a/openbb_platform/extensions/websockets/openbb_websockets/models.py b/openbb_platform/extensions/websockets/openbb_websockets/models.py index 6912f743dc4..e2bfb02314e 100644 --- a/openbb_platform/extensions/websockets/openbb_websockets/models.py +++ b/openbb_platform/extensions/websockets/openbb_websockets/models.py @@ -9,10 +9,9 @@ from openbb_core.provider.abstract.query_params import QueryParams from openbb_core.provider.utils.descriptions import ( DATA_DESCRIPTIONS, ) +from openbb_core.provider.utils.websockets.client import WebSocketClient from pydantic import ConfigDict, Field, field_validator, model_validator -from openbb_websockets.client import WebSocketClient - class WebSocketQueryParams(QueryParams): """Query parameters for WebSocket connection creation.""" 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 b9779e5f0d2..fa305da597e 100644 --- a/openbb_platform/providers/fmp/openbb_fmp/models/websocket_connection.py +++ b/openbb_platform/providers/fmp/openbb_fmp/models/websocket_connection.py @@ -7,7 +7,7 @@ from typing import Any, Literal, Optional from openbb_core.app.model.abstract.error import OpenBBError from openbb_core.provider.abstract.fetcher import Fetcher -from openbb_websockets.client import WebSocketClient +from openbb_core.provider.utils.websockets.client import WebSocketClient from openbb_websockets.models import ( WebSocketConnection, WebSocketData, @@ -15,12 +15,6 @@ from openbb_websockets.models import ( ) from pydantic import Field, field_validator -URL_MAP = { - "stock": "wss://websockets.financialmodelingprep.com", - "fx": "wss://forex.financialmodelingprep.com", - "crypto": "wss://crypto.financialmodelingprep.com", -} - class FmpWebSocketQueryParams(WebSocketQueryParams): """FMP WebSocket query parameters.""" @@ -151,12 +145,10 @@ class FmpWebSocketFetcher(Fetcher[FmpWebSocketQueryParams, FmpWebSocketConnectio import asyncio api_key = credentials.get("fmp_api_key") if credentials else "" - url = URL_MAP[query.asset_type] symbol = query.symbol.lower() - kwargs = { - "url": url, + "asset_type": query.asset_type, "api_key": api_key, "connect_kwargs": query.connect_kwargs, } 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 67b06be549f..44a0240f70e 100644 --- a/openbb_platform/providers/fmp/openbb_fmp/utils/websocket_client.py +++ b/openbb_platform/providers/fmp/openbb_fmp/utils/websocket_client.py @@ -18,11 +18,20 @@ from openbb_core.provider.utils.websockets.message_queue import MessageQueue from openbb_fmp.models.websocket_connection import FmpWebSocketData from pydantic import ValidationError +URL_MAP = { + "stock": "wss://websockets.financialmodelingprep.com", + "fx": "wss://forex.financialmodelingprep.com", + "crypto": "wss://crypto.financialmodelingprep.com", +} + logger = get_logger("openbb.websocket.fmp") kwargs = parse_kwargs() queue = MessageQueue() command_queue = MessageQueue() CONNECT_KWARGS = kwargs.pop("connect_kwargs", {}) +URL = URL_MAP.get(kwargs.pop("asset_type"), None) +if not URL: + raise ValueError("Invalid asset type provided.") DATABASE = Database( results_file=kwargs["results_file"], @@ -143,7 +152,7 @@ async def connect_and_stream(): stdin_task = asyncio.create_task(read_stdin_and_queue_commands()) try: - websocket = await websockets.connect(kwargs["url"], **CONNECT_KWARGS) + websocket = await websockets.connect(URL, **CONNECT_KWARGS) await login(websocket) await subscribe(websocket, kwargs["symbol"], "subscribe") diff --git a/openbb_platform/providers/fmp/tests/test_fmp_fetchers.py b/openbb_platform/providers/fmp/tests/test_fmp_fetchers.py index ec5853ca9ca..ab0b59402a9 100644 --- a/openbb_platform/providers/fmp/tests/test_fmp_fetchers.py +++ b/openbb_platform/providers/fmp/tests/test_fmp_fetchers.py @@ -7,6 +7,7 @@ from unittest.mock import MagicMock, patch import pytest from openbb_core.app.service.user_service import UserService +from openbb_core.provider.utils.websockets.client import WebSocketClient from openbb_fmp.models.analyst_estimates import FMPAnalystEstimatesFetcher from openbb_fmp.models.available_indices import FMPAvailableIndicesFetcher from openbb_fmp.models.balance_sheet import FMPBalanceSheetFetcher @@ -79,7 +80,6 @@ from openbb_fmp.models.websocket_connection import ( ) from openbb_fmp.models.world_news import FMPWorldNewsFetcher from openbb_fmp.models.yield_curve import FMPYieldCurveFetcher -from openbb_websockets.client import WebSocketClient test_credentials = UserService().default_user_settings.credentials.model_dump( mode="json" diff --git a/openbb_platform/providers/intrinio/openbb_intrinio/models/websocket_connection.py b/openbb_platform/providers/intrinio/openbb_intrinio/models/websocket_connection.py index 1672ccdbd35..850a7c792dd 100644 --- a/openbb_platform/providers/intrinio/openbb_intrinio/models/websocket_connection.py +++ b/openbb_platform/providers/intrinio/openbb_intrinio/models/websocket_connection.py @@ -7,8 +7,8 @@ from typing import Any, Literal, Optional from openbb_core.app.model.abstract.error import OpenBBError from openbb_core.provider.abstract.fetcher import Fetcher +from openbb_core.provider.utils.websockets.client import WebSocketClient from openbb_intrinio.utils.references import TRADE_CONDITIONS, VENUES -from openbb_websockets.client import WebSocketClient from openbb_websockets.models import ( WebSocketConnection, WebSocketData, diff --git a/openbb_platform/providers/intrinio/tests/test_intrinio_fetchers.py b/openbb_platform/providers/intrinio/tests/test_intrinio_fetchers.py index a49177d2c89..dd0e3d60f12 100644 --- a/openbb_platform/providers/intrinio/tests/test_intrinio_fetchers.py +++ b/openbb_platform/providers/intrinio/tests/test_intrinio_fetchers.py @@ -7,6 +7,7 @@ from unittest.mock import MagicMock, patch import pytest from openbb_core.app.service.user_service import UserService +from openbb_core.provider.utils.websockets.client import WebSocketClient from openbb_intrinio.models.balance_sheet import IntrinioBalanceSheetFetcher from openbb_intrinio.models.calendar_ipo import IntrinioCalendarIpoFetcher from openbb_intrinio.models.cash_flow import IntrinioCashFlowStatementFetcher @@ -73,7 +74,6 @@ from openbb_intrinio.models.websocket_connection import ( IntrinioWebSocketFetcher, ) from openbb_intrinio.models.world_news import IntrinioWorldNewsFetcher -from openbb_websockets.client import WebSocketClient test_credentials = UserService().default_user_settings.credentials.model_dump( mode="json" 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 a94daca58c4..3277d498546 100644 --- a/openbb_platform/providers/polygon/openbb_polygon/models/websocket_connection.py +++ b/openbb_platform/providers/polygon/openbb_polygon/models/websocket_connection.py @@ -9,6 +9,7 @@ from openbb_core.app.model.abstract.error import OpenBBError from openbb_core.provider.abstract.data import Data from openbb_core.provider.abstract.fetcher import Fetcher from openbb_core.provider.utils.descriptions import DATA_DESCRIPTIONS +from openbb_core.provider.utils.websockets.client import WebSocketClient from openbb_polygon.utils.constants import ( CRYPTO_EXCHANGE_MAP, FX_EXCHANGE_MAP, @@ -20,7 +21,6 @@ from openbb_polygon.utils.constants import ( STOCK_TRADE_CONDITIONS, ) from openbb_polygon.utils.helpers import map_tape -from openbb_websockets.client import WebSocketClient from openbb_websockets.models import ( WebSocketConnection, WebSocketData, @@ -28,16 +28,6 @@ from openbb_websockets.models import ( ) from pydantic import Field, field_validator, model_validator -URL_MAP = { - "stock": "wss://socket.polygon.io/stocks", - "stock_delayed": "wss://delayed.polygon.io/stocks", - "options": "wss://socket.polygon.io/options", - "options_delayed": "wss://delayed.polygon.io/options", - "fx": "wss://socket.polygon.io/forex", - "crypto": "wss://socket.polygon.io/crypto", - "index": "wss://socket.polygon.io/indices", - "index_delayed": "wss://delayed.polygon.io/indices", -} ASSET_CHOICES = [ "stock", @@ -1227,12 +1217,8 @@ class PolygonWebSocketFetcher( ) -> dict: """Extract data from the WebSocket.""" api_key = credentials.get("polygon_api_key") if credentials else "" - url = URL_MAP[query.asset_type] - symbol = query.symbol.upper() - kwargs = { - "url": url, "asset_type": query.asset_type, "feed": query.feed, "api_key": api_key, 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 b6345d32ea8..1c6fd004ec8 100644 --- a/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py +++ b/openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py @@ -22,6 +22,17 @@ from openbb_polygon.models.websocket_connection import ( ) from pydantic import ValidationError +URL_MAP = { + "stock": "wss://socket.polygon.io/stocks", + "stock_delayed": "wss://delayed.polygon.io/stocks", + "options": "wss://socket.polygon.io/options", + "options_delayed": "wss://delayed.polygon.io/options", + "fx": "wss://socket.polygon.io/forex", + "crypto": "wss://socket.polygon.io/crypto", + "index": "wss://socket.polygon.io/indices", + "index_delayed": "wss://delayed.polygon.io/indices", +} + logger = get_logger("openbb.websocket.polygon") queue = MessageQueue(logger=logger) command_queue = MessageQueue(logger=logger) @@ -29,8 +40,12 @@ command_queue = MessageQueue(logger=logger) kwargs = parse_kwargs() CONNECT_KWARGS = kwargs.pop("connect_kwargs", {}) FEED = kwargs.pop("feed", None) -ASSET_TYPE = kwargs.pop("asset_type", None) +ASSET_TYPE = kwargs.pop("asset_type", "crypto") kwargs["results_file"] = os.path.abspath(kwargs.get("results_file")) +URL = URL_MAP.get(ASSET_TYPE) + +if not URL: + raise ValueError("Invalid asset type provided.") DATABASE = Database( results_file=kwargs["results_file"], @@ -208,7 +223,7 @@ async def connect_and_stream(): connect_kwargs["close_timeout"] = None try: - async with websockets.connect(kwargs["url"], **connect_kwargs) as websocket: + async with websockets.connect(URL, **connect_kwargs) as websocket: await login(websocket) response = await websocket.recv() diff --git a/openbb_platform/providers/polygon/tests/test_polygon_fetchers.py b/openbb_platform/providers/polygon/tests/test_polygon_fetchers.py index e62c458c331..76e0db916cf 100644 --- a/openbb_platform/providers/polygon/tests/test_polygon_fetchers.py +++ b/openbb_platform/providers/polygon/tests/test_polygon_fetchers.py @@ -6,6 +6,7 @@ from unittest.mock import MagicMock, patch import pytest from openbb_core.app.service.user_service import UserService +from openbb_core.provider.utils.websockets.client import WebSocketClient from openbb_polygon.models.balance_sheet import PolygonBalanceSheetFetcher from openbb_polygon.models.cash_flow import PolygonCashFlowStatementFetcher from openbb_polygon.models.company_news import PolygonCompanyNewsFetcher @@ -25,7 +26,6 @@ from openbb_polygon.models.websocket_connection import ( PolygonWebSocketData, PolygonWebSocketFetcher, ) -from openbb_websockets.client import WebSocketClient test_credentials = UserService().default_user_settings.credentials.model_dump( mode="json" 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 33c3f1828e2..b8d7073e4b9 100644 --- a/openbb_platform/providers/tiingo/openbb_tiingo/models/websocket_connection.py +++ b/openbb_platform/providers/tiingo/openbb_tiingo/models/websocket_connection.py @@ -10,7 +10,7 @@ from openbb_core.provider.abstract.fetcher import Fetcher from openbb_core.provider.utils.descriptions import ( QUERY_DESCRIPTIONS, ) -from openbb_websockets.client import WebSocketClient +from openbb_core.provider.utils.websockets.client import WebSocketClient from openbb_websockets.models import ( WebSocketConnection, WebSocketData, @@ -18,12 +18,6 @@ from openbb_websockets.models import ( ) from pydantic import Field, field_validator, model_validator -URL_MAP = { - "stock": "wss://api.tiingo.com/iex", - "fx": "wss://api.tiingo.com/fx", - "crypto": "wss://api.tiingo.com/crypto", -} - # These are the data array order of definitions. IEX_FIELDS = [ "type", @@ -244,7 +238,6 @@ class TiingoWebSocketFetcher( from asyncio import sleep api_key = credentials.get("tiingo_token") if credentials else "" - url = URL_MAP[query.asset_type] threshold_level = ( 5 if query.asset_type == "fx" or query.feed == "trade" @@ -254,12 +247,11 @@ class TiingoWebSocketFetcher( else 0 ) ) - symbol = query.symbol.lower() kwargs = { - "url": url, "api_key": api_key, + "asset_type": query.asset_type, "threshold_level": threshold_level, "connect_kwargs": query.connect_kwargs, } 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 4a6a17bf511..ef1664d8e20 100644 --- a/openbb_platform/providers/tiingo/openbb_tiingo/utils/websocket_client.py +++ b/openbb_platform/providers/tiingo/openbb_tiingo/utils/websocket_client.py @@ -19,6 +19,12 @@ from openbb_core.provider.utils.websockets.message_queue import MessageQueue from openbb_tiingo.models.websocket_connection import TiingoWebSocketData from pydantic import ValidationError +URL_MAP = { + "stock": "wss://api.tiingo.com/iex", + "fx": "wss://api.tiingo.com/fx", + "crypto": "wss://api.tiingo.com/crypto", +} + # These are the data array definitions. IEX_FIELDS = [ "type", @@ -74,6 +80,10 @@ queue = MessageQueue(logger=logger) kwargs = parse_kwargs() CONNECT_KWARGS = kwargs.pop("connect_kwargs", {}) kwargs["results_file"] = os.path.abspath(kwargs["results_file"]) +URL = URL_MAP.get(kwargs.pop("asset_type", "crypto")) + +if not URL: + raise ValueError("Invalid asset type provided.") DATABASE = Database( results_file=kwargs["results_file"], @@ -86,8 +96,6 @@ DATABASE = Database( # Subscribe and unsubscribe events are handled in a separate connection using the subscription_id set by the login event. async def update_symbols(symbol, event): """Update the symbols to subscribe to.""" - url = kwargs["url"] - if not SUBSCRIPTION_ID: logger.error( "PROVIDER ERROR: Must be assigned a subscription ID to update symbols. Try logging in." @@ -103,7 +111,7 @@ async def update_symbols(symbol, event): }, } - async with websockets.connect(url) as websocket: + async with websockets.connect(URL) as websocket: await websocket.send(json.dumps(update_event)) response = await websocket.recv() message = json.loads(response) @@ -228,7 +236,7 @@ async def connect_and_stream(): try: try: - async with websockets.connect(kwargs["url"], **connect_kwargs) as websocket: + async with websockets.connect(URL, **connect_kwargs) as websocket: logger.info("PROVIDER INFO: WebSocket connection established.") await websocket.send(json.dumps(subscribe_event)) while True: diff --git a/openbb_platform/providers/tiingo/tests/test_tiingo_fetchers.py b/openbb_platform/providers/tiingo/tests/test_tiingo_fetchers.py index 8fbf39aff7b..26fb140de19 100644 --- a/openbb_platform/providers/tiingo/tests/test_tiingo_fetchers.py +++ b/openbb_platform/providers/tiingo/tests/test_tiingo_fetchers.py @@ -6,6 +6,7 @@ from unittest.mock import MagicMock, patch import pytest from openbb_core.app.service.user_service import UserService +from openbb_provider.utils.websockets.client import WebSocketClient from openbb_tiingo.models.company_news import TiingoCompanyNewsFetcher from openbb_tiingo.models.crypto_historical import TiingoCryptoHistoricalFetcher from openbb_tiingo.models.currency_historical import TiingoCurrencyHistoricalFetcher @@ -17,7 +18,6 @@ from openbb_tiingo.models.websocket_connection import ( TiingoWebSocketFetcher, ) from openbb_tiingo.models.world_news import TiingoWorldNewsFetcher -from openbb_websockets.client import WebSocketClient test_credentials = UserService().default_user_settings.credentials.model_dump( mode="json"