How to get on_change callback func work with data_editor

Hi,
I am trying to use on_change callback function within data_editor. I can’t seem to get my table updated when I edit the editable fields on the table. My expectation is as soon as I edit the data the table should update in place. Here is my code that I could not get working correctly for me. Currently, when I change data, nothing happens (i.e the prices do not get updated).

if final_dfs:
    st.session_state.final_df = pd.concat(final_dfs, ignore_index=True)

    # Add default values and extra columns
    if "Markup, %" not in st.session_state.final_df.columns:
        st.session_state.final_df["Markup, %"] = 30
    if "External margin, %" not in st.session_state.final_df.columns:
        st.session_state.final_df["External margin, %"] = 50
    if "Select" not in st.session_state.final_df.columns:
        st.session_state.final_df["Select"] = False
    
    # Define a callback function to recalculate the prices
    def recalculate_prices():
        st.session_state.final_df = calculate_price_with_markup(st.session_state.final_df)
        st.session_state.final_df = calculate_total_price(st.session_state.final_df)
        st.session_state.final_df = calculate_final_base_margin(st.session_state.final_df)

        # Rearrange the column sequence
        desired_order = [
            "part_id",
            "diameter_in",
            "Total Cost (USD)",
            "Markup, %",
            "Price with markup (USD)",
            "External margin, %",
            "Total price (USD)",
            "Final base margin, %",
            "Select"
        ]
        st.session_state.final_df = st.session_state.final_df[desired_order]

    # Initial calculation of prices
    recalculate_prices()

    # Define the column configuration
    column_config = {
        "part_id": st.column_config.TextColumn(
            "Part ID",
            disabled=True,
        ),
        "diameter_in": st.column_config.NumberColumn(
            "Diameter",
            disabled=True,
        ),
        "Total Cost (USD)": st.column_config.NumberColumn(
            "Total Cost (USD)",
            disabled=True,
        ),
        "Markup, %": st.column_config.NumberColumn(
            "Markup, %",
        ),
        "External margin, %": st.column_config.NumberColumn(
            "External margin, %",

        ),
        "Price with markup (USD)": st.column_config.NumberColumn(
            "Price with markup (USD)",
            disabled=True,

        ),
        "Total price (USD)": st.column_config.NumberColumn(
            "Total price (USD)",
            disabled=True,
        ),
        "Final base margin, %": st.column_config.NumberColumn(
            "Final base margin, %",
            disabled=True,
        ),

        "Select": st.column_config.CheckboxColumn("Select"),
    }

    # Display the data editor
    edited_df = st.data_editor(
        st.session_state.final_df,
        key="data_editor",
        hide_index=True,
        column_config=column_config,
        on_change=recalculate_prices,
    )

The on_change callback lacks access to the most recent edited data values. Consequently, I have implemented a manual update procedure for the session state data frame after an edit, ensuring that the prices are recalculated accordingly.

import streamlit as st
import pandas as pd

# Placeholder functions for price calculations
def calculate_price_with_markup(df):
    df["Price with markup (USD)"] = df["Total Cost (USD)"] * (1 + df["Markup, %"] / 100)
    return df

def calculate_total_price(df):
    df["Total price (USD)"] = df["Price with markup (USD)"] * (1 + df["External margin, %"] / 100)
    return df

def calculate_final_base_margin(df):
    df["Final base margin, %"] = (df["Total price (USD)"] - df["Total Cost (USD)"]) / df["Total price (USD)"] * 100
    return df

# Assuming final_dfs is a list of DataFrames
final_dfs = [
    pd.DataFrame({
        "part_id": ["A1", "B1", "C1"],
        "diameter_in": [1.0, 2.5, 3.0],
        "Total Cost (USD)": [10, 20, 30],
    }),
    pd.DataFrame({
        "part_id": ["D1", "E1", "F1"],
        "diameter_in": [4.0, 5.5, 6.0],
        "Total Cost (USD)": [40, 50, 60],
    })
]

if final_dfs:

    # Define a callback function to recalculate the prices
    def recalculate_prices():
        print("Recalculating prices...")

        st.session_state.final_df = calculate_price_with_markup(st.session_state.final_df)
        st.session_state.final_df = calculate_total_price(st.session_state.final_df)
        st.session_state.final_df = calculate_final_base_margin(st.session_state.final_df)

        # Rearrange the column sequence
        desired_order = [
            "part_id",
            "diameter_in",
            "Total Cost (USD)",
            "Markup, %",
            "Price with markup (USD)",
            "External margin, %",
            "Total price (USD)",
            "Final base margin, %",
            "Select"
        ]
        st.session_state.final_df = st.session_state.final_df[desired_order]


    # Initialize the final DataFrame if not already present
    if 'final_df' not in st.session_state:
        st.session_state.final_df = pd.concat(final_dfs, ignore_index=True)

        # Add default values and extra columns
        if "Markup, %" not in st.session_state.final_df.columns:
            st.session_state.final_df["Markup, %"] = 30
        if "External margin, %" not in st.session_state.final_df.columns:
            st.session_state.final_df["External margin, %"] = 50
        if "Select" not in st.session_state.final_df.columns:
            st.session_state.final_df["Select"] = False

        recalculate_prices()

    # Define the column configuration
    column_config = {
        "part_id": st.column_config.TextColumn(
            "Part ID",
            disabled=True,
        ),
        "diameter_in": st.column_config.NumberColumn(
            "Diameter",
            disabled=True,
        ),
        "Total Cost (USD)": st.column_config.NumberColumn(
            "Total Cost (USD)",
            disabled=True,
        ),
        "Markup, %": st.column_config.NumberColumn(
            "Markup, %",
        ),
        "External margin, %": st.column_config.NumberColumn(
            "External margin, %",

        ),
        "Price with markup (USD)": st.column_config.NumberColumn(
            "Price with markup (USD)",
            disabled=True,

        ),
        "Total price (USD)": st.column_config.NumberColumn(
            "Total price (USD)",
            disabled=True,
        ),
        "Final base margin, %": st.column_config.NumberColumn(
            "Final base margin, %",
            disabled=True,
        ),

        "Select": st.column_config.CheckboxColumn("Select"),
    }

    # Display the data editor
    edited_df = st.data_editor(
        st.session_state.final_df,
        key="data_editor",
        hide_index=True,
        column_config=column_config,
    )

    # Check if the data has been edited and recalculate prices
    if not edited_df.equals(st.session_state.final_df):
        st.session_state.final_df = edited_df
        recalculate_prices()

        # Trigger a rerun to update the displayed data
        st.rerun()