Filter PNL in streamlit dashboard

elfpy/data/postgres.py

@@ -369,6 +369,29 @@ def get_current_wallet_info(
     return current_wallet_info
 
 
+def get_agents(session: Session, start_block: int | None = None, end_block: int | None = None) -> list[str]:
+    """Gets the list of all agents from the WalletInfo table"""
+    query = session.query(WalletInfo.walletAddress)
+    # Support for negative indices
+    if (start_block is not None) and (start_block < 0):
+        start_block = _get_latest_block_number_wallet_info(session) + start_block + 1
+    if (end_block is not None) and (end_block < 0):
+        end_block = _get_latest_block_number_wallet_info(session) + end_block + 1
+
+    if start_block is not None:
+        query = query.filter(WalletInfo.blockNumber >= start_block)
+    if end_block is not None:
+        query = query.filter(WalletInfo.blockNumber < end_block)
+
+    if query is None:
+        return []
+    query = query.distinct()
+
+    results = pd.read_sql(query.statement, con=session.connection())
+
+    return results["walletAddress"].to_list()
+
+
 def get_latest_block_number(session: Session) -> int:
     """Gets the latest block number based on the pool info table in the db
     Arguments

examples/hackweek_demo/extract_data_logs.py

@@ -108,4 +108,6 @@ def decode_prefix(row):
     trade_data["timestamp"] = trade_data["block_timestamp"]
     trade_data["block_timestamp"] = trade_data["block_timestamp"].astype(int)
 
+    trade_data = trade_data.sort_values("block_timestamp")
+
     return trade_data

examples/hackweek_demo/plot_pnl.py

@@ -112,12 +112,18 @@ def get_wallet_from_onchain_trade_info(
     return wallet
 
 
-def calculate_pnl(trade_data):
+def calculate_pnl(trade_data, agent_list):
     # pylint: disable=too-many-locals
     """Calculates the pnl given trade data"""
     # Drop all rows with nan maturity timestamps
     trade_data = trade_data[~trade_data["maturity_time"].isna()]
 
+    # Filter trade_data based on agent_list
+    trade_data = trade_data[trade_data["operator"].isin(agent_list)]
+
+    if len(trade_data) == 0:
+        return (pd.DataFrame([]), pd.DataFrame([]))
+
     # %% estimate position duration and add it in
     position_duration = max(trade_data.maturity_time - trade_data.block_timestamp)
     # print(f"empirically observed position_duration in seconds: {position_duration}")
@@ -143,7 +149,7 @@ def calculate_pnl(trade_data):
     for _, row in trade_data.iterrows():
         new_pnl: dict[str, float] = {}
         for agent in agents:
-            agent_index = agents.index(agent)
+            agent_shortname = agent[:6] + "..." + agent[-4:]
 
             marginal_trades = pd.DataFrame(row).T
 
@@ -180,7 +186,7 @@ def calculate_pnl(trade_data):
             for _, short in agent_wallets[agent].shorts.items():
                 pnl += float(short.balance) * (1 - spot_price)
 
-            new_pnl[f"agent_{agent_index}_pnl"] = pnl
+            new_pnl[f"{agent_shortname}"] = pnl
         pnl_data.append(new_pnl)
         time_data.append(row["timestamp"])
 

examples/hackweek_demo/run_demo.py

@@ -20,22 +20,14 @@
 # None to view all
 view_window = None
 
-# set up streamlit
-# creating a single-element container
-placeholder = st.empty()
 
 # %%
 # near real-time / live feed simulation
 
 ## Get transactions from data
 
 
-def get_ticker(data):
-    """Given transaction data, return a subset of the dataframe"""
-    # Return reverse of methods to put most recent transactions at the top
-    out = data[["blockNumber", "input_method"]].set_index("blockNumber").iloc[::-1]
-    return out
-
+# TODO should these all be seperate figures instead of one figure?
 
 fig = mpf.figure(style="mike", figsize=(15, 15))  # type: ignore
 ax_ohlcv = fig.add_subplot(2, 2, 1)
@@ -56,44 +48,70 @@ def get_ticker(data):
 # pool config data is static, so just read once
 config_data = postgres.get_pool_config(session)
 
+# Set up agent selection state list outside of live updates
+st.session_state.options = []
+
+
+def get_ticker(data):
+    """Given transaction data, return a subset of the dataframe"""
+    # Return reverse of methods to put most recent transactions at the top
+    out = data[["blockNumber", "input_method"]].set_index("blockNumber").iloc[::-1]
+    return out
+
+
+agent_list = postgres.get_agents(session, start_block=start_block)
+
+filter_agents = st.multiselect("PNL Agents", agent_list, key="agent_select")
+
+
+# creating a single-element container
+main_placeholder = st.empty()
+# creating a single-element sidebar
+sidebar_placeholder = st.sidebar.empty()
+
+
 while True:
     txn_data = postgres.get_transactions(session, start_block=start_block)
     pool_info_data = postgres.get_pool_info(session, start_block=start_block)
 
     combined_data = get_combined_data(txn_data, pool_info_data)
 
+    if len(combined_data) == 0:
+        time.sleep(0.1)
+        continue
+
     ohlcv = calc_ohlcv(combined_data, freq="5T")
 
     ticker = get_ticker(txn_data)
 
     (fixed_rate_x, fixed_rate_y) = calc_fixed_rate(combined_data)
 
-    pnl_x, pnl_y = calculate_pnl(combined_data)
+    pnl_x, pnl_y = calculate_pnl(combined_data, filter_agents)
 
     # Plot reserve levels (share and bond reserves, in poolinfo)
 
     # Fix axes labels
 
     # Add ticker
 
-    with placeholder.container():
-        ticker_col, plots_col = st.columns([0.2, 1], gap="small")
-        with ticker_col:
-            st.dataframe(ticker.iloc[:100], height=900, width=170)
-
-        with plots_col:
-            # Clears all axes
-            ax_ohlcv.clear()
-            ax_fixed_rate.clear()
-            ax_vol.clear()
-            ax_pnl.clear()
-
-            # TODO add in volume
-            plot_ohlcv(ohlcv, ax_ohlcv, ax_vol)
-            plot_fixed_rate(fixed_rate_x, fixed_rate_y, ax_fixed_rate)
-            plot_pnl(pnl_x, pnl_y, ax_pnl)
-
-            fig.autofmt_xdate()  # type: ignore
-            st.pyplot(fig=fig)  # type: ignore
+    # Selection options for filtering agent, defined outside of update loop
+    # Session state gets updated inside loop
+    # with plots_col:
+    with sidebar_placeholder.container():
+        st.dataframe(ticker.iloc[:100], height=900, use_container_width=True)
+
+    with main_placeholder.container():
+        # Clears all axes
+        ax_vol.clear()
+        ax_pnl.clear()
+        ax_ohlcv.clear()
+        ax_fixed_rate.clear()
+
+        plot_ohlcv(ohlcv, ax_ohlcv, ax_vol)
+        plot_fixed_rate(fixed_rate_x, fixed_rate_y, ax_fixed_rate)
+        plot_pnl(pnl_x, pnl_y, ax_pnl)
+
+        fig.autofmt_xdate()  # type: ignore
+        st.pyplot(fig=fig)  # type: ignore
 
     time.sleep(0.1)

tests/data/test_wallet_info.py

@@ -142,3 +142,15 @@ def test_current_wallet_info(self, session):
         wallet_info_df = postgres.get_current_wallet_info(session).reset_index()
         assert wallet_info_df["tokenType"].equals(pd.Series(["BASE", "LP"], name="tokenType"))
         assert wallet_info_df["tokenValue"].equals(pd.Series([6.1, 5.1], name="tokenValue"))
+
+    def test_get_agents(self, session):
+        """Testing helper function to get current wallet values"""
+        wallet_info_1 = WalletInfo(blockNumber=0, walletAddress="addr_1")  # add your other columns here...
+        wallet_info_2 = WalletInfo(blockNumber=1, walletAddress="addr_1")  # add your other columns here...
+        wallet_info_3 = WalletInfo(blockNumber=2, walletAddress="addr_2")  # add your other columns here...
+        postgres.add_wallet_infos([wallet_info_1, wallet_info_2, wallet_info_3], session)
+
+        agents = postgres.get_agents(session)
+        assert len(agents) == 2
+        assert "addr_1" in agents
+        assert "addr_2" in agents