[Need Help] Listening to Real Time Updates From Firestore

I followed the helpful blogpost here: https://blog.streamlit.io/streamlit-firestore-continued/ by the awesome @AustinC! However, I am not certain on how to store updates from firestore snapshots into my streamlit session state. Any tips?

import threading
def initialize_session_state():
    if 'field' not in st.session_state:
             st.session_state.field = ''

initialize_session_state()
callback_done = threading.Event()
# Create a callback on_snapshot function to capture changes
def on_snapshot(doc_snapshot, changes, read_time):
    for doc in doc_snapshot:
        print(f'Received document snapshot: {doc.id}')
        print(f'Received document content: {doc.to_dict()["field"]}')
      
        db_dict =  doc.to_dict()

        load_from_db(db_dict, True)


    callback_done.set()

def load_from_db (_doc, from_snapshot = False):

    if from_snapshot:
        db_dict =  _doc

        st.session_state.field= db_dict['field']


# Watch a specific document
doc_watch = db.collection("collection_name").document('document_id').on_snapshot(on_snapshot)

When running something like this – on_snapshot throws an error saying that I did not initialize field in the session state. Do I need to pass the session state into on_snapshot? If so how do I do that?