app.py

from utils import (
    container_emphasis_css,
    container_regular_css,
    question_font,
    modify_sql_guid_injection
)

import streamlit as st
import json
import pandas as pd
from datetime import datetime
from streamlit_extras.stylable_container import stylable_container

@st.cache_data
def load_data(file):
    return pd.read_csv(file)

st.set_page_config(layout="wide", initial_sidebar_state="expanded")
st.title("Relevancy Compare App")
st.divider()
st.markdown(question_font, unsafe_allow_html=True)


def edit_df_data_change():
    edited_rows: dict = st.session_state.edit_df_data['edited_rows']
    st.session_state.previous_row_index = st.session_state.previous_row_index
    st.session_state.current_row_index = next(iter(edited_rows))
    st.session_state.df = st.session_state.df.assign(selected=False)
    update_dict = {idx: values for idx, values in edited_rows.items()}
    st.session_state.df.update(pd.DataFrame.from_dict(update_dict, orient='index'))


st.need_rerun = False
if "df" not in st.session_state:
    role_pick = st.selectbox("You are?", options=["Kyle", "Shadi", "Daniel", "Xin"], index=None, key='role_pick')
    uploaded_file = st.file_uploader("Choose a CSV file", type="csv")
    st.session_state.role = role_pick
    st.need_rerun = True
else:
    uploaded_file = st.session_state.df
