jerry/OpenBB
1
0
Fork 0
mirror of https://github.com/OpenBB-finance/OpenBB.git synced 2026-11-06 09:22:26 +00:00
Code Issues Packages Projects Releases Wiki Activity

small update

Browse Source
This commit is contained in:
Danglewood
2024-11-19 18:31:41 -08:00
parent a43eb868dd
commit 7f3b7bc84c
9 changed files with 70 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
openbb_platform/extensions/websockets/README.md vendored
View File

@@ -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.

6
openbb_platform/extensions/websockets/openbb_websockets/models.py
View File

@@ -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

8
openbb_platform/providers/fmp/openbb_fmp/models/websocket_connection.py
View File

@@ -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"])

16
openbb_platform/providers/fmp/openbb_fmp/utils/websocket_client.py
View File

@@ -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:

17
openbb_platform/providers/fmp/tests/test_fmp_fetchers.py
View File

@@ -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()

8
openbb_platform/providers/polygon/openbb_polygon/models/websocket_connection.py
View File

@@ -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"])

16
openbb_platform/providers/polygon/openbb_polygon/utils/websocket_client.py
View File

@@ -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)

8
openbb_platform/providers/tiingo/openbb_tiingo/models/websocket_connection.py
View File

@@ -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"])

12
openbb_platform/providers/tiingo/openbb_tiingo/utils/websocket_client.py
View File

@@ -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: