Session Ending when button is pressed inside Chat Session

Hi there, I am trying to make a GPT-like application which takes DB Schema input and generate queries for that DB according to the Userโ€™s Natural Language input. Now the application is designed in such a way that the DB is connected to the code, and when a query is generated, it is executed directly on the DB to display the results.

I am trying to put some functionality in here, that when a query is generated, it will ask the user if he wants to execute that query or not. So for that, I am trying to add a button to ask user if he want to execute it, and query will only be executed when user has pressed that button. Here is my code for that functionality:

if prompt := st.chat_input("Hi there, what's up?"):

    st.session_state.messages.append({"role": "user", "content": prompt})

    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):

        response = client.chat.completions.create(
            model=st.session_state["openai_model"],
            messages=[
                {"role": m["role"], "content": m["content"] + formatted_system_message}
                for m in st.session_state.messages
            ],
        )


        if response and response.choices:

            # -------------------- HANDLING RESPONSE FROM GPT -------------------- #

            message_content = response.choices[0].message.content  # Corrected access to the message content
            #st.write('Hi World')
            response = message_content
            data_res = handle_response(response)

            formatted_message = """
                                    SQL Query:
                                    {query}

                                    Description:
                                    {description}

                                    Critical Level:
                                    {critical_level}
                                    """.format(query=data_res['Query'], description=data_res['Description'], critical_level=data_res['Critical Level'])

        print('Response we get :',response)

        #--------------------------------------------------------#    
        #--------------- Calling Data Display Function Response -------------#

        try:
            st.button('Run', on_click=display_on_interface(data_res))

        except json.JSONDecodeError as e:

            st.error(f"Error decoding the JSON response: {e}")

        except KeyError as e:

            st.error(f"Missing key in JSON response: {e}")       



        #------------------------------------------------------#
        #-------------- DATABASE EXECUTION CODE ---------------#

        #------------------------------------------------------#



    st.session_state.messages.append({"role": "assistant", "content": formatted_message})

Now when the button is added is display, and I click on the button, It instantly ends my session without display anything.



Now you see, when I press โ€˜RUNโ€™, my session ends and no results are displayed like the dataframe which i am trying to display in the display_on_interface function.

In here:

def display_on_interface(data):
    a = False
    st.subheader("SQL Query")
    st.code(data['Query'], language="sql")

    st.subheader("Description")
    st.write(data['Description'])

    st.subheader("Critical Level")
    st.write(data['Critical Level'])

    query = data['Query']



        try:
            st.button('Run', on_click=display_on_interface(data_res))

        except json.JSONDecodeError as e:

            st.error(f"Error decoding the JSON response: {e}")

        except KeyError as e:

            st.error(f"Missing key in JSON response: {e}")   


Kindly assist me in this, what could be the issue, I have tried several solution like storing state_session, various GPT solution but nothing is working.โ€™

Here is my complete code for reference:

from openai import OpenAI
import streamlit as st
import pandas as pd
import db_connection
from openaiapi import get_completion_from_messages
from prompts import SYSTEM_MESSAGE
import json
import random


# Defining Database Functions

def query_database(query, conn):
    """Run SQL query and return results in a dataframe"""
    return pd.read_sql_query(query, conn)


def handle_response(response):

    response = response.replace('```', '').strip()

    if response.startswith('json'):
        response = response[4:].strip()  # Remove the first four characters and any whitespace

    json_response = json.loads(response)

    query = json_response['Query']  # Make sure the key matches exactly
    description = json_response['Description']
    critical_level = json_response['Critical Level']

    combined_dict = {
    "Query": query,
    "Description": description,
    "Critical Level": critical_level
    }   

    return combined_dict


## ---------------------------------- ##
# ---------- Data To Display -------- ##

def display_on_interface(data):
    a = False
    st.subheader("SQL Query")
    st.code(data['Query'], language="sql")

    st.subheader("Description")
    st.write(data['Description'])

    st.subheader("Critical Level")
    st.write(data['Critical Level'])

    query = data['Query']

    # try:
    # # Run the SQL query and display the results
    #   sql_results = query_database(query, conn)
    #   st.subheader("Query Results:")
    #   st.dataframe(sql_results)
            
    # except Exception as e:
    #         st.error(f"An error occurred: {e}")
    
   
#---------------------------------------------------------#


# -------------- Connections -------------- #

client = OpenAI(api_key=st.secrets["OPENAI_API_KEY"])

conn = db_connection.create_connection()
schemas = db_connection.get_schema_representation()
formatted_system_message = SYSTEM_MESSAGE.format(schema=schemas)


# ----------------------------------------- #

st.title("AskMatic")


# ----------------- CODE ----------------------- #


if "openai_model" not in st.session_state:
    st.session_state["openai_model"] = "gpt-4-1106-preview"

# Prepend the formatted system message to st.session_state.messages at the start of the session
if 'messages' not in st.session_state:
    st.session_state.messages = [{"role": "system", "content": formatted_system_message}]


for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

if prompt := st.chat_input("Hi there, what's up?"):

    st.session_state.messages.append({"role": "user", "content": prompt})

    with st.chat_message("user"):
        st.markdown(prompt)

    with st.chat_message("assistant"):

        response = client.chat.completions.create(
            model=st.session_state["openai_model"],
            messages=[
                {"role": m["role"], "content": m["content"] + formatted_system_message}
                for m in st.session_state.messages
            ],
        )


        if response and response.choices:

            # -------------------- HANDLING RESPONSE FROM GPT -------------------- #

            message_content = response.choices[0].message.content  # Corrected access to the message content
            #st.write('Hi World')
            response = message_content
            data_res = handle_response(response)

            formatted_message = """
                                    SQL Query:
                                    {query}

                                    Description:
                                    {description}

                                    Critical Level:
                                    {critical_level}
                                    """.format(query=data_res['Query'], description=data_res['Description'], critical_level=data_res['Critical Level'])

        print('Response we get :',response)

        #--------------------------------------------------------#    
        #--------------- Calling Data Display Function Response -------------#

        try:
            st.button('Run', on_click=display_on_interface(data_res))

        except json.JSONDecodeError as e:

            st.error(f"Error decoding the JSON response: {e}")

        except KeyError as e:

            st.error(f"Missing key in JSON response: {e}")       



        #------------------------------------------------------#
        #-------------- DATABASE EXECUTION CODE ---------------#

        #------------------------------------------------------#



    st.session_state.messages.append({"role": "assistant", "content": formatted_message})