if uploaded_file is not None and 'role' in st.session_state:
    "You are modifying as :", st.session_state.role
    with st.container():
        if "df" not in st.session_state:
            df = load_data(uploaded_file)
            st.session_state.df_shape = df.shape
            # selected load
            select_values = [False] * len(df)
            
            if "Comment" not in df.columns:
                comments = [''] * len(df)
                df.insert(0, 'Comment', comments)
                st.session_state.default_row_index = 0
            
            if "LastModified" not in df.columns:
                modifies = [''] * len(df)
                df.insert(0, 'LastModified', modifies)

            else:
                df['Comment'] = df['Comment'].astype(str)
                df['LastModified'] = df['LastModified'].astype(str)
                comments = [c if c != 'nan' else '' for c in df['Comment'].tolist()]
                modifies = [c if c != 'nan' else '' for c in df['LastModified'].tolist()]
                df['Comment'] = comments
                df['LastModified'] = modifies
                first_nan = 0
                for i, c in enumerate(comments):
                    if not c:
                        first_nan = i
                        break
                st.session_state.default_row_index = first_nan
            
            st.session_state.current_row_index = st.session_state.default_row_index
            st.session_state.previous_row_index = st.session_state.current_row_index

            df.insert(0, 'Select', select_values)
            df.at[st.session_state.current_row_index, 'Select'] = True
            st.session_state['df'] = df
        
        df_col, = st.columns(1)

    st.session_state.previous_row_index = st.session_state.current_row_index
    with st.container():
        bcol1, bcol2, _ = st.columns([1, 1, 5])
        with bcol1:
            previous_row = st.button("Prev Row")
        with bcol2:
            next_row = st.button("Next Row")
            

    if "current_row_index" in st.session_state and previous_row:
        if st.session_state.current_row_index > 0:
            st.session_state.previous_row_index = st.session_state.current_row_index
            st.session_state.current_row_index -= 1

    if "current_row_index" in st.session_state and next_row:
        if st.session_state.current_row_index < st.session_state.df_shape[0] - 1:
            st.session_state.previous_row_index = st.session_state.current_row_index
            st.session_state.current_row_index += 1

    row_index = st.session_state.current_row_index
    row_data = st.session_state.df.loc[row_index]
    question = row_data['question']
    pred_relevant_choice = row_data['pred_relevant_choice']
    correct_response = row_data['correct_response']
    relative_relevancy = json.loads(row_data['relative_relevancy'])
    reason = relative_relevancy['reason']
    most_relevant = relative_relevancy['most_relevant']
    options=['first', 'second', 'both']
    default_index = options.index(most_relevant)
    default_comment = row_data['Comment']
    first_css = container_regular_css
    second_css = container_regular_css
    first_color = ':black'
    second_color = ':black'
    
        
    st.write(f"##### Row {row_index}/{st.session_state.df_shape[0]-1}")
    with st.container(border=True):
        # st.write(f"Question: {question}")
        st.markdown(f'<p class="question-font">[Q]: {question}</p>', unsafe_allow_html=True)
        select_by = st.selectbox("Mark By: ", ["pred_relevant_choice", "correct_response"], key='mark_by_dropdown')
        col1_res, col2_res, _ = st.columns([0.2, 0.2, 0.6])
        if select_by == 'pred_relevant_choice':
            col1_res.write(f":blue[pred_relevant_choice: {pred_relevant_choice}]")
            col2_res.write(f"correct_response: {correct_response}")
        else:
            col1_res.write(f"pred_relevant_choice: {pred_relevant_choice}")
            col2_res.write(f":blue[correct_response: {correct_response}]")


        if row_data[select_by].lower() == 'first' or row_data[select_by].lower() == 'both':
            first_css = container_emphasis_css
            first_color = ':blue'

            if row_data[select_by].lower() == 'both':
                second_css = container_emphasis_css
                second_color = ':blue'
        else:
            second_css = container_emphasis_css
            second_color = ':blue'

        col1, col2 = st.columns(2)
        with col1:
            st.write("Target SQL:", use_container_width=True)
            with stylable_container(
                key="first_sql_container",
                css_styles=first_css
                ):
                target_sql_text = modify_sql_guid_injection(row_data['target_sql'].strip())
                #st.write(f"{first_color}[{target_sql_text}]", use_container_width=True)
                col1, _ = st.columns([0.99, 0.01])
                with col1:
                    st.code(
                        target_sql_text,
                        language="sql"
                    )
                    # st.code(
                    #     "\n".join(
                    #         tw.wrap(
                    #         target_sql_text,
                    #         width=60,
                    #         )
                    #     ),
                    #     language="sql"
                    # )

        with col2:
            st.write("Predicted SQL:", use_container_width=True)
            with stylable_container(
                key="second_sql_container",
                css_styles=second_css
                ):
                predicted_sql_text = modify_sql_guid_injection(row_data['predicted_sql'].strip())
                # st.write(f"{second_color}[{predicted_sql_text}]")
                col1, _ = st.columns([0.99, 0.01])
                with col1:
                     st.code(
                        predicted_sql_text,
                        language="sql"
                    )
        col1_comment, col2_comment = st.columns([1, 1])
        with col1_comment:
            select_most_relevant = st.selectbox("most_relevant", options=options, index=default_index, key='most_relevant_drop_down')
            reason_box = st.text_area("reason", value=reason, key='reason_text')
        with col2_comment:
            comment = st.text_area("Comment", value=default_comment, key='comment_text')
    
    with st.container():
        bcol1, bcol2, _, _, _, _, _, _, _, _ = st.columns(10)
        with bcol1:
            save = st.button("Save")
        with bcol2:
            download = st.button("Download")

    if save:
        updated_relative_relevancy = json.dumps({'most_relevant': select_most_relevant, "reason": reason_box})
        st.session_state.df.at[row_index, 'Comment'] = comment
        st.session_state.df.at[row_index, 'relative_relevancy'] = updated_relative_relevancy
        st.session_state.df.at[row_index, 'correct_response'] = select_most_relevant
        current_time = datetime.now()
        current_time_str = current_time.strftime("%Y-%m-%d %H:%M:%S")
        st.session_state.df.at[row_index, 'LastModified'] = f"{st.session_state.role}@{current_time_str}"
        df = st.session_state.df.drop('Select', axis=1)
        df.to_csv('modified.csv', index=False)
        st.rerun()

    if download:
        with open('modified.csv') as f:
            st.download_button('Download modified CSV', f, file_name='modified.csv')
    
    st.session_state.df['Select'] = False
    st.session_state.df.at[st.session_state.current_row_index, 'Select'] = True

    edited_df = df_col.data_editor(st.session_state.df, num_rows="fixed", use_container_width=True, key='edit_df_data', on_change=edit_df_data_change)

    #st.session_state.mannul_select = edited_df[edited_df['Select']].index
    st.session_state.df = edited_df

    if st.need_rerun:
        st.need_rerun = False
        st.rerun()