st_components/session_manager.py

import streamlit
from utils import load_session_cache

def init_session_state(st: streamlit) -> None:
    # TODO apply pydantic
    init_values = {
        'selected_dataset': None, # TODO make default datasets
        'selected_model' : 'gpt-4o',
        'selected_method' : 'llm',
        'use_cache' : True,
        'temperature' : 0.2,
        'selected_library' : 'plotly',
        'lida_ko': None,
        'summary': None,
        'goals': [],
        'selected_goal_object': None,
        'selected_dataframe': None,
        'num_visualizations': 2,
        'num_goals':4,
        'own_goal':None,
        'messages': [],
        'visualizations': [],
        'update_viz': False,
    }

    for key, value in init_values.items():
        if key not in st.session_state:
            st.session_state[key] = value

    session_id = st.query_params.get("session_id")
    session_data = load_session_cache(session_id)

    if session_data:
        for key, value in session_data.items():
            st.session_state[key] = value